package types

import (
	"encoding/json"
	"os"
	"path/filepath"
	"testing"

	abci "github.com/cometbft/cometbft/abci/types"
	dbm "github.com/cosmos/cosmos-db"
	"github.com/stretchr/testify/require"

	"cosmossdk.io/log"
	"cosmossdk.io/store/metrics"
	pruningtypes "cosmossdk.io/store/pruning/types"
	"cosmossdk.io/store/rootmulti"
	storetypes "cosmossdk.io/store/types"

	"github.com/cosmos/cosmos-sdk/baseapp"
)

func useUpgradeLoader(height int64, upgrades *storetypes.StoreUpgrades) func(*baseapp.BaseApp) {
	return func(app *baseapp.BaseApp) {
		app.SetStoreLoader(UpgradeStoreLoader(height, upgrades))
	}
}

func initStore(t *testing.T, db dbm.DB, storeKey string, k, v []byte) {
	t.Helper()

	rs := rootmulti.NewStore(db, log.NewNopLogger(), metrics.NewNoOpMetrics())
	rs.SetPruning(pruningtypes.NewPruningOptions(pruningtypes.PruningNothing))
	key := storetypes.NewKVStoreKey(storeKey)
	rs.MountStoreWithDB(key, storetypes.StoreTypeIAVL, nil)
	err := rs.LoadLatestVersion()
	require.Nil(t, err)
	require.Equal(t, int64(0), rs.LastCommitID().Version)

	// write some data in substore
	kv, _ := rs.GetStore(key).(storetypes.KVStore)
	require.NotNil(t, kv)
	kv.Set(k, v)
	commitID := rs.Commit()
	require.Equal(t, int64(1), commitID.Version)
}

func checkStore(t *testing.T, db dbm.DB, ver int64, storeKey string, k, v []byte) {
	t.Helper()
	rs := rootmulti.NewStore(db, log.NewNopLogger(), metrics.NewNoOpMetrics())
	rs.SetPruning(pruningtypes.NewPruningOptions(pruningtypes.PruningNothing))
	key := storetypes.NewKVStoreKey(storeKey)
	rs.MountStoreWithDB(key, storetypes.StoreTypeIAVL, nil)
	err := rs.LoadLatestVersion()
	require.Nil(t, err)
	require.Equal(t, ver, rs.LastCommitID().Version)

	// query data in substore
	kv, _ := rs.GetStore(key).(storetypes.KVStore)

	require.NotNil(t, kv)
	require.Equal(t, v, kv.Get(k))
}

// Test that we can make commits and then reload old versions.
// Test that LoadLatestVersion actually does.
func TestSetLoader(t *testing.T) {
	upgradeHeight := int64(5)

	// set a temporary home dir
	homeDir := t.TempDir()
	upgradeInfoFilePath := filepath.Join(homeDir, UpgradeInfoFilename)
	upgradeInfo := &Plan{
		Name: "test", Height: upgradeHeight,
	}

	data, err := json.Marshal(upgradeInfo)
	require.NoError(t, err)

	err = os.WriteFile(upgradeInfoFilePath, data, 0o600)
	require.NoError(t, err)

	// make sure it exists before running everything
	_, err = os.Stat(upgradeInfoFilePath)
	require.NoError(t, err)

	cases := map[string]struct {
		setLoader    func(*baseapp.BaseApp)
		origStoreKey string
		loadStoreKey string
	}{
		"don't set loader": {
			setLoader:    nil,
			origStoreKey: "foo",
			loadStoreKey: "foo",
		},
		"rename with inline opts": {
			setLoader: useUpgradeLoader(upgradeHeight, &storetypes.StoreUpgrades{
				Renamed: []storetypes.StoreRename{{
					OldKey: "foo",
					NewKey: "bar",
				}},
			}),
			origStoreKey: "foo",
			loadStoreKey: "bar",
		},
	}

	k := []byte("key")
	v := []byte("value")

	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			// prepare a db with some data
			db := dbm.NewMemDB()

			initStore(t, db, tc.origStoreKey, k, v)

			// load the app with the existing db
			opts := []func(*baseapp.BaseApp){baseapp.SetPruning(pruningtypes.NewPruningOptions(pruningtypes.PruningNothing))}

			logger := log.NewTestLogger(t)

			oldApp := baseapp.NewBaseApp(t.Name(), logger.With("instance", "orig"), db, nil, opts...)
			oldApp.MountStores(storetypes.NewKVStoreKey(tc.origStoreKey))

			err := oldApp.LoadLatestVersion()
			require.Nil(t, err)
			require.Equal(t, int64(1), oldApp.LastBlockHeight())

			for i := int64(2); i <= upgradeHeight-1; i++ {
				_, err = oldApp.FinalizeBlock(&abci.RequestFinalizeBlock{Height: i})
				require.NoError(t, err)
				_, err := oldApp.Commit()
				require.NoError(t, err)
			}

			require.Equal(t, upgradeHeight-1, oldApp.LastBlockHeight())

			if tc.setLoader != nil {
				opts = append(opts, tc.setLoader)
			}

			// load the new newApp with the original newApp db
			newApp := baseapp.NewBaseApp(t.Name(), logger.With("instance", "new"), db, nil, opts...)
			newApp.MountStores(storetypes.NewKVStoreKey(tc.loadStoreKey))

			err = newApp.LoadLatestVersion()
			require.Nil(t, err)

			require.Equal(t, upgradeHeight-1, newApp.LastBlockHeight())

			// "execute" one block
			_, err = newApp.FinalizeBlock(&abci.RequestFinalizeBlock{Height: upgradeHeight})
			require.NoError(t, err)
			_, err = newApp.Commit()
			require.NoError(t, err)
			require.Equal(t, upgradeHeight, newApp.LastBlockHeight())

			// check db is properly updated
			checkStore(t, db, upgradeHeight, tc.loadStoreKey, k, v)
		})
	}
}
