package wallet

import (
	context "context"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"slices"
	"strconv"
	"strings"

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

type GprcHandler struct {
	nockchain.UnimplementedWalletServiceServer
	client NockchainClient
}

func NewGprcHandler(client NockchainClient) GprcHandler {
	return GprcHandler{
		client: client,
	}
}
func (h *GprcHandler) Keygen(ctx context.Context, req *nockchain.KeygenRequest) (*nockchain.KeygenResponse, error) {
	var entropy [32]byte
	_, err := rand.Read(entropy[:]) // Fill the slice with random bytes
	if err != nil {
		return nil, err
	}

	var salt [16]byte
	_, err = rand.Read(salt[:])
	if err != nil {
		return nil, err
	}

	argonBytes := crypto.DeriveKey(0, entropy[:], salt[:], nil, nil, 6, 786432, 4, 32)
	slices.Reverse(argonBytes)
	mnemonic, err := bip39.NewMnemonic(argonBytes)
	if err != nil {
		return nil, fmt.Errorf("failed to generate mnemonic: %v", err)
	}
	masterKey, err := crypto.MasterKeyFromSeed(mnemonic)
	if err != nil {
		return nil, err
	}
	privBytes := append([]byte{0x00}, masterKey.PrivateKey...)
	return &nockchain.KeygenResponse{
		Seed:             mnemonic,
		PrivateKey:       base58.Encode(masterKey.PrivateKey),
		PublicKey:        base58.Encode(masterKey.PublicKey),
		ChainCode:        base58.Encode(masterKey.ChainCode),
		ImportPrivateKey: base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, privBytes, crypto.PrivateKeyStart)),
		ImportPublicKey:  base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, masterKey.PublicKey, crypto.PublicKeyStart)),
	}, nil
}

func (h *GprcHandler) ImportKeys(ctx context.Context, req *nockchain.ImportKeysRequest) (*nockchain.ImportKeysResponse, error) {
	switch req.ImportType {
	case nockchain.ImportType_UNDEFINED:
		return nil, fmt.Errorf("invalid import type")
	case nockchain.ImportType_EXTENDED_KEY:
		// metadata layout: [version][depth][parent-fp][index][chain-code][key-data][checksum]
		data := base58.Decode(req.Key)
		switch {
		case strings.HasPrefix(req.Key, "zprv"):
			if len(data) != 82 {
				return nil, fmt.Errorf("invalid extended private key length: %d (expected 82)", len(data))
			}
			if data[45] != 0x00 {
				return nil, fmt.Errorf("invalid private key prefix at byte 45: 0x%02x (expected 0x00)", data[45])
			}
			hash := sha256.Sum256(data[:78])
			hash = sha256.Sum256(hash[:])
			if !slices.Equal(hash[:4], data[78:]) {
				return nil, fmt.Errorf("invalid checksum")
			}
			chainCode := make([]byte, 32)
			copy(chainCode, data[13:45])
			privateKey := make([]byte, 32)
			copy(privateKey, data[46:78])
			masterKey, err := crypto.MasterKeyFromPrivKey(chainCode, privateKey)
			if err != nil {
				return nil, err
			}
			privBytes := append([]byte{0x00}, masterKey.PrivateKey...)
			return &nockchain.ImportKeysResponse{
				Seed:             "",
				PrivateKey:       base58.Encode(masterKey.PrivateKey),
				PublicKey:        base58.Encode(masterKey.PublicKey),
				ChainCode:        base58.Encode(masterKey.ChainCode),
				ImportPrivateKey: base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, privBytes, crypto.PrivateKeyStart)),
				ImportPublicKey:  base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, masterKey.PublicKey, crypto.PublicKeyStart)),
			}, nil
		case strings.HasPrefix(req.Key, "zpub"):
			if len(data) != 145 {
				return nil, fmt.Errorf("invalid extended public key length: %d (expected 145)", len(data))
			}

			hash := sha256.Sum256(data[:141])
			hash = sha256.Sum256(hash[:])
			if !slices.Equal(hash[:4], data[141:]) {
				return nil, fmt.Errorf("invalid checksum")
			}

			chainCode := make([]byte, 32)
			copy(chainCode, data[13:45])
			publicKey := make([]byte, 97)
			copy(publicKey, data[45:141])
			return &nockchain.ImportKeysResponse{
				Seed:             "",
				PrivateKey:       "",
				PublicKey:        base58.Encode(publicKey),
				ChainCode:        base58.Encode(chainCode),
				ImportPrivateKey: "",
				ImportPublicKey:  base58.Encode(crypto.SerializeExtend(chainCode, publicKey, crypto.PublicKeyStart)),
			}, nil
		default:
			return nil, fmt.Errorf("invalid extended key")
		}
	case nockchain.ImportType_MASTER_PRIVKEY:
		splits := strings.Split(req.Key, ",")
		if len(splits) != 2 {
			return nil, fmt.Errorf("master key must be in [chain_code],[key] format")
		}
		chainCode := base58.Decode(splits[0])
		key := base58.Decode(splits[1])
		masterKey, err := crypto.MasterKeyFromPrivKey(chainCode, key)
		if err != nil {
			return nil, err
		}
		privBytes := append([]byte{0x00}, masterKey.PrivateKey...)
		return &nockchain.ImportKeysResponse{
			Seed:             "",
			PrivateKey:       base58.Encode(masterKey.PrivateKey),
			PublicKey:        base58.Encode(masterKey.PublicKey),
			ChainCode:        base58.Encode(masterKey.ChainCode),
			ImportPrivateKey: base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, privBytes, crypto.PrivateKeyStart)),
			ImportPublicKey:  base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, masterKey.PublicKey, crypto.PublicKeyStart)),
		}, nil
	case nockchain.ImportType_SEEDPHRASE:
		masterKey, err := crypto.MasterKeyFromSeed(req.Key)
		if err != nil {
			return nil, err
		}
		privBytes := append([]byte{0x00}, masterKey.PrivateKey...)
		return &nockchain.ImportKeysResponse{
			Seed:             "",
			PrivateKey:       base58.Encode(masterKey.PrivateKey),
			PublicKey:        base58.Encode(masterKey.PublicKey),
			ChainCode:        base58.Encode(masterKey.ChainCode),
			ImportPrivateKey: base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, privBytes, crypto.PrivateKeyStart)),
			ImportPublicKey:  base58.Encode(crypto.SerializeExtend(masterKey.ChainCode, masterKey.PublicKey, crypto.PublicKeyStart)),
		}, nil
	case nockchain.ImportType_WATCH_ONLY:
		pubKey := base58.Decode(req.Key)
		return &nockchain.ImportKeysResponse{
			Seed:             "",
			PrivateKey:       "",
			PublicKey:        base58.Encode(pubKey),
			ChainCode:        "",
			ImportPrivateKey: "",
			ImportPublicKey:  "",
		}, nil
	default:
		return nil, fmt.Errorf("invalid import type")
	}
}

