From a93ec9fc045db0afcb05533ba75429001fe95d89 Mon Sep 17 00:00:00 2001
From: rustdev <placex.com@gmail.com>
Date: Wed, 20 Dec 2023 23:18:53 +0000
Subject: [PATCH] override end block custom staking module implementation

---
 custom/staking/keeper/keeper.go     |  4 ++--
 custom/staking/keeper/msg_server.go | 33 +++++++++++---------------
 custom/staking/module.go            | 36 +++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/custom/staking/keeper/keeper.go b/custom/staking/keeper/keeper.go
index 7717e2249..0eb850289 100644
--- a/custom/staking/keeper/keeper.go
+++ b/custom/staking/keeper/keeper.go
@@ -14,7 +14,7 @@ type Keeper struct {
 	cdc               codec.BinaryCodec
 	acck              accountkeeper.AccountKeeper
 	mintkeeper        *mintkeeper.Keeper
-	stakingmiddleware *stakingmiddleware.Keeper
+	Stakingmiddleware *stakingmiddleware.Keeper
 	authority         string
 }
 
@@ -46,7 +46,7 @@ func NewKeeper(
 		acck:              acck,
 		authority:         authority,
 		mintkeeper:        mintkeeper,
-		stakingmiddleware: stakingmiddleware,
+		Stakingmiddleware: stakingmiddleware,
 		cdc:               cdc,
 	}
 	return keeper
diff --git a/custom/staking/keeper/msg_server.go b/custom/staking/keeper/msg_server.go
index 79cccfdc2..6547c5046 100644
--- a/custom/staking/keeper/msg_server.go
+++ b/custom/staking/keeper/msg_server.go
@@ -3,7 +3,6 @@ package keeper
 import (
 	"context"
 
-	"cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
@@ -32,28 +31,24 @@ func (k msgServer) EditValidator(goCtx context.Context, msg *types.MsgEditValida
 func (k msgServer) Delegate(goCtx context.Context, msg *types.MsgDelegate) (*types.MsgDelegateResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	// bondDenom := k.BondDenom(ctx)
-	// if msg.Amount.Denom != bondDenom {
+	bondDenom := k.BondDenom(ctx)
+	if msg.Amount.Denom != bondDenom {
+		return nil, sdkerrors.Wrapf(
+			sdkerrors.ErrInvalidRequest, "invalid coin denomination: got %s, expected %s", msg.Amount.Denom, bondDenom,
+		)
+	}
+
+	// k.mintkeeper.SetLastTotalPower(ctx, math.Int{})
+	// k.stakingmiddleware.SetLastTotalPower(ctx, math.Int{})
+
+	// delegations := k.Stakingmiddleware.DequeueAllDelegation(ctx)
+	// if len(delegations) > 2 {
 	// 	return nil, sdkerrors.Wrapf(
-	// 		sdkerrors.ErrInvalidRequest, "invalid coin denomination: got %s, expected %s", msg.Amount.Denom, bondDenom,
+	// 		sdkerrors.ErrInvalidRequest, "should always be less then X : got %s, expected %s", len(delegations), 1,
 	// 	)
 	// }
 
-	// delegation := types.Delegation{
-	// 	DelegatorAddress: msg.DelegatorAddress,
-	// 	ValidatorAddress: msg.ValidatorAddress,
-	// 	Shares:           msg.Amount.Amount.ToLegacyDec(),
-	// }
-	k.mintkeeper.SetLastTotalPower(ctx, math.Int{})
-	k.stakingmiddleware.SetLastTotalPower(ctx, math.Int{})
-
-	k.stakingmiddleware.SetDelegation(ctx, msg.DelegatorAddress, msg.ValidatorAddress, msg.Amount.Denom, msg.Amount.Amount)
-	delegations := k.stakingmiddleware.DequeueAllDelegation(ctx)
-	if len(delegations) >= 1 {
-		return nil, sdkerrors.Wrapf(
-			sdkerrors.ErrInvalidRequest, "should always be less then 2 : got %s, expected %s", len(delegations),
-		)
-	}
+	k.Stakingmiddleware.SetDelegation(ctx, msg.DelegatorAddress, msg.ValidatorAddress, msg.Amount.Denom, msg.Amount.Amount)
 
 	return &types.MsgDelegateResponse{}, nil
 	// return nil, fmt.Errorf("My custom error: Nikita")
diff --git a/custom/staking/module.go b/custom/staking/module.go
index 8c2ec5b75..3c31fe4fa 100644
--- a/custom/staking/module.go
+++ b/custom/staking/module.go
@@ -3,6 +3,7 @@ package bank
 import (
 	"fmt"
 
+	abcitype "github.com/cometbft/cometbft/abci/types"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	stakingmodule "github.com/cosmos/cosmos-sdk/x/staking"
@@ -11,6 +12,7 @@ import (
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	// custombankkeeper "github.com/notional-labs/composable/v6/custom/bank/keeper"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	customstakingkeeper "github.com/notional-labs/composable/v6/custom/staking/keeper"
 )
 
@@ -53,3 +55,37 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 		panic(fmt.Sprintf("failed to migrate x/staking from version 3 to 4: %v", err))
 	}
 }
+
+// func (am AppModule) BeginBlock(ctx sdk.Context, _abc abcitype.RequestBeginBlock) {
+// 	//Define the logic around the batching.
+// 	am.AppModule.BeginBlock(ctx, _abc)
+// }
+
+func (am AppModule) EndBlock(ctx sdk.Context, _abc abcitype.RequestEndBlock) []abcitype.ValidatorUpdate {
+	//Define the logic around the batching.
+	//TODO!!!
+
+	println("EndBlock Custom Staking Module")
+
+	delegations := am.keeper.Stakingmiddleware.DequeueAllDelegation(ctx)
+	println("Delegations: ", delegations)
+	println("Delegations len: ", len(delegations))
+	//for delegations print the delegator address and the validator address
+	for _, delegation := range delegations {
+		println("Delegator Address: ", delegation.DelegatorAddress)
+		println("Validator Address: ", delegation.ValidatorAddress)
+		fmt.Println("Amount", delegation.Amount.Amount)
+	}
+
+	// am.keeper.Delegate()
+	ctx.EventManager().EmitEvent(
+		sdk.NewEvent(
+			"DequeueAllDelegation",
+			// sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
+			// sdk.NewAttribute(types.AttributeKeyValidator, dvPair.ValidatorAddress),
+			// sdk.NewAttribute(types.AttributeKeyDelegator, dvPair.DelegatorAddress),
+		),
+	)
+
+	return am.AppModule.EndBlock(ctx, _abc)
+}