package wallet_test

import (
	"context"
	"fmt"
	"math/big"
	"slices"
	"testing"

	"github.com/btcsuite/btcd/btcutil/base58"
	"github.com/cosmos/go-bip39"
	"github.com/phamminh0811/private-grpc/crypto"
	"github.com/phamminh0811/private-grpc/nockchain"
	"github.com/phamminh0811/private-grpc/wallet"
	"github.com/stretchr/testify/assert"
)

// The entropy, salt and result is taken from "nockchain-wallet keygen" command
func TestKeyGen(t *testing.T) {
	entropyBigInt, isOk := new(big.Int).SetString("29615235796517918707367078072007441124337225858809749976291970867443501879006", 10)
	assert.True(t, isOk)

	entropy := entropyBigInt.Bytes()
	assert.Len(t, entropy, 32)

	saltBigInt, isOk := new(big.Int).SetString("212311808188922973323281316240858086116", 10)
	assert.True(t, isOk)

	salt := saltBigInt.Bytes()
	assert.Len(t, salt, 16)

	argonBytes := crypto.DeriveKey(0, entropy[:], salt[:], nil, nil, 6, 786432, 4, 32)
	slices.Reverse(argonBytes)
	mnemonic, err := bip39.NewMnemonic(argonBytes)
	assert.NoError(t, err)
	assert.Equal(t, mnemonic, "brass vacuum stairs hurt brisk govern describe enforce fly exact rescue capable belt flavor lottery sauce easy frame orange legal injury border obey novel")

	masterKey, err := crypto.MasterKeyFromSeed(mnemonic)
	assert.NoError(t, err)

	assert.Equal(t,
		base58.Encode(masterKey.PublicKey),
		"39DL6YA1kSRCKMjzpFEtC8rmnxVuseUrP2LnViwY7YEhZYZkX2HmnAZ63Uwy1DwuXstmF1VeJDucg719xw49j9CKL3bsKq3A6SZN918CowcgQroHsgohj7dYgpGRWk41s42F",
	)
	assert.Equal(t,
		base58.Encode(masterKey.PrivateKey),
		"4SyUrsbGKPRknzvGakWmFbYefzHzb1r4LUmJpQD8WPcR",
	)
	assert.Equal(t,
		base58.Encode(masterKey.ChainCode),
		"58SARPmADHvUcpq7XfBoCgwzy5QC8Kb3JrezpHqA85x2",
	)

	// assert import priv/pubkey
	privBytes := append([]byte{0x00}, masterKey.PrivateKey...)
	importPrivKey := crypto.SerializeExtend(masterKey.ChainCode, privBytes, crypto.PrivateKeyStart)
	assert.Len(t, importPrivKey, 82)
	importPubKey := crypto.SerializeExtend(masterKey.ChainCode, masterKey.PublicKey, crypto.PublicKeyStart)
	assert.Len(t, importPubKey, 145)

	assert.Equal(t,
		base58.Encode(importPrivKey),
		"zprv2CyrSHEkzQzu4HCtJRFiP4t2rVMauZwLfJDFrNbqS8Pz3nsmXy5bAUx2HYUykaMuU4MiQTHsDcKYjLCjrPfpceNT9XBHgx1pUjKzBrF6Wdo",
	)
	assert.Equal(t,
		base58.Encode(importPubKey),
		"zpubUQwNTNE3hsCkMpBnD37W5QirkyVryokAVPLnPin1c6M13RRsq3yEJbwp5ies6qXF6DvJq5Woxw6ygT53PSVrmrsQgtHhbfMEixKNFm7qb4mELhpyoovpFEV1YPHFZx4xQGYBNF6qvXU6AHNh4TLrUdkYAdXKS2J5rPiSVPrXKGo8fLG6ZBCGBjJfPcwDb2VEJC",
	)
}

