// Package hd is BIP32 private derivation, and the account paths a swap wallet uses.
//
// It derives from a seed and never from a public key: the wallet holds its own keys, so it
// has no reason to derive a child it cannot spend, and private derivation is the smaller
// surface. A future watch-only export needs public derivation and is not in this cut.
package hd

import (
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/binary"
	"fmt"
	"strconv"
	"strings"

	"git.smesh.lol/nostrino/pkg/crypto/decred/secp256k1"
	"golang.org/x/crypto/ripemd160"
)

// Hardened is the index offset that marks a hardened child, as BIP32 defines it.
const Hardened uint32 = 0x80000000

// Key is one node of a BIP32 tree: a private scalar and the chain code that derives the
// next level.
type Key struct {
	// Private is the 32-byte scalar.
	Private [32]byte
	// ChainCode is the 32-byte chain code, which is what makes a child unique to its
	// parent even when two parents share a key.
	ChainCode [32]byte
	// Depth is 0 at the master key.
	Depth byte
	// ParentFingerprint is the first four bytes of the parent's hash160.
	ParentFingerprint [4]byte
	// Child is the index this key was derived at, zero at the master key.
	Child uint32
	// Path is every index from the master key down to this one. It is carried rather than
	// looked up, because a key that cannot say where it came from cannot be written down
	// and found again, and a wallet's addresses are exactly that.
	Path []uint32
}

// NewMaster derives the master key from a BIP39 seed.
//
// The seed is 16 to 64 bytes per BIP32; a BIP39 seed is 64.
func NewMaster(seed []byte) (*Key, error) {
	if len(seed) < 16 || len(seed) > 64 {
		return nil, fmt.Errorf("a BIP32 seed is 16 to 64 bytes, not %d", len(seed))
	}

	mac := hmac.New(sha512.New, []byte("Bitcoin seed"))
	mac.Write(seed)

	sum := mac.Sum(nil)

	key, err := keyFromParts(sum[:32], sum[32:], 0, [4]byte{}, 0, nil)

	if err != nil {
		return nil, fmt.Errorf("this seed has no master key: %w", err)
	}

	return key, nil
}

// keyFromParts validates a private scalar and assembles a key.
func keyFromParts(
	private, chainCode []byte, depth byte, parent [4]byte, child uint32, path []uint32,
) (*Key, error) {
	if len(private) != 32 || len(chainCode) != 32 {
		return nil, fmt.Errorf("a key is a 32-byte scalar and a 32-byte chain code")
	}

	var scalar secp256k1.ModNScalar

	if overflow := scalar.SetByteSlice(private); overflow {
		return nil, fmt.Errorf("the scalar is not below the curve order")
	}

	if scalar.IsZero() {
		return nil, fmt.Errorf("the scalar is zero")
	}

	out := &Key{Depth: depth, ParentFingerprint: parent, Child: child}

	if len(path) > 0 {
		out.Path = append([]uint32(nil), path...)
	}

	copy(out.ChainCode[:], chainCode)

	scalar.PutBytesUnchecked(out.Private[:])

	// The scalar is stored as the value the curve uses, so a caller that reads Private
	// gets the canonical 32 bytes whatever the input's leading zeroes were.
	return out, nil
}

// Valid reports whether a key is usable.
func (k *Key) Valid() bool { return k != nil }

// Public returns the 33-byte compressed public key.
func (k *Key) Public() []byte {
	private := secp256k1.PrivKeyFromBytes(k.Private[:])

	return private.PubKey().SerializeCompressed()
}

