From 62462e2e7826b8da2a4d0275f9ec73701c48e008 Mon Sep 17 00:00:00 2001
From: omerlavanet <omer@lavanet.xyz>
Date: Mon, 23 Dec 2024 13:45:07 +0200
Subject: [PATCH] add pool-rewards-breakdown command

---
 x/subscription/client/cli/query.go            |   1 +
 .../client/cli/query_estimated_rewards.go     | 119 ++++++++++++++++++
 2 files changed, 120 insertions(+)

diff --git a/x/subscription/client/cli/query.go b/x/subscription/client/cli/query.go
index 1c78247f2a..9b696a97eb 100644
--- a/x/subscription/client/cli/query.go
+++ b/x/subscription/client/cli/query.go
@@ -27,6 +27,7 @@ func GetQueryCmd(queryRoute string) *cobra.Command {
 	cmd.AddCommand(CmdQueryParams())
 	cmd.AddCommand(CmdCurrent())
 	cmd.AddCommand(CmdEstimatedProviderRewards())
+	cmd.AddCommand(CmdPoolRewards())
 	cmd.AddCommand(CmdEstimatedValidatorsRewards())
 	cmd.AddCommand(CmdTrackedCuUsage())
 
diff --git a/x/subscription/client/cli/query_estimated_rewards.go b/x/subscription/client/cli/query_estimated_rewards.go
index 1f82c3a345..4c3f529198 100644
--- a/x/subscription/client/cli/query_estimated_rewards.go
+++ b/x/subscription/client/cli/query_estimated_rewards.go
@@ -1,9 +1,14 @@
 package cli
 
 import (
+	"fmt"
+
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/flags"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/lavanet/lava/v4/utils"
+	dualstakingtypes "github.com/lavanet/lava/v4/x/dualstaking/types"
+	epochstoragetypes "github.com/lavanet/lava/v4/x/epochstorage/types"
 	"github.com/lavanet/lava/v4/x/subscription/types"
 	"github.com/spf13/cobra"
 )
@@ -62,3 +67,117 @@ func CmdEstimatedProviderRewards() *cobra.Command {
 
 	return cmd
 }
+
+func CmdPoolRewards() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "pool-rewards-breakdown",
+		Short: "Calculates estimated rewards for all pools",
+		Long:  `estimate the total rewards a pool will give to all providers if the month ends now`,
+		Args:  cobra.ExactArgs(0),
+		RunE: func(cmd *cobra.Command, args []string) (err error) {
+			clientCtx, err := client.GetClientQueryContext(cmd)
+			if err != nil {
+				return err
+			}
+
+			queryClient := types.NewQueryClient(clientCtx)
+
+			// get all provider addresses
+			provider := ""
+			epochStorageEueryClient := epochstoragetypes.NewQueryClient(clientCtx)
+			res, err := epochStorageEueryClient.ProviderMetaData(cmd.Context(), &epochstoragetypes.QueryProviderMetaDataRequest{Provider: provider})
+			if err != nil {
+				return err
+			}
+			addresses := []string{}
+
+			for _, meta := range res.GetMetaData() {
+				addresses = append(addresses, meta.GetProvider())
+			}
+			runEstimateWithRetries := func(req types.QueryEstimatedProviderRewardsRequest) (*types.QueryEstimatedRewardsResponse, error) {
+				res, err := queryClient.EstimatedProviderRewards(cmd.Context(), &req)
+				if err != nil {
+					res, err = queryClient.EstimatedProviderRewards(cmd.Context(), &req)
+					if err != nil {
+						res, err = queryClient.EstimatedProviderRewards(cmd.Context(), &req)
+						if err != nil {
+							return nil, err
+						}
+					}
+				}
+				return res, err
+			}
+
+			summary := map[string]types.EstimatedRewardInfo{}
+			total := sdk.DecCoins{}
+			for idx, provider := range addresses {
+				fmt.Printf("\rProgress: %d/%d", idx+1, len(addresses))
+				req := types.QueryEstimatedProviderRewardsRequest{Provider: provider}
+				res, err := runEstimateWithRetries(req)
+				if err != nil {
+					utils.LavaFormatError("failed to query provider", err, utils.Attribute{Key: "provider", Value: provider})
+					continue
+				}
+				total = summarizeForRes(res, summary, total)
+				dualStakingQueryClient := dualstakingtypes.NewQueryClient(clientCtx)
+				resDel, err := dualStakingQueryClient.ProviderDelegators(cmd.Context(), &dualstakingtypes.QueryProviderDelegatorsRequest{Provider: provider})
+				if err != nil {
+					resDel, err = dualStakingQueryClient.ProviderDelegators(cmd.Context(), &dualstakingtypes.QueryProviderDelegatorsRequest{Provider: provider})
+					if err != nil {
+						resDel, err = dualStakingQueryClient.ProviderDelegators(cmd.Context(), &dualstakingtypes.QueryProviderDelegatorsRequest{Provider: provider})
+						if err != nil {
+							return err
+						}
+
+					}
+				}
+				delegations := resDel.GetDelegations()
+				for idx2, del := range delegations {
+					fmt.Printf("\rProgress: %d/%d %d/%d", idx+1, len(addresses), idx2+1, len(delegations))
+					delegatorName := del.Delegator
+					req := types.QueryEstimatedProviderRewardsRequest{Provider: provider, AmountDelegator: delegatorName}
+					res, err := runEstimateWithRetries(req)
+					if err != nil {
+						utils.LavaFormatError("failed to query delegator rewards", err, utils.LogAttr("provider", provider), utils.LogAttr("delegator", delegatorName))
+						continue
+					}
+					total = summarizeForRes(res, summary, total)
+				}
+			}
+			fmt.Printf("\n---- results ----\n\n")
+			info := []types.EstimatedRewardInfo{}
+			for _, sumEntry := range summary {
+				info = append(info, sumEntry)
+			}
+			printMe := &types.QueryEstimatedRewardsResponse{
+				Info:  info,
+				Total: total,
+			}
+			return clientCtx.PrintProto(printMe)
+		},
+	}
+
+	flags.AddQueryFlagsToCmd(cmd)
+
+	return cmd
+}
+
+func summarizeForRes(res *types.QueryEstimatedRewardsResponse, summary map[string]types.EstimatedRewardInfo, total sdk.DecCoins) sdk.DecCoins {
+	info := res.Info
+	for _, entry := range info {
+		if _, ok := summary[entry.Source]; !ok {
+			summary[entry.Source] = entry
+		} else {
+			entryIn := summary[entry.Source]
+			coinsArr := entry.Amount
+			for _, coin := range coinsArr {
+				entryIn.Amount = entryIn.Amount.Add(coin)
+			}
+			summary[entry.Source] = entryIn
+		}
+	}
+	for _, coin := range res.Total {
+		total = total.Add(coin)
+	}
+	return total
+}