func TestImportKey(t *testing.T) {
	type Input struct {
		req        *nockchain.ImportKeysRequest
		expectResp *nockchain.ImportKeysResponse
		isErr      bool
		errStr     string
	}

	correctImportPrivKey := base58.Decode("zprv2CyrSHEkzQzu4HCtJRFiP4t2rVMauZwLfJDFrNbqS8Pz3nsmXy5bAUx2HYUykaMuU4MiQTHsDcKYjLCjrPfpceNT9XBHgx1pUjKzBrF6Wdo")
	invalidImportPrivKeyPrefix := make([]byte, 82)
	copy(invalidImportPrivKeyPrefix[:], correctImportPrivKey)
	invalidImportPrivKeyPrefix[45] = 0x01

	invalidImportPrivKeyChecksum := make([]byte, 82)
	copy(invalidImportPrivKeyChecksum[:], correctImportPrivKey)
	copy(invalidImportPrivKeyChecksum[78:], []byte{1, 2, 3, 4})

	correctImportPubkey := base58.Decode("zpubUQwNTNE3hsCkMpBnD37W5QirkyVryokAVPLnPin1c6M13RRsq3yEJbwp5ies6qXF6DvJq5Woxw6ygT53PSVrmrsQgtHhbfMEixKNFm7qb4mELhpyoovpFEV1YPHFZx4xQGYBNF6qvXU6AHNh4TLrUdkYAdXKS2J5rPiSVPrXKGo8fLG6ZBCGBjJfPcwDb2VEJC")
	invalidImportPubkeyChecksum := make([]byte, 145)
	copy(invalidImportPubkeyChecksum[:], correctImportPubkey)
	copy(invalidImportPubkeyChecksum[141:], []byte{1, 2, 3, 4})

	response := &nockchain.ImportKeysResponse{
		PublicKey:        "39DL6YA1kSRCKMjzpFEtC8rmnxVuseUrP2LnViwY7YEhZYZkX2HmnAZ63Uwy1DwuXstmF1VeJDucg719xw49j9CKL3bsKq3A6SZN918CowcgQroHsgohj7dYgpGRWk41s42F",
		PrivateKey:       "4SyUrsbGKPRknzvGakWmFbYefzHzb1r4LUmJpQD8WPcR",
		ChainCode:        "58SARPmADHvUcpq7XfBoCgwzy5QC8Kb3JrezpHqA85x2",
		ImportPrivateKey: base58.Encode(correctImportPrivKey),
		ImportPublicKey:  base58.Encode(correctImportPubkey),
	}

	responseReadOnly := &nockchain.ImportKeysResponse{
		PublicKey:        "39DL6YA1kSRCKMjzpFEtC8rmnxVuseUrP2LnViwY7YEhZYZkX2HmnAZ63Uwy1DwuXstmF1VeJDucg719xw49j9CKL3bsKq3A6SZN918CowcgQroHsgohj7dYgpGRWk41s42F",
		PrivateKey:       "",
		ChainCode:        "58SARPmADHvUcpq7XfBoCgwzy5QC8Kb3JrezpHqA85x2",
		ImportPrivateKey: "",
		ImportPublicKey:  base58.Encode(correctImportPubkey),
	}
	inputs := []Input{
		// case invalid type
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "",
				ImportType: nockchain.ImportType_UNDEFINED,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid import type",
		},
		// case invalid extended key
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "some wrong string",
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid extended key",
		},
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "zprv wrong priv import key length",
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid extended private key length",
		},
		{
			req: &nockchain.ImportKeysRequest{
				Key:        base58.Encode(invalidImportPrivKeyPrefix),
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid private key prefix at byte 45",
		},
		{
			req: &nockchain.ImportKeysRequest{
				Key:        base58.Encode(invalidImportPrivKeyChecksum),
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid checksum",
		},
		// case success import priv key
		{
			req: &nockchain.ImportKeysRequest{
				Key:        base58.Encode(correctImportPrivKey),
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: response,
			isErr:      false,
			errStr:     "",
		},
		// case invalid import pub key
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "zpub wrong public import key length",
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid extended public key length",
		},
		{
			req: &nockchain.ImportKeysRequest{
				Key:        base58.Encode(invalidImportPubkeyChecksum),
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid checksum",
		},
		// case success import pub key
		{
			req: &nockchain.ImportKeysRequest{
				Key:        base58.Encode(correctImportPubkey),
				ImportType: nockchain.ImportType_EXTENDED_KEY,
			},
			expectResp: responseReadOnly,
			isErr:      false,
			errStr:     "",
		},
		// case missing chaincode when import master privkey
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "4SyUrsbGKPRknzvGakWmFbYefzHzb1r4LUmJpQD8WPcR",
				ImportType: nockchain.ImportType_MASTER_PRIVKEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "master key must be in [chain_code],[key] format",
		},
		// case invalid length
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "abcdxyz,4SyUrsbGKPRknzvGakWmFbYefzHzb1r4LUmJpQD8WPcR",
				ImportType: nockchain.ImportType_MASTER_PRIVKEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid chain code length",
		},
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "58SARPmADHvUcpq7XfBoCgwzy5QC8Kb3JrezpHqA85x2,abcdxyz",
				ImportType: nockchain.ImportType_MASTER_PRIVKEY,
			},
			expectResp: nil,
			isErr:      true,
			errStr:     "invalid priv key length",
		},
		// case success import master privkey
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "58SARPmADHvUcpq7XfBoCgwzy5QC8Kb3JrezpHqA85x2,4SyUrsbGKPRknzvGakWmFbYefzHzb1r4LUmJpQD8WPcR",
				ImportType: nockchain.ImportType_MASTER_PRIVKEY,
			},
			expectResp: response,
			isErr:      false,
			errStr:     "",
		},
		// case success import seed
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "brass vacuum stairs hurt brisk govern describe enforce fly exact rescue capable belt flavor lottery sauce easy frame orange legal injury border obey novel",
				ImportType: nockchain.ImportType_SEEDPHRASE,
			},
			expectResp: response,
			isErr:      false,
			errStr:     "",
		},
		// case sucess import pubkey
		{
			req: &nockchain.ImportKeysRequest{
				Key:        "39DL6YA1kSRCKMjzpFEtC8rmnxVuseUrP2LnViwY7YEhZYZkX2HmnAZ63Uwy1DwuXstmF1VeJDucg719xw49j9CKL3bsKq3A6SZN918CowcgQroHsgohj7dYgpGRWk41s42F",
				ImportType: nockchain.ImportType_WATCH_ONLY,
			},
			expectResp: &nockchain.ImportKeysResponse{
				PublicKey: "39DL6YA1kSRCKMjzpFEtC8rmnxVuseUrP2LnViwY7YEhZYZkX2HmnAZ63Uwy1DwuXstmF1VeJDucg719xw49j9CKL3bsKq3A6SZN918CowcgQroHsgohj7dYgpGRWk41s42F",
			},
			isErr:  false,
			errStr: "",
		},
	}

	for _, input := range inputs {
		nc, err := wallet.NewNockchainClient("nockchain-api.zorp.io:443")
		assert.NoError(t, err)

		handler := wallet.NewGprcHandler(*nc)

		resp, err := handler.ImportKeys(context.Background(), input.req)
		if input.isErr {
			assert.ErrorContains(t, err, input.errStr)
		} else {
			assert.NoError(t, err)
			assert.Equal(t, resp, input.expectResp)
		}
	}
}

