package bip85

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"strings"

	"github.com/btcsuite/btcd/btcec/v2"
	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/btcutil/base58"
	"github.com/btcsuite/btcd/btcutil/hdkeychain"
	"github.com/btcsuite/btcd/chaincfg"
	"golang.org/x/crypto/sha3"
)

const (
	// BIP85_MASTER_PATH is the derivation path prefix for all BIP85 applications
	BIP85_MASTER_PATH = "m/83696968'"

	// BIP85_KEY_HMAC_KEY is the HMAC key used for deriving the entropy
	BIP85_KEY_HMAC_KEY = "bip-entropy-from-k"

	// Application numbers
	APP_BIP39  = 39 // BIP39 mnemonics
	APP_HD_WIF = 2  // WIF for Bitcoin Core
	APP_XPRV   = 32 // Extended private key
	APP_HEX    = 128169
	APP_PWD64  = 707764 // Base64 passwords
	APP_PWD85  = 707785 // Base85 passwords
	APP_RSA    = 828365
)

// Version bytes for extended keys
var (
	// MainNetPrivateKey is the version for mainnet private keys
	MainNetPrivateKey = []byte{0x04, 0x88, 0xAD, 0xE4}
	// TestNetPrivateKey is the version for testnet private keys
	TestNetPrivateKey = []byte{0x04, 0x35, 0x83, 0x94}
)

// BIP85DRNG is a deterministic random number generator seeded by BIP85 entropy
type BIP85DRNG struct {
	shake io.Reader
}

// NewBIP85DRNG creates a new DRNG seeded with BIP85 entropy
func NewBIP85DRNG(entropy []byte) *BIP85DRNG {
	// The entropy must be exactly 64 bytes (512 bits)
	if len(entropy) != 64 {
		panic("BIP85DRNG entropy must be 64 bytes")
	}

	// Initialize SHAKE256 with the entropy
	shake := sha3.NewShake256()
	shake.Write(entropy)

	return &BIP85DRNG{
		shake: shake,
	}
}

// Read implements the io.Reader interface
func (d *BIP85DRNG) Read(p []byte) (n int, err error) {
	return d.shake.Read(p)
}

// DeriveChildKey returns the private key and chain code bytes
func DeriveChildKey(masterKey *hdkeychain.ExtendedKey, path string) ([]byte, error) {
	// Validate the masterKey is a private key
	if !masterKey.IsPrivate() {
		return nil, fmt.Errorf("master key must be a private key")
	}

	// Derive the child key at the specified path
	childKey, err := deriveChildKey(masterKey, path)
	if err != nil {
		return nil, fmt.Errorf("failed to derive child key: %w", err)
	}

	// Get the private key bytes
	ecPrivKey, err := childKey.ECPrivKey()
	if err != nil {
		return nil, fmt.Errorf("failed to get EC private key: %w", err)
	}

	// Serialize the private key to get the bytes
	return ecPrivKey.Serialize(), nil
}

// DeriveBIP85Entropy derives entropy from a BIP32 master key using the BIP85 method
func DeriveBIP85Entropy(masterKey *hdkeychain.ExtendedKey, path string) ([]byte, error) {
	// Get the child key bytes
	privKeyBytes, err := DeriveChildKey(masterKey, path)
	if err != nil {
		return nil, err
	}

	// Apply HMAC-SHA512
	h := hmac.New(sha512.New, []byte(BIP85_KEY_HMAC_KEY))
	h.Write(privKeyBytes)
	entropy := h.Sum(nil)

	return entropy, nil
}

// deriveChildKey derives a child key from a parent key using the given path
func deriveChildKey(parent *hdkeychain.ExtendedKey, path string) (*hdkeychain.ExtendedKey, error) {
	if path == "" || path == "m" || path == "/" {
		return parent, nil
	}

	// Remove the "m/" or "/" prefix if present
	path = strings.TrimPrefix(path, "m/")
	path = strings.TrimPrefix(path, "/")

	// Split the path into individual components
	components := strings.Split(path, "/")

	// Start with the parent key
	key := parent

	// Derive each component
	for _, component := range components {
		// Check if the component is hardened
		hardened := strings.HasSuffix(component, "'") || strings.HasSuffix(component, "h")
		if hardened {
			component = strings.TrimSuffix(component, "'")
			component = strings.TrimSuffix(component, "h")
		}

		// Parse the index
		var index uint32
		_, err := fmt.Sscanf(component, "%d", &index)
		if err != nil {
			return nil, fmt.Errorf("invalid path component: %s", component)
		}

		// Apply hardening if needed
		if hardened {
			index += hdkeychain.HardenedKeyStart
		}

		// Derive the child key
		child, err := key.Derive(index)
		if err != nil {
			return nil, fmt.Errorf("failed to derive child key at index %d: %w", index, err)
		}

		key = child
	}

	return key, nil
}

