keeper.go

package keeper

import (
	"context"
	"errors"

	"github.com/cosmos/cosmos-sdk/codec"
	sdk "github.com/cosmos/cosmos-sdk/types"
	"github.com/ethereum/go-ethereum/common"

	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"

	"cosmossdk.io/collections"
	storetypes "cosmossdk.io/core/store"
	"cosmossdk.io/log"

	"github.com/pushchain/push-chain-node/x/uexecutor/types"
)

type Keeper struct {
	cdc codec.BinaryCodec

	logger        log.Logger
	schemaBuilder *collections.SchemaBuilder

	// state management
	storeService      storetypes.KVStoreService
	Params            collections.Item[types.Params]
	authority         string
	evmKeeper         types.EVMKeeper
	feemarketKeeper   types.FeeMarketKeeper
	bankKeeper        types.BankKeeper
	accountKeeper     types.AccountKeeper
	uregistryKeeper  types.UregistryKeeper
	uvalidatorKeeper types.UValidatorKeeper

	// Inbound trackers
	PendingInbounds collections.KeySet[string]

	// UniversalTx collection
	UniversalTx collections.Map[string, types.UniversalTx]

	// Module account manual nonce
	ModuleAccountNonce collections.Item[uint64]

	// GasPrices — deprecated, replaced by ChainMetas. Kept for genesis backward compat.
	GasPrices collections.Map[string, types.GasPrice]

	// ChainMetas collection stores aggregated chain metadata (gas price + block height) for each chain
	ChainMetas collections.Map[string, types.ChainMeta]

	// PendingOutbounds is a secondary index of outbounds with PENDING status.
	// Key: outbound ID -> Value: PendingOutboundEntry
	PendingOutbounds collections.Map[string, types.PendingOutboundEntry]
}

// NewKeeper creates a new Keeper instance
func NewKeeper(
	cdc codec.BinaryCodec,
	storeService storetypes.KVStoreService,
	logger log.Logger,
	authority string,
	evmKeeper types.EVMKeeper,
	feemarketKeeper types.FeeMarketKeeper,
	bankKeeper types.BankKeeper,
	accountKeeper types.AccountKeeper,
	uregistryKeeper types.UregistryKeeper,
	uvalidatorKeeper types.UValidatorKeeper,
) Keeper {
	logger = logger.With(log.ModuleKey, "x/"+types.ModuleName)

	sb := collections.NewSchemaBuilder(storeService)

	if authority == "" {
		authority = authtypes.NewModuleAddress(govtypes.ModuleName).String()
	}

	k := Keeper{
		cdc:           cdc,
		logger:        logger,
		schemaBuilder: sb,
		storeService:  storeService,
		Params:        collections.NewItem(sb, types.ParamsKey, types.ParamsName, codec.CollValue[types.Params](cdc)),

		authority:         authority,
		evmKeeper:         evmKeeper,
		feemarketKeeper:   feemarketKeeper,
		bankKeeper:        bankKeeper,
		accountKeeper:     accountKeeper,
		uregistryKeeper:  uregistryKeeper,
		uvalidatorKeeper: uvalidatorKeeper,

		PendingInbounds: collections.NewKeySet(
			sb,
			types.InboundsKey,
			types.InboundsName,
			collections.StringKey,
		),

		UniversalTx: collections.NewMap(
			sb,
			types.UniversalTxKey,
			types.UniversalTxName,
			collections.StringKey,
			codec.CollValue[types.UniversalTx](cdc),
		),

		ModuleAccountNonce: collections.NewItem(
			sb,
			types.ModuleAccountNonceKey,
			types.ModuleAccountNonceName,
			collections.Uint64Value,
		),

		GasPrices: collections.NewMap(
			sb,
			types.GasPricesKey,
			types.GasPricesName,
			collections.StringKey,
			codec.CollValue[types.GasPrice](cdc),
		),

		ChainMetas: collections.NewMap(
			sb,
			types.ChainMetaKey,
			types.ChainMetasName,
			collections.StringKey,
			codec.CollValue[types.ChainMeta](cdc),
		),

		PendingOutbounds: collections.NewMap(
			sb,
			types.PendingOutboundsKey,
			types.PendingOutboundsName,
			collections.StringKey,
			codec.CollValue[types.PendingOutboundEntry](cdc),
		),
	}

	return k
}

func (k Keeper) Logger() log.Logger {
	return k.logger
}

