package bsv

import (
	"encoding/binary"
	"strings"
)

// BIP32 constants.
const (
	HDHardened     = 0x80000000
	HDMaxIndex     = 2 * HDHardened
	HDFingerprintSize = 4
)

// hdBuffers holds the raw serialized fields of an extended key.
type hdBuffers struct {
	Version          []byte
	Depth            []byte
	ParentFingerPrint []byte
	ChildIndex       []byte
	ChainCode        []byte
	PrivateKey       []byte // 32 bytes or nil for public
	PublicKey        []byte // 33 bytes or nil for private
	Checksum         []byte
}

// HDPrivateKey mirrors lib/hdprivatekey (BIP32 extended private key).
type HDPrivateKey struct {
	Network    *Network
	Depth      int
	PrivateKey *PrivateKey
	PublicKey  *PublicKey
	FingerPrint []byte

	Xprivkey string
	buffers  *hdBuffers
	hdPub    *HDPublicKey
}

// Xpubkey returns the serialized extended public key (lazily computed).
func (k *HDPrivateKey) Xpubkey() string {
	if k.hdPub == nil {
		pub, err := hdPublicKeyFromPrivate(k)
		if err == nil {
			k.hdPub = pub
		}
	}
	if k.hdPub != nil {
		return k.hdPub.Xpubkey
	}
	return ""
}

// NewHDPrivateKey builds an HDPrivateKey from a serialized string, a buffer,
// a seed (hex string or buffer), or an object.
func NewHDPrivateKey(arg interface{}) (*HDPrivateKey, error) {
	switch v := arg.(type) {
	case *HDPrivateKey:
		return v, nil
	case nil:
		return HDPrivateKeyFromSeed(nil, nil)
	case string:
		return hdPrivateKeyFromSerialized(v)
	case []byte:
		if isHexa(string(v)) {
			return HDPrivateKeyFromSeed(v, nil)
		}
		return hdPrivateKeyFromSerialized(string(v))
	default:
		return nil, newError("UnrecognizedArgument", "Unrecognized argument for HDPrivateKey")
	}
}

// hdPrivateKeyFromSerialized parses an xprv string.
func hdPrivateKeyFromSerialized(xprv string) (*HDPrivateKey, error) {
	decoded, err := Base58CheckDecode(xprv)
	if err != nil {
		return nil, err
	}
	if len(decoded) != 78 {
		return nil, newError("InvalidLength", "invalid xprv length: %d", len(decoded))
	}
	buffers := &hdBuffers{
		Version:          decoded[0:4],
		Depth:            decoded[4:5],
		ParentFingerPrint: decoded[5:9],
		ChildIndex:       decoded[9:13],
		ChainCode:        decoded[13:45],
		PrivateKey:       decoded[46:78],
		Checksum:         nil,
	}
	return hdPrivateKeyFromBuffers(buffers, xprv)
}

// hdPrivateKeyFromBuffers validates and builds the key structure.
func hdPrivateKeyFromBuffers(buffers *hdBuffers, xprv string) (*HDPrivateKey, error) {
	network := NetworkGet(binary.BigEndian.Uint32(buffers.Version), []string{"xprivkey"})
	if network == nil {
		return nil, newError("InvalidNetwork", "invalid xprivkey version")
	}
	priv, err := NewPrivateKey(NewBNFromBuffer(buffers.PrivateKey, false), network)
	if err != nil {
		return nil, err
	}
	pub := priv.ToPublicKey()

	seq := concatBytes(buffers.Version, buffers.Depth, buffers.ParentFingerPrint,
		buffers.ChildIndex, buffers.ChainCode, []byte{0}, buffers.PrivateKey)
	cksum := Base58CheckChecksum(seq)
	if len(buffers.Checksum) > 0 && !byteEq(buffers.Checksum, cksum) {
		return nil, newError("InvalidB58Checksum", "invalid checksum")
	}
	buffers.Checksum = cksum

	if xprv == "" {
		xprv = Base58CheckEncode(seq)
	}

	fingerPrint := (Hash{}).Sha256ripemd160(pub.ToDER())[:HDFingerprintSize]

	key := &HDPrivateKey{
		Network:     network,
		Depth:       int(buffers.Depth[0]),
		PrivateKey:  priv,
		PublicKey:   pub,
		FingerPrint: fingerPrint,
		Xprivkey:    xprv,
		buffers:     buffers,
	}
	return key, nil
}

