package simulation

import (
	"bytes"
	"fmt"

	"github.com/cosmos/cosmos-sdk/baseapp"
	sdk "github.com/cosmos/cosmos-sdk/types"
	"github.com/cosmos/cosmos-sdk/x/auth"
	"github.com/cosmos/cosmos-sdk/x/bank"
	"github.com/cosmos/cosmos-sdk/x/distribution"
	"github.com/cosmos/cosmos-sdk/x/mock/simulation"
	"github.com/cosmos/cosmos-sdk/x/stake"
	"github.com/cosmos/cosmos-sdk/x/stake/keeper"
	stakeTypes "github.com/cosmos/cosmos-sdk/x/stake/types"
	abci "github.com/tendermint/tendermint/abci/types"
)

// AllInvariants runs all invariants of the stake module.
// Currently: total supply, positive power
func AllInvariants(ck bank.Keeper, k stake.Keeper,
	f auth.FeeCollectionKeeper, d distribution.Keeper,
	am auth.AccountKeeper) simulation.Invariant {

	return func(app *baseapp.BaseApp) error {
		err := SupplyInvariants(ck, k, f, d, am)(app)
		if err != nil {
			return err
		}

		err = PositivePowerInvariant(k)(app)
		if err != nil {
			return err
		}

		err = ValidatorSetInvariant(k)(app)
		return err
	}
}

// SupplyInvariants checks that the total supply reflects all held loose tokens, bonded tokens, and unbonding delegations
// nolint: unparam
func SupplyInvariants(ck bank.Keeper, k stake.Keeper,
	f auth.FeeCollectionKeeper, d distribution.Keeper, am auth.AccountKeeper) simulation.Invariant {
	return func(app *baseapp.BaseApp) error {
		ctx := app.NewContext(false, abci.Header{})
		pool := k.GetPool(ctx)

		loose := sdk.ZeroDec()
		bonded := sdk.ZeroDec()
		am.IterateAccounts(ctx, func(acc auth.Account) bool {
			loose = loose.Add(sdk.NewDecFromInt(acc.GetCoins().AmountOf(stakeTypes.DefaultBondDenom)))
			return false
		})
		k.IterateUnbondingDelegations(ctx, func(_ int64, ubd stake.UnbondingDelegation) bool {
			loose = loose.Add(sdk.NewDecFromInt(ubd.Balance.Amount))
			return false
		})
		k.IterateValidators(ctx, func(_ int64, validator sdk.Validator) bool {
			switch validator.GetStatus() {
			case sdk.Bonded:
				bonded = bonded.Add(validator.GetPower())
			case sdk.Unbonding:
				loose = loose.Add(validator.GetTokens())
			case sdk.Unbonded:
				loose = loose.Add(validator.GetTokens())
			}
			return false
		})

		feePool := d.GetFeePool(ctx)

		// add outstanding fees
		loose = loose.Add(sdk.NewDecFromInt(f.GetCollectedFees(ctx).AmountOf(stakeTypes.DefaultBondDenom)))

		// add community pool
		loose = loose.Add(feePool.CommunityPool.AmountOf(stakeTypes.DefaultBondDenom))

		// add validator distribution pool
		loose = loose.Add(feePool.ValPool.AmountOf(stakeTypes.DefaultBondDenom))

		// add validator distribution commission and yet-to-be-withdrawn-by-delegators
		d.IterateValidatorDistInfos(ctx,
			func(_ int64, distInfo distribution.ValidatorDistInfo) (stop bool) {
				loose = loose.Add(distInfo.DelPool.AmountOf(stakeTypes.DefaultBondDenom))
				loose = loose.Add(distInfo.ValCommission.AmountOf(stakeTypes.DefaultBondDenom))
				return false
			},
		)

		// Loose tokens should equal coin supply plus unbonding delegations
		// plus tokens on unbonded validators
		if !pool.LooseTokens.Equal(loose) {
			return fmt.Errorf("loose token invariance:\n\tpool.LooseTokens: %v"+
				"\n\tsum of account tokens: %v", pool.LooseTokens, loose)
		}

		// Bonded tokens should equal sum of tokens with bonded validators
		if !pool.BondedTokens.Equal(bonded) {
			return fmt.Errorf("bonded token invariance:\n\tpool.BondedTokens: %v"+
				"\n\tsum of account tokens: %v", pool.BondedTokens, bonded)
		}

		return nil
	}
}

// PositivePowerInvariant checks that all stored validators have > 0 power.
func PositivePowerInvariant(k stake.Keeper) simulation.Invariant {
	return func(app *baseapp.BaseApp) error {
		ctx := app.NewContext(false, abci.Header{})

		iterator := k.ValidatorsPowerStoreIterator(ctx)
		pool := k.GetPool(ctx)

		for ; iterator.Valid(); iterator.Next() {
			validator, found := k.GetValidator(ctx, iterator.Value())
			if !found {
				panic(fmt.Sprintf("validator record not found for address: %X\n", iterator.Value()))
			}

			powerKey := keeper.GetValidatorsByPowerIndexKey(validator, pool)

			if !bytes.Equal(iterator.Key(), powerKey) {
				return fmt.Errorf("power store invariance:\n\tvalidator.Power: %v"+
					"\n\tkey should be: %v\n\tkey in store: %v", validator.GetPower(), powerKey, iterator.Key())
			}
		}
		iterator.Close()
		return nil
	}
}

// ValidatorSetInvariant checks equivalence of Tendermint validator set and SDK validator set
func ValidatorSetInvariant(k stake.Keeper) simulation.Invariant {
	return func(app *baseapp.BaseApp) error {
		// TODO
		return nil
	}
}