// This test should be run with timeout 120s
func TestScan(t *testing.T) {
	// some random seed we take to get empty notes
	seed1 := "foster chicken claw fade income frown junior abandon price lesson mango wrap dry clay loyal camera caught during property useless puppy royal soccer arm"
	seed2 := "brass vacuum stairs hurt brisk govern describe enforce fly exact rescue capable belt flavor lottery sauce easy frame orange legal injury border obey novel"

	masterKey1, err := crypto.MasterKeyFromSeed(seed1)
	assert.NoError(t, err)

	nc, err := wallet.NewNockchainClient("nockchain-api.zorp.io:443")
	assert.NoError(t, err)

	masterKey1Scan, err := nc.WalletGetBalance(base58.Encode(masterKey1.PublicKey))
	assert.NoError(t, err)
	assert.Empty(t, masterKey1Scan.Notes)

	masterKey2, err := crypto.MasterKeyFromSeed(seed2)
	assert.NoError(t, err)

	masterKey2Scan, err := nc.WalletGetBalance(base58.Encode(masterKey2.PublicKey))
	assert.NoError(t, err)
	assert.Len(t, masterKey2Scan.Notes, 1)
	assert.Equal(t, masterKey2Scan.Notes[0].Note.Assets.Value, uint64(100000))
	assert.Equal(t, masterKey2Scan.Notes[0].Note.OriginPage.Value, uint64(35054))
}