func (h *GprcHandler) DeriveChild(ctx context.Context, req *nockchain.DeriveChildRequest) (*nockchain.DeriveChildResponse, error) {
	data := base58.Decode(req.ImportedKey)
	index := req.Index
	if index > 1<<32 {
		return nil, fmt.Errorf("child index %d out of range, child indices are capped to values between [0, 2^32)", index)
	}
	if req.Hardened {
		index += 1 << 31
	}
	switch {
	case strings.HasPrefix(req.ImportedKey, "zprv"):
		if len(data) != 82 {
			return nil, fmt.Errorf("invalid extended private key length: %d (expected 82)", len(data))
		}
		if data[45] != 0x00 {
			return nil, fmt.Errorf("invalid private key prefix at byte 45: 0x%02x (expected 0x00)", data[45])
		}
		hash := sha256.Sum256(data[:78])
		hash = sha256.Sum256(hash[:])
		if !slices.Equal(hash[:4], data[78:]) {
			return nil, fmt.Errorf("invalid checksum")
		}
		chainCode := make([]byte, 32)
		copy(chainCode, data[13:45])
		privateKey := make([]byte, 32)
		copy(privateKey, data[46:78])
		masterKey, err := crypto.MasterKeyFromPrivKey(chainCode, privateKey)
		if err != nil {
			return nil, err
		}

		childKey, err := masterKey.DeriveChild(index)
		if err != nil {
			return nil, err
		}
		return &nockchain.DeriveChildResponse{
			PublicKey:  base58.Encode(childKey.PublicKey),
			PrivateKey: base58.Encode(childKey.PrivateKey),
			ChainCode:  base58.Encode(childKey.ChainCode),
		}, nil

	case strings.HasPrefix(req.ImportedKey, "zpub"):
		if len(data) != 145 {
			return nil, fmt.Errorf("invalid extended public key length: %d (expected 145)", len(data))
		}

		hash := sha256.Sum256(data[:141])
		hash = sha256.Sum256(hash[:])
		if !slices.Equal(hash[:4], data[141:]) {
			return nil, fmt.Errorf("invalid checksum")
		}

		chainCode := make([]byte, 32)
		copy(chainCode, data[13:45])
		publicKey := make([]byte, 97)
		copy(publicKey, data[45:141])

		masterKey := crypto.MasterKey{
			PublicKey:  publicKey,
			ChainCode:  chainCode,
			PrivateKey: []byte{},
		}
		childKey, err := masterKey.DeriveChild(index)
		if err != nil {
			return nil, err
		}
		return &nockchain.DeriveChildResponse{
			PublicKey:  base58.Encode(childKey.PublicKey),
			PrivateKey: "",
			ChainCode:  base58.Encode(childKey.ChainCode),
		}, nil
	default:
		return nil, fmt.Errorf("invalid extended key")
	}
}

