package ante_test

import (
	"testing"
	"time"

	abci "git.cw.tr/mukan-network/mukan-consensus/abci/types"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	// TODO We don't need to import these API types if we use gogo's registry
	// ref: https://git.cw.tr/mukan-network/mukan-sdk/issues/14647
	_ "cosmossdk.io/api/cosmos/bank/v1beta1"
	_ "cosmossdk.io/api/cosmos/crypto/secp256k1"
	storetypes "cosmossdk.io/store/types"

	"git.cw.tr/mukan-network/mukan-sdk/client"
	"git.cw.tr/mukan-network/mukan-sdk/client/tx"
	cryptotypes "git.cw.tr/mukan-network/mukan-sdk/crypto/types"
	"git.cw.tr/mukan-network/mukan-sdk/runtime"
	"git.cw.tr/mukan-network/mukan-sdk/testutil"
	clitestutil "git.cw.tr/mukan-network/mukan-sdk/testutil/cli"
	"git.cw.tr/mukan-network/mukan-sdk/testutil/testdata"
	_ "git.cw.tr/mukan-network/mukan-sdk/testutil/testdata/testpb"
	sdk "git.cw.tr/mukan-network/mukan-sdk/types"
	moduletestutil "git.cw.tr/mukan-network/mukan-sdk/types/module/testutil"
	"git.cw.tr/mukan-network/mukan-sdk/types/tx/signing"
	"git.cw.tr/mukan-network/mukan-sdk/x/auth"
	"git.cw.tr/mukan-network/mukan-sdk/x/auth/ante"
	antetestutil "git.cw.tr/mukan-network/mukan-sdk/x/auth/ante/testutil"
	authcodec "git.cw.tr/mukan-network/mukan-sdk/x/auth/codec"
	"git.cw.tr/mukan-network/mukan-sdk/x/auth/keeper"
	xauthsigning "git.cw.tr/mukan-network/mukan-sdk/x/auth/signing"
	authtestutil "git.cw.tr/mukan-network/mukan-sdk/x/auth/testutil"
	txtestutil "git.cw.tr/mukan-network/mukan-sdk/x/auth/tx/testutil"
	"git.cw.tr/mukan-network/mukan-sdk/x/auth/types"
	"git.cw.tr/mukan-network/mukan-sdk/x/bank"
)

// TestAccount represents an account used in the tests in x/auth/ante.
type TestAccount struct {
	acc  sdk.AccountI
	priv cryptotypes.PrivKey
}

// AnteTestSuite is a test suite to be used with ante handler tests.
type AnteTestSuite struct {
	anteHandler    sdk.AnteHandler
	ctx            sdk.Context
	clientCtx      client.Context
	txBuilder      client.TxBuilder
	accountKeeper  keeper.AccountKeeper
	bankKeeper     *authtestutil.MockBankKeeper
	txBankKeeper   *txtestutil.MockBankKeeper
	feeGrantKeeper *antetestutil.MockFeegrantKeeper
	encCfg         moduletestutil.TestEncodingConfig
}