// This test should be run with timeout 120s
func TestFullFlow(t *testing.T) {
	mnemonic1 := "rail nurse smile angle uphold gun kitten spoon quick frozen trigger cable decorate episode blame tray off bag arena taxi approve breeze job letter"
	masterKey1, err := crypto.MasterKeyFromSeed(mnemonic1)
	assert.NoError(t, err)

	mnemonic2 := "brass vacuum stairs hurt brisk govern describe enforce fly exact rescue capable belt flavor lottery sauce easy frame orange legal injury border obey novel"
	masterKey2, err := crypto.MasterKeyFromSeed(mnemonic2)
	assert.NoError(t, err)

	inputName := "[4taoqkpysafnp64WBQyzHDKVrqkMeNrdAiVSbWdzZmj7yQYZgQtCq4W 9cjUFbdtaFHeXNWCAKjsTphBchHmCoUU6a1aDbJAFz9qHqeG8osh4wF]"

	nc, err := wallet.NewNockchainClient("nockchain-api.zorp.io:443")
	assert.NoError(t, err)

	masterKeyScan, err := nc.WalletGetBalance(base58.Encode(masterKey1.PublicKey))
	assert.NoError(t, err)

	var note *nockchain.NockchainNote
	for _, balanceEntry := range masterKeyScan.Notes {
		firstName := crypto.Tip5HashToBase58([5]uint64{
			balanceEntry.Name.First.Belt_1.Value,
			balanceEntry.Name.First.Belt_2.Value,
			balanceEntry.Name.First.Belt_3.Value,
			balanceEntry.Name.First.Belt_4.Value,
			balanceEntry.Name.First.Belt_5.Value,
		})
		lastName := crypto.Tip5HashToBase58([5]uint64{
			balanceEntry.Name.Last.Belt_1.Value,
			balanceEntry.Name.Last.Belt_2.Value,
			balanceEntry.Name.Last.Belt_3.Value,
			balanceEntry.Name.Last.Belt_4.Value,
			balanceEntry.Name.Last.Belt_5.Value,
		})
		nname := "[" + firstName + " " + lastName + "]"
		if nname == inputName {
			nnote := wallet.ParseBalanceEntry(balanceEntry)
			note = &nnote
			break
		}
	}
	assert.NotNil(t, note)

	parentHash, err := wallet.HashNote(note)
	assert.NoError(t, err)

	gift := uint64(100000)
	seed1 := &nockchain.NockchainSeed{
		OutputSource: nil,
		Recipient: &nockchain.NockchainLock{
			KeysRequired: 1,
			Pubkeys:      []string{base58.Encode(masterKey2.PublicKey)},
		},
		TimelockIntent: nil,
		Gift:           gift,
		ParentHash:     crypto.Tip5HashToBase58(parentHash),
	}

	gift = note.Asset - gift - 100
	seed2 := &nockchain.NockchainSeed{
		OutputSource: nil,
		Recipient: &nockchain.NockchainLock{
			KeysRequired: 1,
			Pubkeys:      []string{base58.Encode(masterKey1.PublicKey)},
		},
		TimelockIntent: nil,
		Gift:           gift,
		ParentHash:     crypto.Tip5HashToBase58(parentHash),
	}

	spend := nockchain.NockchainSpend{
		Signatures: nil,
		Seeds: []*nockchain.NockchainSeed{
			seed1, seed2,
		},
		Fee: 100,
	}

	msg, err := wallet.HashMsg(&spend)
	assert.NoError(t, err)

	chalT8, sigT8, err := wallet.ComputeSig(*masterKey1, msg)
	assert.NoError(t, err)

	spend.Signatures = []*nockchain.NockchainSignature{
		{
			Pubkey: base58.Encode(masterKey1.PublicKey),
			Chal:   chalT8[:],
			Sig:    sigT8[:],
		},
	}

	input := nockchain.NockchainInput{
		Name: &nockchain.NockchainName{
			First: "4taoqkpysafnp64WBQyzHDKVrqkMeNrdAiVSbWdzZmj7yQYZgQtCq4W",
			Last:  "9cjUFbdtaFHeXNWCAKjsTphBchHmCoUU6a1aDbJAFz9qHqeG8osh4wF",
		},
		Note:  note,
		Spend: &spend,
	}

	id, err := wallet.ComputeTxId([]*nockchain.NockchainInput{&input}, &nockchain.TimelockRange{
		Min: nil,
		Max: nil,
	}, 100)
	assert.NoError(t, err)

	rawTx := nockchain.RawTx{
		TxId:   crypto.Tip5HashToBase58(id),
		Inputs: []*nockchain.NockchainInput{&input},
		TimelockRange: &nockchain.TimelockRange{
			Min: nil,
			Max: nil,
		},
		TotalFees: 100,
	}
	resp, err := nc.WalletSendTransaction(&rawTx)
	assert.NoError(t, err)
	assert.Equal(t, resp.Result, &nockchain.WalletSendTransactionResponse_Ack{
		Ack: &nockchain.Acknowledged{},
	})

	txAcceptedResp, err := nc.TxAccepted(rawTx.TxId)
	assert.NoError(t, err)
	assert.Equal(t, txAcceptedResp.Result, &nockchain.TransactionAcceptedResponse_Accepted{
		Accepted: true,
	})
}