// ChildKey derives one child key.
//
// A hardened index (at least 2^31) is derived from the private key; a normal index from the
// parent's public key. Both are the same arithmetic on the parent's scalar, which is why
// this package never needs point addition.
func (k *Key) ChildKey(index uint32) (*Key, error) {
	mac := hmac.New(sha512.New, k.ChainCode[:])

	if index >= Hardened {
		mac.Write([]byte{0})
		mac.Write(k.Private[:])
	} else {
		mac.Write(k.Public())
	}

	var encoded [4]byte

	binary.BigEndian.PutUint32(encoded[:], index)

	mac.Write(encoded[:])

	sum := mac.Sum(nil)

	var offset secp256k1.ModNScalar

	if overflow := offset.SetByteSlice(sum[:32]); overflow {
		// Roughly one seed in 2^127. BIP32 says to derive the next index; a caller with a
		// fixed path cannot, so it is an error and the wallet's index loop can skip it.
		return nil, fmt.Errorf("child %d is unrepresentable at this path", index)
	}

	var parent secp256k1.ModNScalar

	parent.SetByteSlice(k.Private[:])

	var child secp256k1.ModNScalar

	child.Add2(&offset, &parent)

	if child.IsZero() {
		return nil, fmt.Errorf("child %d derives the zero key", index)
	}

	var private [32]byte

	child.PutBytesUnchecked(private[:])

	next := &Key{
		Private:           private,
		Depth:             k.Depth + 1,
		ParentFingerprint: k.Fingerprint(),
		Child:             index,
		Path:              append(append([]uint32(nil), k.Path...), index),
	}

	copy(next.ChainCode[:], sum[32:])

	return next, nil
}

// Fingerprint is the first four bytes of the key's hash160, which identifies it as a
// parent in a serialised extended key.
func (k *Key) Fingerprint() [4]byte {
	hash := hash160(k.Public())

	var out [4]byte

	copy(out[:], hash[:4])

	return out
}

// hash160 is sha256 then ripemd160, the identifier a P2WPKH address carries.
func hash160(data []byte) [20]byte {
	first := sha256.Sum256(data)

	hasher := ripemd160.New()
	hasher.Write(first[:])

	var out [20]byte

	copy(out[:], hasher.Sum(nil))

	return out
}

// Path is a derivation path such as m/84'/0'/0'/0/0.
type Path []uint32

// String renders a path in the usual notation.
func (p Path) String() string {
	var out strings.Builder

	out.WriteString("m")

	for _, index := range p {
		out.WriteString("/")
		out.WriteString(strconv.FormatUint(uint64(index&^Hardened), 10))

		if index >= Hardened {
			out.WriteString("'")
		}
	}

	return out.String()
}

// ParsePath reads a derivation path.
//
// Both the apostrophe and an "h" are accepted for a hardened index, because a person who
// copies a path from another wallet has no way to know which one this client wants.
func ParsePath(text string) (Path, error) {
	trimmed := strings.TrimSpace(text)

	if trimmed == "" {
		return nil, fmt.Errorf("a derivation path is required")
	}

	parts := strings.Split(trimmed, "/")

	if parts[0] != "m" && parts[0] != "M" {
		return nil, fmt.Errorf("a derivation path starts with m: %q", text)
	}

	path := make(Path, 0, len(parts)-1)

	for _, part := range parts[1:] {
		if part == "" {
			return nil, fmt.Errorf("the derivation path has an empty level: %q", text)
		}

		hardened := false

		switch {
		case strings.HasSuffix(part, "'"):
			hardened = true
			part = strings.TrimSuffix(part, "'")
		case strings.HasSuffix(part, "h"), strings.HasSuffix(part, "H"):
			hardened = true
			part = part[:len(part)-1]
		}

		index, err := strconv.ParseUint(part, 10, 32)

		if err != nil {
			return nil, fmt.Errorf("the derivation level %q is not a number: %q", part, text)
		}

		if index >= uint64(Hardened) {
			return nil, fmt.Errorf("the derivation level %q is out of range", part)
		}

		value := uint32(index)

		if hardened {
			value |= Hardened
		}

		path = append(path, value)
	}

	return path, nil
}

// Derive walks a path from this key.
func (k *Key) Derive(path Path) (*Key, error) {
	current := k

	for _, index := range path {
		next, err := current.ChildKey(index)

		if err != nil {
			return nil, fmt.Errorf("%s at %s: %w", path, current.ChildPath(), err)
		}

		current = next
	}

	return current, nil
}

// ChildPath renders the path this key was derived at, for an error message or a record.
func (k *Key) ChildPath() string { return Path(k.Path).String() }