// InitGenesis initializes the module's state from a genesis state.
func (k *Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) error {

	if err := data.Params.ValidateBasic(); err != nil {
		return err
	}

	// Only deploy factory contracts on fresh genesis, not on import from export.
	// Re-deploying on import would overwrite existing EVM state or cause nonce collisions.
	if !data.Exported {
		deployFactoryEA(ctx, k.evmKeeper)
	}

	if err := k.Params.Set(ctx, data.Params); err != nil {
		return err
	}

	// Restore PendingInbounds
	for _, key := range data.PendingInbounds {
		if err := k.PendingInbounds.Set(ctx, key); err != nil {
			return err
		}
	}

	// Restore UniversalTx
	for _, entry := range data.UniversalTxs {
		if err := k.UniversalTx.Set(ctx, entry.Key, entry.Value); err != nil {
			return err
		}
	}

	// Restore ModuleAccountNonce
	if data.ModuleAccountNonce > 0 {
		if err := k.ModuleAccountNonce.Set(ctx, data.ModuleAccountNonce); err != nil {
			return err
		}
	}

	// Restore GasPrices
	for _, entry := range data.GasPrices {
		if err := k.GasPrices.Set(ctx, entry.Key, entry.Value); err != nil {
			return err
		}
	}

	// Restore ChainMetas
	for _, entry := range data.ChainMetas {
		if err := k.ChainMetas.Set(ctx, entry.Key, entry.Value); err != nil {
			return err
		}
	}

	// Restore PendingOutbounds
	for _, entry := range data.PendingOutbounds {
		if err := k.PendingOutbounds.Set(ctx, entry.OutboundId, entry); err != nil {
			return err
		}
	}

	return nil
}

// ExportGenesis exports the module's state to a genesis state.
func (k *Keeper) ExportGenesis(ctx context.Context) *types.GenesisState {
	params, err := k.Params.Get(ctx)
	if err != nil {
		panic(err)
	}

	// Export PendingInbounds
	var pendingInbounds []string
	err = k.PendingInbounds.Walk(ctx, nil, func(key string) (bool, error) {
		pendingInbounds = append(pendingInbounds, key)
		return false, nil
	})
	if err != nil {
		panic(err)
	}

	// Export UniversalTx
	var universalTxs []types.UniversalTxEntry
	err = k.UniversalTx.Walk(ctx, nil, func(key string, value types.UniversalTx) (bool, error) {
		universalTxs = append(universalTxs, types.UniversalTxEntry{Key: key, Value: value})
		return false, nil
	})
	if err != nil {
		panic(err)
	}

	// Export ModuleAccountNonce
	moduleAccountNonce, err := k.ModuleAccountNonce.Get(ctx)
	if err != nil {
		if !errors.Is(err, collections.ErrNotFound) {
			panic(err)
		}
		moduleAccountNonce = 0
	}

	// Export GasPrices
	var gasPrices []types.GasPriceEntry
	err = k.GasPrices.Walk(ctx, nil, func(key string, value types.GasPrice) (bool, error) {
		gasPrices = append(gasPrices, types.GasPriceEntry{Key: key, Value: value})
		return false, nil
	})
	if err != nil {
		panic(err)
	}

	// Export ChainMetas
	var chainMetas []types.ChainMetaEntry
	err = k.ChainMetas.Walk(ctx, nil, func(key string, value types.ChainMeta) (bool, error) {
		chainMetas = append(chainMetas, types.ChainMetaEntry{Key: key, Value: value})
		return false, nil
	})
	if err != nil {
		panic(err)
	}

	// Export PendingOutbounds
	var pendingOutbounds []types.PendingOutboundEntry
	err = k.PendingOutbounds.Walk(ctx, nil, func(key string, value types.PendingOutboundEntry) (bool, error) {
		pendingOutbounds = append(pendingOutbounds, value)
		return false, nil
	})
	if err != nil {
		panic(err)
	}

	return &types.GenesisState{
		Params:             params,
		PendingInbounds:    pendingInbounds,
		UniversalTxs:       universalTxs,
		ModuleAccountNonce: moduleAccountNonce,
		GasPrices:          gasPrices,
		ChainMetas:         chainMetas,
		PendingOutbounds:   pendingOutbounds,
		Exported:           true,
	}
}

func (k *Keeper) GetUeModuleAddress(ctx context.Context) (common.Address, string) {
	ueModuleAcc := k.accountKeeper.GetModuleAccount(ctx, types.ModuleName) // "ue"
	ueModuleAddr := ueModuleAcc.GetAddress()
	var ethSenderUEAddr common.Address
	copy(ethSenderUEAddr[:], ueModuleAddr.Bytes())

	return ethSenderUEAddr, ethSenderUEAddr.Hex()
}

func (k Keeper) SchemaBuilder() *collections.SchemaBuilder {
	return k.schemaBuilder
}

// GetModuleAccountNonce returns the current module account nonce.
// If not set yet, it safely defaults to 0.
func (k Keeper) GetModuleAccountNonce(ctx sdk.Context) (uint64, error) {
	nonce, err := k.ModuleAccountNonce.Get(ctx)
	if err != nil {
		// If the key is missing, return 0 instead of error
		if errors.Is(err, collections.ErrNotFound) {
			return 0, nil
		}
		return 0, err
	}
	return nonce, nil
}

// IncrementModuleAccountNonce increases the nonce by 1 and stores it back.
func (k Keeper) IncrementModuleAccountNonce(ctx sdk.Context) (uint64, error) {
	nonce, err := k.GetModuleAccountNonce(ctx)
	if err != nil {
		return 0, err
	}
	newNonce := nonce + 1
	if err := k.ModuleAccountNonce.Set(ctx, newNonce); err != nil {
		return 0, err
	}
	return newNonce, nil
}

// SetModuleAccountNonce allows explicitly setting the nonce (optional, for migration or testing).
func (k Keeper) SetModuleAccountNonce(ctx sdk.Context, nonce uint64) error {
	return k.ModuleAccountNonce.Set(ctx, nonce)
}