//   - `names` - Comma-separated list of note name pairs in format "[first last]"
//     Example: "[first1 last1],[first2 last2]"
//
//   - `recipients` - Comma-separated list of recipient $locks
//     Example: "[1 pk1],[2 pk2,pk3,pk4]"
//     A simple comma-separated list is also supported: "pk1,pk2,pk3",
//     where it is presumed that all recipients are single-signature,
//     that is to say, it is the same as "[1 pk1],[1 pk2],[1 pk3]"
//
//   - `gifts` - Comma-separated list of amounts to send to each recipient
//     Example: "100,200"
//
// - `fee` - Transaction fee to be subtracted from one of the input notes
func (h *GprcHandler) CreateTx(ctx context.Context, req *nockchain.CreateTxRequest) (*nockchain.CreateTxResponse, error) {
	firstNames := [][5]uint64{}
	lastNames := [][5]uint64{}
	names := strings.Split(req.Names, ",")
	notes := make([]*nockchain.NockchainNote, len(names))
	for _, name := range names {
		name = strings.TrimSpace(name)
		if strings.HasPrefix(name, "[") && strings.HasSuffix(name, "]") {
			inner := name[1 : len(name)-1]
			part := strings.Split(inner, " ")
			if len(part) == 2 {
				firstNames = append(firstNames, crypto.Base58ToTip5Hash(part[0]))
				lastNames = append(lastNames, crypto.Base58ToTip5Hash(part[1]))
			}
		}
	}

	recipents := []*nockchain.NockchainLock{}
	if strings.Contains(req.Recipients, "[") {
		pairs := strings.Split(req.Recipients, ",")
		for _, pair := range pairs {
			pair = strings.TrimSpace(pair)

			if strings.HasPrefix(pair, "[") && strings.HasSuffix(pair, "]") {
				inner := pair[1 : len(pair)-1]
				parts := strings.SplitN(inner, " ", 2)

				if len(parts) == 2 {
					number, err := strconv.ParseUint(parts[0], 10, 64)
					if err != nil {
						continue
					}

					pubkeysStr := strings.Split(parts[1], ",")

					recipents = append(recipents, &nockchain.NockchainLock{KeysRequired: number, Pubkeys: pubkeysStr})
				}
			}
		}
	} else {
		// Parse simple format: "pk1,pk2,pk3"
		addrs := strings.Split(req.Recipients, ",")

		for _, addr := range addrs {
			recipents = append(recipents, &nockchain.NockchainLock{KeysRequired: uint64(1), Pubkeys: []string{strings.TrimSpace(addr)}})
		}
	}
	gifts := []uint64{}
	for _, gift := range strings.Split(req.Gifts, ",") {
		gift, err := strconv.ParseUint(gift, 10, 64)
		if err != nil {
			continue
		}

		gifts = append(gifts, gift)
	}

	// Verify lengths based on single vs multiple mode
	if len(recipents) == 1 && len(gifts) == 1 {
		// Single mode: can spend from multiple notes to single recipient
		// No additional validation needed - any number of names is allowed
	} else {
		// Multiple mode: all lengths must match
		if len(firstNames) != len(recipents) || len(firstNames) != len(gifts) {
			return nil, fmt.Errorf("multiple recipient mode requires names, recipients, and gifts to have the same length")
		}
	}

	var masterKey *crypto.MasterKey
	chainCode := base58.Decode(req.ChainCode)
	key := base58.Decode(req.Key)
	masterKey, err := crypto.MasterKeyFromPrivKey(chainCode, key)
	if err != nil {
		return nil, err
	}

	if !req.IsMasterKey {
		index := req.Index
		if index > 1<<32 {
			return nil, fmt.Errorf("child index %d out of range, child indices are capped to values between [0, 2^32)", index)
		}
		if req.Hardened {
			index += 1 << 31
		}
		childKey, err := masterKey.DeriveChild(index)
		if err != nil {
			return nil, err
		}
		masterKey = &childKey
	}

	// Scan key to get notes
	masterKeyScan, err := h.client.WalletGetBalance(base58.Encode(masterKey.PublicKey))
	if err != nil {
		return nil, err
	}
	if len(masterKeyScan.Notes) != 0 {
		for _, note := range masterKeyScan.Notes {
			firstName := crypto.Tip5HashToBase58([5]uint64{
				note.Name.First.Belt_1.Value,
				note.Name.First.Belt_2.Value,
				note.Name.First.Belt_3.Value,
				note.Name.First.Belt_4.Value,
				note.Name.First.Belt_5.Value,
			})
			lastName := crypto.Tip5HashToBase58([5]uint64{
				note.Name.Last.Belt_1.Value,
				note.Name.Last.Belt_2.Value,
				note.Name.Last.Belt_3.Value,
				note.Name.Last.Belt_4.Value,
				note.Name.Last.Belt_5.Value,
			})
			name := "[" + firstName + " " + lastName + "]"
			if i := slices.Index(names, name); i != -1 {
				balanceEntry := ParseBalanceEntry(note)
				notes[i] = &balanceEntry
			}
		}
	}

	inputs := []*nockchain.NockchainInput{}
	for i := 0; i < len(names); i++ {
		if notes[i] == nil {
			return nil, fmt.Errorf("notes scanned is missing at name: %s", names[i])
		}

		if notes[i].Asset < gifts[i]+req.Fee {
			return nil, fmt.Errorf("note not enough balance")
		}

		parentHash, err := HashNote(notes[i])
		if err != nil {
			return nil, err
		}
		seeds := []*nockchain.NockchainSeed{
			{
				OutputSource:   nil,
				Recipient:      recipents[i],
				TimelockIntent: req.TimelockIntent,
				Gift:           gifts[i],
				ParentHash:     crypto.Tip5HashToBase58(parentHash),
			},
		}

		if notes[i].Asset < gifts[i]+req.Fee {
			seeds = append(seeds, &nockchain.NockchainSeed{
				OutputSource: nil,
				Recipient: &nockchain.NockchainLock{
					KeysRequired: 1,
					Pubkeys:      []string{base58.Encode(masterKey.PublicKey)},
				},
				TimelockIntent: req.TimelockIntent,
				Gift:           notes[i].Asset - gifts[i] - req.Fee,
				ParentHash:     crypto.Tip5HashToBase58(parentHash),
			})
		}

		spend := nockchain.NockchainSpend{
			Signatures: nil,
			Seeds:      seeds,
			Fee:        req.Fee,
		}

		msg, err := HashMsg(&spend)
		if err != nil {
			return nil, err
		}

		// sign
		chalT8, sigT8, err := ComputeSig(*masterKey, msg)
		if err != nil {
			return nil, err
		}

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

		input := nockchain.NockchainInput{
			Name: &nockchain.NockchainName{
				First: crypto.Tip5HashToBase58(firstNames[i]),
				Last:  crypto.Tip5HashToBase58(lastNames[i]),
			},
			Note:  notes[i],
			Spend: &spend,
		}
		inputs = append(inputs, &input)
	}

	var timelockRange *nockchain.TimelockRange
	if req.TimelockIntent == nil {
		timelockRange = nil
	} else {
		if req.TimelockIntent.Absolute != nil {
			timelockRange = req.TimelockIntent.Absolute
		}
		if req.TimelockIntent.Relative != nil {
			timelockRange = req.TimelockIntent.Relative

		}
	}
	rawTx := nockchain.RawTx{
		TxId:          "",
		Inputs:        inputs,
		TimelockRange: timelockRange,
		TotalFees:     req.Fee,
	}
	txId, err := ComputeTxId(inputs, timelockRange, req.Fee)
	if err != nil {
		return nil, err
	}
	rawTx.TxId = crypto.Tip5HashToBase58(txId)
	return &nockchain.CreateTxResponse{
		RawTx: &rawTx,
	}, nil
}