func byteEq(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// HDPrivateKeyFromSeed generates the master key from a seed (BIP32).
func HDPrivateKeyFromSeed(seed []byte, network interface{}) (*HDPrivateKey, error) {
	if seed == nil {
		var err error
		seed, err = (Random{}).GetRandomBuffer(64)
		if err != nil {
			return nil, err
		}
	}
	if len(seed) < 16 {
		return nil, newError("NotEnoughEntropy", "not enough entropy: %d bytes", len(seed))
	}
	if len(seed) > 64 {
		return nil, newError("TooMuchEntropy", "too much entropy: %d bytes", len(seed))
	}
	hash := (Hash{}).Sha512hmac(seed, []byte("Bitcoin seed"))
	var net *Network
	if network != nil {
		net = NetworkGet(network, nil)
	}
	if net == nil {
		net = defaultNetwork
	}
	version := make([]byte, 4)
	binary.BigEndian.PutUint32(version, net.Xprivkey)
	buffers := &hdBuffers{
		Version:          version,
		Depth:            []byte{0},
		ParentFingerPrint: []byte{0, 0, 0, 0},
		ChildIndex:       []byte{0, 0, 0, 0},
		ChainCode:        hash[32:64],
		PrivateKey:       hash[0:32],
	}
	return hdPrivateKeyFromBuffers(buffers, "")
}

// IsValidHDPath validates a derivation path string or index.
func IsValidHDPath(arg interface{}, hardened bool) bool {
	switch v := arg.(type) {
	case string:
		indexes := getDerivationIndexes(v)
		if indexes == nil {
			return false
		}
		for _, idx := range indexes {
			if !IsValidHDPath(idx, false) {
				return false
			}
		}
		return true
	case int:
		idx := v
		if idx < HDHardened && hardened {
			idx += HDHardened
		}
		return idx >= 0 && idx < HDMaxIndex
	}
	return false
}

// getDerivationIndexes parses a path string like "m/0/1/2'" into indexes.
// Returns nil when malformed.
func getDerivationIndexes(path string) []int {
	steps := strings.Split(path, "/")
	rootAliases := map[string]bool{"m": true, "M": true, "m'": true, "M'": true}
	if rootAliases[path] {
		return []int{}
	}
	if !rootAliases[steps[0]] {
		return nil
	}
	var indexes []int
	for _, step := range steps[1:] {
		isHardened := strings.HasSuffix(step, "'")
		if isHardened {
			step = step[:len(step)-1]
		}
		if step == "" || strings.HasPrefix(step, "-") {
			return nil
		}
		index := 0
		valid := true
		for _, c := range step {
			if c < '0' || c > '9' {
				valid = false
				break
			}
			index = index*10 + int(c-'0')
		}
		if !valid {
			return nil
		}
		if isHardened {
			index += HDHardened
		}
		indexes = append(indexes, index)
	}
	return indexes
}

// DeriveChild derives a child key from an index or path string.
func (k *HDPrivateKey) DeriveChild(arg interface{}, hardened bool) (*HDPrivateKey, error) {
	switch v := arg.(type) {
	case int:
		return k.deriveWithNumber(v, hardened, false)
	case string:
		return k.deriveFromString(v, false)
	default:
		return nil, newError("InvalidDerivationArgument", "invalid derivation argument")
	}
}

// DeriveNonCompliantChild derives using the old non-padded serialization.
func (k *HDPrivateKey) DeriveNonCompliantChild(arg interface{}, hardened bool) (*HDPrivateKey, error) {
	switch v := arg.(type) {
	case int:
		return k.deriveWithNumber(v, hardened, true)
	case string:
		return k.deriveFromString(v, true)
	default:
		return nil, newError("InvalidDerivationArgument", "invalid derivation argument")
	}
}

func (k *HDPrivateKey) deriveWithNumber(index int, hardened bool, nonCompliant bool) (*HDPrivateKey, error) {
	if !IsValidHDPath(index, hardened) {
		return nil, newError("InvalidPath", "invalid path: %d", index)
	}
	if index >= HDHardened {
		hardened = true
	}
	if index < HDHardened && hardened {
		index += HDHardened
	}

	indexBuffer := make([]byte, 4)
	binary.BigEndian.PutUint32(indexBuffer, uint32(index))
	var data []byte
	if hardened && nonCompliant {
		nonZeroPadded := k.PrivateKey.Bn.ToBuffer(0, false)
		data = concatBytes([]byte{0}, nonZeroPadded, indexBuffer)
	} else if hardened {
		privateKeyBuffer := k.PrivateKey.Bn.ToBuffer(32, false)
		data = concatBytes([]byte{0}, privateKeyBuffer, indexBuffer)
	} else {
		data = concatBytes(k.PublicKey.ToDER(), indexBuffer)
	}
	hash := (Hash{}).Sha512hmac(data, k.buffers.ChainCode)
	leftPart := NewBNFromBuffer(hash[0:32], false)
	chainCode := hash[32:64]

	privateKey := leftPart.Add(k.PrivateKey.Bn.Int).Umod(GetN().Int).ToBuffer(32, false)

	if !PrivateKeyIsValid(privateKey, nil) {
		return k.deriveWithNumber(index+1, false, nonCompliant)
	}

	version := make([]byte, 4)
	binary.BigEndian.PutUint32(version, k.Network.Xprivkey)
	parentFp := make([]byte, 4)
	copy(parentFp, k.FingerPrint)
	childIdx := make([]byte, 4)
	binary.BigEndian.PutUint32(childIdx, uint32(index))

	buffers := &hdBuffers{
		Version:          version,
		Depth:            []byte{byte(k.Depth + 1)},
		ParentFingerPrint: parentFp,
		ChildIndex:       childIdx,
		ChainCode:        chainCode,
		PrivateKey:       privateKey,
	}
	return hdPrivateKeyFromBuffers(buffers, "")
}

func (k *HDPrivateKey) deriveFromString(path string, nonCompliant bool) (*HDPrivateKey, error) {
	if !IsValidHDPath(path, false) {
		return nil, newError("InvalidPath", "invalid path: %s", path)
	}
	indexes := getDerivationIndexes(path)
	derived := k
	var err error
	for _, index := range indexes {
		derived, err = derived.deriveWithNumber(index, false, nonCompliant)
		if err != nil {
			return nil, err
		}
	}
	return derived, nil
}

// ToString returns the xprv serialization.
func (k *HDPrivateKey) ToString() string { return k.Xprivkey }

// ToObject returns a plain object representation.
func (k *HDPrivateKey) ToObject() map[string]interface{} {
	return map[string]interface{}{
		"network":          k.Network.Name,
		"depth":            k.Depth,
		"fingerPrint":      binary.BigEndian.Uint32(k.FingerPrint),
		"parentFingerPrint": binary.BigEndian.Uint32(k.buffers.ParentFingerPrint),
		"childIndex":       binary.BigEndian.Uint32(k.buffers.ChildIndex),
		"chainCode":        toHex(k.buffers.ChainCode),
		"privateKey":       k.PrivateKey.ToHex(),
		"checksum":         binary.BigEndian.Uint32(k.buffers.Checksum),
		"xprivkey":         k.Xprivkey,
	}
}

// ToBuffer returns the xprv string as bytes.
func (k *HDPrivateKey) ToBuffer() []byte { return []byte(k.Xprivkey) }

// ToHex returns the hex of the xprv string.
func (k *HDPrivateKey) ToHex() string { return toHex([]byte(k.Xprivkey)) }

// HDPublicKey mirrors lib/hdpublickey (BIP32 extended public key).
type HDPublicKey struct {
	Network    *Network
	Depth      int
	PublicKey  *PublicKey
	FingerPrint []byte

	Xpubkey string
	buffers *hdBuffers
}

// NewHDPublicKey builds an HDPublicKey from an xpub string, buffer, object,
// or HDPrivateKey.
func NewHDPublicKey(arg interface{}) (*HDPublicKey, error) {
	switch v := arg.(type) {
	case *HDPublicKey:
		return v, nil
	case *HDPrivateKey:
		return hdPublicKeyFromPrivate(v)
	case string:
		return hdPublicKeyFromSerialized(v)
	case []byte:
		return hdPublicKeyFromSerialized(string(v))
	default:
		return nil, newError("UnrecognizedArgument", "Unrecognized argument for HDPublicKey")
	}
}

// HDPublicKeyFromHDPrivateKey derives the public key from a private key.
func HDPublicKeyFromHDPrivateKey(hd *HDPrivateKey) (*HDPublicKey, error) {
	return hdPublicKeyFromPrivate(hd)
}

func hdPublicKeyFromPrivate(arg *HDPrivateKey) (*HDPublicKey, error) {
	buffers := &hdBuffers{
		Version:          make([]byte, 4),
		Depth:            arg.buffers.Depth,
		ParentFingerPrint: arg.buffers.ParentFingerPrint,
		ChildIndex:       arg.buffers.ChildIndex,
		ChainCode:        arg.buffers.ChainCode,
		PublicKey:        arg.PublicKey.ToDER(),
	}
	binary.BigEndian.PutUint32(buffers.Version, arg.Network.Xpubkey)
	return hdPublicKeyFromBuffers(buffers, "")
}

func hdPublicKeyFromSerialized(xpub string) (*HDPublicKey, error) {
	decoded, err := Base58CheckDecode(xpub)
	if err != nil {
		return nil, err
	}
	if len(decoded) != 78 {
		return nil, newError("InvalidLength", "invalid xpub length: %d", len(decoded))
	}
	buffers := &hdBuffers{
		Version:          decoded[0:4],
		Depth:            decoded[4:5],
		ParentFingerPrint: decoded[5:9],
		ChildIndex:       decoded[9:13],
		ChainCode:        decoded[13:45],
		PublicKey:        decoded[45:78],
	}
	return hdPublicKeyFromBuffers(buffers, xpub)
}

func hdPublicKeyFromBuffers(buffers *hdBuffers, xpub string) (*HDPublicKey, error) {
	network := NetworkGet(binary.BigEndian.Uint32(buffers.Version), []string{"xpubkey"})
	if network == nil {
		return nil, newError("InvalidNetwork", "invalid xpubkey version")
	}
	pub, err := PublicKeyFromDER(buffers.PublicKey, true)
	if err != nil {
		return nil, err
	}
	seq := concatBytes(buffers.Version, buffers.Depth, buffers.ParentFingerPrint,
		buffers.ChildIndex, buffers.ChainCode, buffers.PublicKey)
	cksum := Base58CheckChecksum(seq)
	if len(buffers.Checksum) > 0 && !byteEq(buffers.Checksum, cksum) {
		return nil, newError("InvalidB58Checksum", "invalid checksum")
	}
	buffers.Checksum = cksum
	if xpub == "" {
		xpub = Base58CheckEncode(seq)
	}
	fingerPrint := (Hash{}).Sha256ripemd160(pub.ToDER())[:HDFingerprintSize]

	return &HDPublicKey{
		Network:     network,
		Depth:       int(buffers.Depth[0]),
		PublicKey:   pub,
		FingerPrint: fingerPrint,
		Xpubkey:     xpub,
		buffers:     buffers,
	}, nil
}

// DeriveChild derives a child public key (non-hardened only).
func (k *HDPublicKey) DeriveChild(arg interface{}, hardened bool) (*HDPublicKey, error) {
	switch v := arg.(type) {
	case int:
		return k.deriveWithNumber(v, hardened)
	case string:
		return k.deriveFromString(v)
	default:
		return nil, newError("InvalidDerivationArgument", "invalid derivation argument")
	}
}

func (k *HDPublicKey) deriveWithNumber(index int, hardened bool) (*HDPublicKey, error) {
	if index >= HDHardened || hardened {
		return nil, newError("InvalidIndexCantDeriveHardened", "cannot derive hardened child from public key")
	}
	if index < 0 {
		return nil, newError("InvalidPath", "invalid path: %d", index)
	}

	indexBuffer := make([]byte, 4)
	binary.BigEndian.PutUint32(indexBuffer, uint32(index))
	data := concatBytes(k.PublicKey.ToDER(), indexBuffer)
	hash := (Hash{}).Sha512hmac(data, k.buffers.ChainCode)
	leftPart := NewBNFromBuffer(hash[0:32], false)
	chainCode := hash[32:64]

	var publicKey *PublicKey
	var err error
	point := GetG().Mul(leftPart.Int).Add(k.PublicKey.Point)
	publicKey, err = PublicKeyFromPoint(point, true)
	if err != nil {
		return k.deriveWithNumber(index+1, false)
	}

	version := make([]byte, 4)
	binary.BigEndian.PutUint32(version, k.Network.Xpubkey)
	parentFp := make([]byte, 4)
	copy(parentFp, k.FingerPrint)
	childIdx := make([]byte, 4)
	binary.BigEndian.PutUint32(childIdx, uint32(index))

	buffers := &hdBuffers{
		Version:          version,
		Depth:            []byte{byte(k.Depth + 1)},
		ParentFingerPrint: parentFp,
		ChildIndex:       childIdx,
		ChainCode:        chainCode,
		PublicKey:        publicKey.ToDER(),
	}
	return hdPublicKeyFromBuffers(buffers, "")
}

func (k *HDPublicKey) deriveFromString(path string) (*HDPublicKey, error) {
	if strings.Contains(path, "'") {
		return nil, newError("InvalidIndexCantDeriveHardened", "cannot derive hardened child from public key")
	}
	if !IsValidHDPath(path, false) {
		return nil, newError("InvalidPath", "invalid path: %s", path)
	}
	indexes := getDerivationIndexes(path)
	derived := k
	var err error
	for _, index := range indexes {
		derived, err = derived.deriveWithNumber(index, false)
		if err != nil {
			return nil, err
		}
	}
	return derived, nil
}

// ToString returns the xpub serialization.
func (k *HDPublicKey) ToString() string { return k.Xpubkey }

// ToObject returns a plain object representation.
func (k *HDPublicKey) ToObject() map[string]interface{} {
	return map[string]interface{}{
		"network":           k.Network.Name,
		"depth":             k.Depth,
		"fingerPrint":       binary.BigEndian.Uint32(k.FingerPrint),
		"parentFingerPrint": binary.BigEndian.Uint32(k.buffers.ParentFingerPrint),
		"childIndex":        binary.BigEndian.Uint32(k.buffers.ChildIndex),
		"chainCode":         toHex(k.buffers.ChainCode),
		"publicKey":         k.PublicKey.ToString(),
		"checksum":          binary.BigEndian.Uint32(k.buffers.Checksum),
		"xpubkey":           k.Xpubkey,
	}
}

// ToBuffer returns the xpub string as bytes.
func (k *HDPublicKey) ToBuffer() []byte { return []byte(k.Xpubkey) }

// ToHex returns the hex of the xpub string.
func (k *HDPublicKey) ToHex() string { return toHex([]byte(k.Xpubkey)) }