// This test should be run with timeout 300s
func TestCreateTx(t *testing.T) {
	seed1 := "pledge vessel toilet sunny hockey skirt spend wire disorder attitude crumble lecture problem bundle bone rather address over suit ancient primary gospel silent repair"
	masterKey1, err := crypto.MasterKeyFromSeed(seed1)
	assert.NoError(t, err)

	seed2 := "brass vacuum stairs hurt brisk govern describe enforce fly exact rescue capable belt flavor lottery sauce easy frame orange legal injury border obey novel"
	masterKey2, err := crypto.MasterKeyFromSeed(seed2)
	assert.NoError(t, err)

	nc, err := wallet.NewNockchainClient("nockchain-api.zorp.io:443")
	assert.NoError(t, err)
	handler := wallet.NewGprcHandler(*nc)

	inputs := "[CxiTK4HjqPRebUkoy6rH89ZcNGuH4goHkmKgmgCJxZEFS2C3qrDoh4y 91z5muQKgHZDcChdnCSTEtvuN8dbbXp5wzNs5xrCkce2YvSX1q6fu3d],[CxiTK4HjqPRebUkoy6rH89ZcNGuH4goHkmKgmgCJxZEFS2C3qrDoh4y 3a9VWigUM1TuS79yM4dAqkiUg6WJMUPGwVKJpQUHCLJcFZsCWM2q6pF],[CxiTK4HjqPRebUkoy6rH89ZcNGuH4goHkmKgmgCJxZEFS2C3qrDoh4y 9RkPvHMtuYtjV2qvB86u7zcnr26SdVcKzFnHgfhUSEG1W3FgKgLpbBm]"
	recipients := fmt.Sprintf("%s,%s,%s", base58.Encode(masterKey2.PublicKey), base58.Encode(masterKey2.PublicKey), base58.Encode(masterKey2.PublicKey))
	gifts := "100,100,100"
	fee := 100

	req := &nockchain.CreateTxRequest{
		Names:          inputs,
		Recipients:     recipients,
		Gifts:          gifts,
		Fee:            uint64(fee),
		IsMasterKey:    true,
		Key:            base58.Encode(masterKey1.PrivateKey),
		ChainCode:      base58.Encode(masterKey1.ChainCode),
		Index:          0,
		Hardened:       false,
		TimelockIntent: nil,
	}

	res, err := handler.CreateTx(context.Background(), req)
	assert.NoError(t, err)

	// the result is taken from create-tx scripts
	assert.Equal(t, res.RawTx.TxId, "8gjTa6H1MrKXPWgNJCF9fsYE7PUfqhYzxetUoTftz7zyHCv24ZYHDM3")
}