func setupSuite(t *testing.T, isCheckTx, enableUnorderedTxs bool) *AnteTestSuite {
	t.Helper()

	suite := &AnteTestSuite{}
	ctrl := gomock.NewController(t)
	suite.bankKeeper = authtestutil.NewMockBankKeeper(ctrl)
	suite.txBankKeeper = txtestutil.NewMockBankKeeper(ctrl)
	suite.feeGrantKeeper = antetestutil.NewMockFeegrantKeeper(ctrl)

	key := storetypes.NewKVStoreKey(types.StoreKey)
	testCtx := testutil.DefaultContextWithDB(t, key, storetypes.NewTransientStoreKey("transient_test"))
	suite.ctx = testCtx.Ctx.WithIsCheckTx(isCheckTx).WithBlockHeight(1) // app.BaseApp.NewContext(isCheckTx, cmtproto.Header{}).WithBlockHeight(1)
	suite.encCfg = moduletestutil.MakeTestEncodingConfig(auth.AppModuleBasic{}, bank.AppModuleBasic{})

	maccPerms := map[string][]string{
		"fee_collector":          nil,
		"mint":                   {"minter"},
		"bonded_tokens_pool":     {"burner", "staking"},
		"not_bonded_tokens_pool": {"burner", "staking"},
		"multiPerm":              {"burner", "minter", "staking"},
		"random":                 {"random"},
	}

	suite.accountKeeper = keeper.NewAccountKeeper(
		suite.encCfg.Codec, runtime.NewKVStoreService(key), types.ProtoBaseAccount, maccPerms, authcodec.NewBech32Codec("cosmos"),
		sdk.Bech32MainPrefix, types.NewModuleAddress("gov").String(), keeper.WithUnorderedTransactions(enableUnorderedTxs),
	)
	suite.accountKeeper.GetModuleAccount(suite.ctx, types.FeeCollectorName)
	err := suite.accountKeeper.Params.Set(suite.ctx, types.DefaultParams())
	require.NoError(t, err)

	// We're using TestMsg encoding in some tests, so register it here.
	suite.encCfg.Amino.RegisterConcrete(&testdata.TestMsg{}, "testdata.TestMsg", nil)
	testdata.RegisterInterfaces(suite.encCfg.InterfaceRegistry)

	suite.clientCtx = client.Context{}.
		WithTxConfig(suite.encCfg.TxConfig).
		WithClient(clitestutil.NewMockCometRPC(abci.ResponseQuery{}))

	anteHandler, err := ante.NewAnteHandler(
		ante.HandlerOptions{
			AccountKeeper:   suite.accountKeeper,
			BankKeeper:      suite.bankKeeper,
			FeegrantKeeper:  suite.feeGrantKeeper,
			SignModeHandler: suite.encCfg.TxConfig.SignModeHandler(),
			SigGasConsumer:  ante.DefaultSigVerificationGasConsumer,
		},
	)

	require.NoError(t, err)
	suite.anteHandler = anteHandler

	suite.txBuilder = suite.clientCtx.TxConfig.NewTxBuilder()

	return suite
}

// SetupTest setups a new test, with new app, context, and anteHandler.
func SetupTestSuite(t *testing.T, isCheckTx bool) *AnteTestSuite {
	t.Helper()
	return setupSuite(t, isCheckTx, false)
}

// SetupTest setups a new test, with new app, context, and anteHandler.
func SetupTestSuiteWithUnordered(t *testing.T, isCheckTx, enableUnorderedTxs bool) *AnteTestSuite {
	t.Helper()
	return setupSuite(t, isCheckTx, enableUnorderedTxs)
}

func (suite *AnteTestSuite) CreateTestAccounts(numAccs int) []TestAccount {
	var accounts []TestAccount

	for i := range numAccs {
		priv, _, addr := testdata.KeyTestPubAddr()
		acc := suite.accountKeeper.NewAccountWithAddress(suite.ctx, addr)
		err := acc.SetAccountNumber(uint64(i + 1000))
		if err != nil {
			panic(err)
		}
		suite.accountKeeper.SetAccount(suite.ctx, acc)
		accounts = append(accounts, TestAccount{acc, priv})
	}

	return accounts
}

// TestCase represents a test case used in test tables.
type TestCase struct {
	desc     string
	malleate func(*AnteTestSuite) TestCaseArgs
	simulate bool
	expPass  bool
	expErr   error
}

type TestCaseArgs struct {
	chainID   string
	accNums   []uint64
	accSeqs   []uint64
	feeAmount sdk.Coins
	gasLimit  uint64
	msgs      []sdk.Msg
	privs     []cryptotypes.PrivKey
}

func (t TestCaseArgs) WithAccountsInfo(accs []TestAccount) TestCaseArgs {
	newT := t
	for _, acc := range accs {
		newT.accNums = append(newT.accNums, acc.acc.GetAccountNumber())
		newT.accSeqs = append(newT.accSeqs, acc.acc.GetSequence())
		newT.privs = append(newT.privs, acc.priv)
	}
	return newT
}

// DeliverMsgs constructs a tx and runs it through the ante handler. This is used to set the context for a test case, for
// example to test for replay protection.
func (suite *AnteTestSuite) DeliverMsgs(t *testing.T, privs []cryptotypes.PrivKey, msgs []sdk.Msg, feeAmount sdk.Coins, gasLimit uint64, accNums, accSeqs []uint64, chainID string, simulate bool) (sdk.Context, error) {
	t.Helper()

	require.NoError(t, suite.txBuilder.SetMsgs(msgs...))
	suite.txBuilder.SetFeeAmount(feeAmount)
	suite.txBuilder.SetGasLimit(gasLimit)

	tx, txErr := suite.CreateTestTx(suite.ctx, privs, accNums, accSeqs, chainID, signing.SignMode_SIGN_MODE_DIRECT)
	require.NoError(t, txErr)
	txBytes, err := suite.clientCtx.TxConfig.TxEncoder()(tx)
	bytesCtx := suite.ctx.WithTxBytes(txBytes)
	require.NoError(t, err)
	return suite.anteHandler(bytesCtx, tx, simulate)
}