func (h *GprcHandler) Scan(ctx context.Context, req *nockchain.ScanRequest) (*nockchain.ScanResponse, error) {
	scanData := []*nockchain.ScanData{}
	keyBytes := base58.Decode(req.MasterPubkey)

	chainCode := base58.Decode(req.ChainCode)
	masterKey := crypto.MasterKey{
		PublicKey:  keyBytes,
		ChainCode:  chainCode,
		PrivateKey: []byte{},
	}

	masterKeyScan, err := h.client.WalletGetBalance(req.MasterPubkey)
	if err != nil {
		return nil, err
	}
	if len(masterKeyScan.Notes) != 0 {
		scanData = append(scanData, &nockchain.ScanData{
			Pubkey: req.MasterPubkey,
			Data:   masterKeyScan,
		})
	}
	for i := uint64(0); i < req.SearchDepth; i++ {
		childKey, err := masterKey.DeriveChild(i)
		if err != nil {
			continue
		}
		childKeyScan, err := h.client.WalletGetBalance(base58.Encode(childKey.PublicKey))
		if err != nil {
			continue
		}
		if len(childKeyScan.Notes) != 0 {
			scanData = append(scanData, &nockchain.ScanData{
				Pubkey: base58.Encode(childKey.PublicKey),
				Data:   childKeyScan,
			})
		}
	}

	return &nockchain.ScanResponse{
		ScanData: scanData,
	}, nil
}

func (h *GprcHandler) SignTx(context.Context, *nockchain.SignTxRequest) (*nockchain.SignTxResponse, error) {
	return nil, nil
}