// DeriveBIP39Entropy derives entropy for a BIP39 mnemonic
func DeriveBIP39Entropy(masterKey *hdkeychain.ExtendedKey, language, words, index uint32) ([]byte, error) {
	path := fmt.Sprintf("%s/%d'/%d'/%d'/%d'", BIP85_MASTER_PATH, APP_BIP39, language, words, index)

	entropy, err := DeriveBIP85Entropy(masterKey, path)
	if err != nil {
		return nil, err
	}

	// Determine how many bits of entropy to use based on the words
	var bits int
	switch words {
	case 12:
		bits = 128
	case 15:
		bits = 160
	case 18:
		bits = 192
	case 21:
		bits = 224
	case 24:
		bits = 256
	default:
		return nil, fmt.Errorf("invalid BIP39 word count: %d", words)
	}

	// Truncate to the required number of bits (bytes = bits / 8)
	entropy = entropy[:bits/8]

	return entropy, nil
}

// DeriveWIFKey derives a private key in WIF format
func DeriveWIFKey(masterKey *hdkeychain.ExtendedKey, index uint32) (string, error) {
	path := fmt.Sprintf("%s/%d'/%d'", BIP85_MASTER_PATH, APP_HD_WIF, index)

	entropy, err := DeriveBIP85Entropy(masterKey, path)
	if err != nil {
		return "", err
	}

	// Use the first 32 bytes as the key
	keyBytes := entropy[:32]

	// Convert to WIF format
	privKey, _ := btcec.PrivKeyFromBytes(keyBytes)
	wif, err := btcutil.NewWIF(privKey, &chaincfg.MainNetParams, true) // compressed=true
	if err != nil {
		return "", fmt.Errorf("failed to create WIF: %w", err)
	}

	return wif.String(), nil
}

// DeriveXPRV derives an extended private key (XPRV)
func DeriveXPRV(masterKey *hdkeychain.ExtendedKey, index uint32) (*hdkeychain.ExtendedKey, error) {
	path := fmt.Sprintf("%s/%d'/%d'", BIP85_MASTER_PATH, APP_XPRV, index)

	entropy, err := DeriveBIP85Entropy(masterKey, path)
	if err != nil {
		return nil, err
	}

	// The first 32 bytes are the chain code, the second 32 bytes are the private key
	chainCode := entropy[:32]
	privateKey := entropy[32:64]

	// Create serialized extended key
	var serialized bytes.Buffer

	// Add version bytes (4 bytes)
	// Default to mainnet version
	version := MainNetPrivateKey

	// Check if the master key serialization starts with the testnet version bytes
	masterKeyStr := masterKey.String()
	if strings.HasPrefix(masterKeyStr, "tprv") {
		version = TestNetPrivateKey
	}

	// Write serialized data
	serialized.Write(version)            // 4 bytes: version
	serialized.WriteByte(0)              // 1 byte: depth (0 for master)
	serialized.Write([]byte{0, 0, 0, 0}) // 4 bytes: parent fingerprint (0 for master)
	serialized.Write([]byte{0, 0, 0, 0}) // 4 bytes: child number (0 for master)
	serialized.Write(chainCode)          // 32 bytes: chain code
	serialized.WriteByte(0)              // 1 byte: 0x00 prefix for private key
	serialized.Write(privateKey)         // 32 bytes: private key

	// Calculate checksum (first 4 bytes of double-SHA256)
	serializedBytes := serialized.Bytes()
	checksum := doubleSHA256(serializedBytes)[:4]

	// Append checksum
	serializedWithChecksum := append(serializedBytes, checksum...)

	// Base58 encode
	xprvStr := base58.Encode(serializedWithChecksum)

	// Parse the serialized xprv back to an ExtendedKey
	return hdkeychain.NewKeyFromString(xprvStr)
}