func (suite *AnteTestSuite) RunTestCase(t *testing.T, tc TestCase, args TestCaseArgs) {
	t.Helper()

	require.NoError(t, suite.txBuilder.SetMsgs(args.msgs...))
	suite.txBuilder.SetFeeAmount(args.feeAmount)
	suite.txBuilder.SetGasLimit(args.gasLimit)

	// Theoretically speaking, ante handler unit tests should only test
	// ante handlers, but here we sometimes also test the tx creation
	// process.
	tx, txErr := suite.CreateTestTx(suite.ctx, args.privs, args.accNums, args.accSeqs, args.chainID, signing.SignMode_SIGN_MODE_DIRECT)
	txBytes, err := suite.clientCtx.TxConfig.TxEncoder()(tx)
	require.NoError(t, err)
	bytesCtx := suite.ctx.WithTxBytes(txBytes)
	newCtx, anteErr := suite.anteHandler(bytesCtx, tx, tc.simulate)

	if tc.expPass {
		require.NoError(t, txErr)
		require.NoError(t, anteErr)
		require.NotNil(t, newCtx)

		suite.ctx = newCtx
	} else {
		switch {
		case txErr != nil:
			require.Error(t, txErr)
			require.ErrorIs(t, txErr, tc.expErr)

		case anteErr != nil:
			require.Error(t, anteErr)
			require.ErrorIs(t, anteErr, tc.expErr)

		default:
			t.Fatal("expected one of txErr, anteErr to be an error")
		}
	}
}

// CreateTestTx is a helper function to create a tx given multiple inputs.
func (suite *AnteTestSuite) createTx(
	ctx sdk.Context, privs []cryptotypes.PrivKey,
	accNums, accSeqs []uint64,
	chainID string, signMode signing.SignMode, unordered bool, unorderedTimeout time.Time,
) (xauthsigning.Tx, error) {
	suite.txBuilder.SetUnordered(unordered)
	suite.txBuilder.SetTimeoutTimestamp(unorderedTimeout)

	// First round: we gather all the signer infos. We use the "set empty
	// signature" hack to do that.
	var sigsV2 []signing.SignatureV2
	for i, priv := range privs {
		sigV2 := signing.SignatureV2{
			PubKey: priv.PubKey(),
			Data: &signing.SingleSignatureData{
				SignMode:  signMode,
				Signature: nil,
			},
			Sequence: accSeqs[i],
		}

		sigsV2 = append(sigsV2, sigV2)
	}
	err := suite.txBuilder.SetSignatures(sigsV2...)
	if err != nil {
		return nil, err
	}

	// Second round: all signer infos are set, so each signer can sign.
	sigsV2 = []signing.SignatureV2{}
	for i, priv := range privs {
		signerData := xauthsigning.SignerData{
			Address:       sdk.AccAddress(priv.PubKey().Address()).String(),
			ChainID:       chainID,
			AccountNumber: accNums[i],
			Sequence:      accSeqs[i],
			PubKey:        priv.PubKey(),
		}
		sigV2, err := tx.SignWithPrivKey(
			ctx, signMode, signerData,
			suite.txBuilder, priv, suite.clientCtx.TxConfig, accSeqs[i])
		if err != nil {
			return nil, err
		}

		sigsV2 = append(sigsV2, sigV2)
	}
	err = suite.txBuilder.SetSignatures(sigsV2...)
	if err != nil {
		return nil, err
	}

	return suite.txBuilder.GetTx(), nil
}

// CreateTestTx is a helper function to create a tx given multiple inputs.
func (suite *AnteTestSuite) CreateTestTx(
	ctx sdk.Context, privs []cryptotypes.PrivKey,
	accNums, accSeqs []uint64,
	chainID string, signMode signing.SignMode,
) (xauthsigning.Tx, error) {
	return suite.createTx(ctx, privs, accNums, accSeqs, chainID, signMode, false, time.Time{})
}

func (suite *AnteTestSuite) CreateTestUnorderedTx(
	ctx sdk.Context, privs []cryptotypes.PrivKey,
	accNums, accSeqs []uint64,
	chainID string, signMode signing.SignMode,
	unordered bool, unorderedTimeout time.Time,
) (xauthsigning.Tx, error) {
	return suite.createTx(ctx, privs, accNums, accSeqs, chainID, signMode, unordered, unorderedTimeout)
}