package key

import (
	"bytes"
	"io"
	"os"
	"path/filepath"
	"testing"

	"github.com/nspcc-dev/neo-go/cli/input"
	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
	"github.com/nspcc-dev/neo-go/pkg/wallet"
	"github.com/nspcc-dev/neofs-node/cmd/neofs-cli/internal/commonflags"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/require"
	"golang.org/x/term"
)

func Test_getOrGenerate(t *testing.T) {
	dir := t.TempDir()

	wallPath := filepath.Join(dir, "wallet.json")
	w, err := wallet.NewWallet(wallPath)
	require.NoError(t, err)

	acc1, err := wallet.NewAccount()
	require.NoError(t, err)
	require.NoError(t, acc1.Encrypt("pass", keys.NEP2ScryptParams()))
	w.AddAccount(acc1)

	acc2, err := wallet.NewAccount()
	require.NoError(t, err)
	require.NoError(t, acc2.Encrypt("pass", keys.NEP2ScryptParams()))
	acc2.Default = true
	w.AddAccount(acc2)
	require.NoError(t, w.Save())
	w.Close()

	keyPath := filepath.Join(dir, "binary.key")
	rawKey, err := keys.NewPrivateKey()
	require.NoError(t, err)
	require.NoError(t, os.WriteFile(keyPath, rawKey.Bytes(), os.ModePerm))

	wifKey, err := keys.NewPrivateKey()
	require.NoError(t, err)

	nep2Key, err := keys.NewPrivateKey()
	require.NoError(t, err)
	nep2, err := keys.NEP2Encrypt(nep2Key, "pass", keys.NEP2ScryptParams())
	require.NoError(t, err)

	in := bytes.NewBuffer(nil)
	input.Terminal = term.NewTerminal(input.ReadWriter{
		Reader: in,
		Writer: io.Discard,
	}, "")

	checkKeyError(t, filepath.Join(dir, "badfile"), ErrInvalidKey)

	t.Run("wallet", func(t *testing.T) {
		checkKeyError(t, wallPath, ErrInvalidPassword)

		in.WriteString("invalid\r")
		checkKeyError(t, wallPath, ErrInvalidPassword)

		in.WriteString("pass\r")
		checkKey(t, wallPath, acc2.PrivateKey()) // default account

		viper.Set(commonflags.Account, acc1.Address)
		in.WriteString("pass\r")
		checkKey(t, wallPath, acc1.PrivateKey())

		viper.Set(commonflags.Account, "not an address")
		checkKeyError(t, wallPath, ErrInvalidAddress)

		acc, err := wallet.NewAccount()
		require.NoError(t, err)
		viper.Set(commonflags.Account, acc.Address)
		checkKeyError(t, wallPath, ErrInvalidAddress)
	})

	t.Run("WIF", func(t *testing.T) {
		checkKey(t, wifKey.WIF(), wifKey)
	})

	t.Run("NEP-2", func(t *testing.T) {
		checkKeyError(t, nep2, ErrInvalidPassword)

		in.WriteString("invalid\r")
		checkKeyError(t, nep2, ErrInvalidPassword)

		in.WriteString("pass\r")
		checkKey(t, nep2, nep2Key)

		t.Run("password from config", func(t *testing.T) {
			viper.Set("password", "invalid")
			in.WriteString("pass\r")
			checkKeyError(t, nep2, ErrInvalidPassword)

			viper.Set("password", "pass")
			in.WriteString("invalid\r")
			checkKey(t, nep2, nep2Key)
		})
	})

	t.Run("raw key", func(t *testing.T) {
		checkKey(t, keyPath, rawKey)
	})

	t.Run("generate", func(t *testing.T) {
		viper.Set(commonflags.GenerateKey, true)
		actual, err := getOrGenerate()
		require.NoError(t, err)
		require.NotNil(t, actual)
		for _, p := range []*keys.PrivateKey{nep2Key, rawKey, wifKey, acc1.PrivateKey(), acc2.PrivateKey()} {
			require.NotEqual(t, p, actual, "expected new key to be generated")
		}
	})
}

func checkKeyError(t *testing.T, desc string, err error) {
	viper.Set(commonflags.WalletPath, desc)
	_, actualErr := getOrGenerate()
	require.ErrorIs(t, actualErr, err)
}

func checkKey(t *testing.T, desc string, expected *keys.PrivateKey) {
	viper.Set(commonflags.WalletPath, desc)
	actual, err := getOrGenerate()
	require.NoError(t, err)
	require.Equal(t, &expected.PrivateKey, actual)
}