// doubleSHA256 calculates sha256(sha256(data))
func doubleSHA256(data []byte) []byte {
	hash1 := sha256.Sum256(data)
	hash2 := sha256.Sum256(hash1[:])
	return hash2[:]
}

// DeriveHex derives a raw hex string of specified length
func DeriveHex(masterKey *hdkeychain.ExtendedKey, numBytes, index uint32) (string, error) {
	if numBytes < 16 || numBytes > 64 {
		return "", fmt.Errorf("numBytes must be between 16 and 64")
	}

	path := fmt.Sprintf("%s/%d'/%d'/%d'", BIP85_MASTER_PATH, APP_HEX, numBytes, index)

	entropy, err := DeriveBIP85Entropy(masterKey, path)
	if err != nil {
		return "", err
	}

	// Truncate to the required number of bytes
	entropy = entropy[:numBytes]

	return hex.EncodeToString(entropy), nil
}

// DeriveBase64Password derives a password encoded in Base64
func DeriveBase64Password(masterKey *hdkeychain.ExtendedKey, pwdLen, index uint32) (string, error) {
	if pwdLen < 20 || pwdLen > 86 {
		return "", fmt.Errorf("pwdLen must be between 20 and 86")
	}

	path := fmt.Sprintf("%s/%d'/%d'/%d'", BIP85_MASTER_PATH, APP_PWD64, pwdLen, index)

	entropy, err := DeriveBIP85Entropy(masterKey, path)
	if err != nil {
		return "", err
	}

	// Base64 encode all 64 bytes of entropy
	encodedStr := base64.StdEncoding.EncodeToString(entropy)

	// Remove any padding
	encodedStr = strings.TrimRight(encodedStr, "=")

	// Slice to the desired password length
	if uint32(len(encodedStr)) < pwdLen {
		return "", fmt.Errorf("derived password length %d is shorter than requested length %d", len(encodedStr), pwdLen)
	}

	return encodedStr[:pwdLen], nil
}

// DeriveBase85Password derives a password encoded in Base85
func DeriveBase85Password(masterKey *hdkeychain.ExtendedKey, pwdLen, index uint32) (string, error) {
	if pwdLen < 10 || pwdLen > 80 {
		return "", fmt.Errorf("pwdLen must be between 10 and 80")
	}

	path := fmt.Sprintf("%s/%d'/%d'/%d'", BIP85_MASTER_PATH, APP_PWD85, pwdLen, index)

	entropy, err := DeriveBIP85Entropy(masterKey, path)
	if err != nil {
		return "", err
	}

	// Base85 encode all 64 bytes of entropy using the RFC1924 character set
	encoded := encodeBase85WithRFC1924Charset(entropy)

	// Slice to the desired password length
	if uint32(len(encoded)) < pwdLen {
		return "", fmt.Errorf("encoded length %d is less than requested length %d", len(encoded), pwdLen)
	}

	return encoded[:pwdLen], nil
}

// encodeBase85WithRFC1924Charset encodes data using Base85 with the RFC1924 character set
func encodeBase85WithRFC1924Charset(data []byte) string {
	// RFC1924 character set
	charset := "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!#$%&()*+-;<=>?@^_`{|}~"

	// Pad data to multiple of 4
	padded := make([]byte, ((len(data)+3)/4)*4)
	copy(padded, data)

	var buf strings.Builder
	buf.Grow(len(padded) * 5 / 4) // Each 4 bytes becomes 5 Base85 characters

	// Process in 4-byte chunks
	for i := 0; i < len(padded); i += 4 {
		// Convert 4 bytes to uint32 (big-endian)
		chunk := binary.BigEndian.Uint32(padded[i : i+4])

		// Convert to 5 base-85 digits
		digits := make([]byte, 5)
		for j := 4; j >= 0; j-- {
			idx := chunk % 85
			digits[j] = charset[idx]
			chunk /= 85
		}

		buf.Write(digits)
	}

	return buf.String()
}

// ParseMasterKey parses an extended key from a string
func ParseMasterKey(xprv string) (*hdkeychain.ExtendedKey, error) {
	return hdkeychain.NewKeyFromString(xprv)
}