// Package bip32 implements the BIP-32 private child key derivation (CKDpriv)
// used to derive TRON keys from a BIP-39 seed. Only private derivation is
// required, so extended-key serialization is intentionally omitted.
package bip32

import (
  "crypto/hmac"
  "crypto/sha512"
  "encoding/binary"
  "errors"
  "math/big"
  "strconv"
  "strings"

  "code.luxferre.top/luxferre/kisstron/secp256k1"
)

// HardenedKeyStart is the BIP-32 hardened derivation offset (2^31).
const HardenedKeyStart uint32 = 0x80000000

// DefaultPath is the standard TRON BIP-44 derivation path.
const DefaultPath = "m/44'/195'/0'/0/0"

// ErrInvalidPath is returned for unparsable derivation paths.
var ErrInvalidPath = errors.New("bip32: invalid derivation path")

// ErrInvalidSeed is returned when the derived key is not usable.
var ErrInvalidSeed = errors.New("bip32: invalid seed")

// ParsePath decodes a BIP-32 path such as m/44'/195'/0'/0/0 into indices.
// Nodes may use ' or H (or h) to mark hardened derivation.
func ParsePath(path string) ([]uint32, error) {
  parts := strings.Split(path, "/")
  if len(parts) == 0 || (parts[0] != "m" && parts[0] != "M") {
    return nil, ErrInvalidPath
  }
  out := make([]uint32, 0, len(parts)-1)
  for _, part := range parts[1:] {
    if part == "" {
      return nil, ErrInvalidPath
    }
    hardened := false
    if last := part[len(part)-1]; last == '\'' || last == 'H' || last == 'h' {
      hardened = true
      part = part[:len(part)-1]
    }
    idx, err := strconv.ParseUint(part, 10, 31)
    if err != nil {
      return nil, ErrInvalidPath
    }
    v := uint32(idx)
    if hardened {
      v |= HardenedKeyStart
    }
    out = append(out, v)
  }
  return out, nil
}

// MasterKey derives the master private key and chain code from a seed.
func MasterKey(seed []byte) (key, chainCode []byte, err error) {
  mac := hmac.New(sha512.New, []byte("Bitcoin seed"))
  mac.Write(seed)
  sum := mac.Sum(nil)
  key = make([]byte, 32)
  chainCode = make([]byte, 32)
  copy(key, sum[:32])
  copy(chainCode, sum[32:])
  if new(big.Int).SetBytes(key).Sign() == 0 {
    return nil, nil, ErrInvalidSeed
  }
  return key, chainCode, nil
}

// DerivePath derives the private key at the given path from a seed.
func DerivePath(seed []byte, path string) ([]byte, error) {
  indices, err := ParsePath(path)
  if err != nil {
    return nil, err
  }
  key, chainCode, err := MasterKey(seed)
  if err != nil {
    return nil, err
  }
  for _, idx := range indices {
    key, chainCode, err = CKDpriv(key, chainCode, idx)
    if err != nil {
      return nil, err
    }
  }
  return key, nil
}

// CKDpriv computes the child private key and chain code for a single index.
func CKDpriv(key, chainCode []byte, index uint32) (childKey, childChain []byte, err error) {
  if len(key) != 32 || len(chainCode) != 32 {
    return nil, nil, ErrInvalidSeed
  }
  var data []byte
  if index >= HardenedKeyStart {
    data = append([]byte{0x00}, key...)
  } else {
    pub, perr := compressedPublicKey(key)
    if perr != nil {
      return nil, nil, perr
    }
    data = pub
  }
  var ser [4]byte
  binary.BigEndian.PutUint32(ser[:], index)
  data = append(data, ser[:]...)

  mac := hmac.New(sha512.New, chainCode)
  mac.Write(data)
  sum := mac.Sum(nil)

  il := new(big.Int).SetBytes(sum[:32])
  child := new(big.Int).Add(il, new(big.Int).SetBytes(key))
  child.Mod(child, secpOrder)
  if il.Cmp(secpOrder) >= 0 || child.Sign() == 0 {
    return nil, nil, ErrInvalidSeed
  }
  childKey = make([]byte, 32)
  child.FillBytes(childKey)
  childChain = make([]byte, 32)
  copy(childChain, sum[32:])
  return childKey, childChain, nil
}

// compressedPublicKey returns the 33-byte SEC1 compressed public key for a
// 32-byte private key.
func compressedPublicKey(key []byte) ([]byte, error) {
  p, err := secp256k1.PublicKey(key)
  if err != nil {
    return nil, err
  }
  out := make([]byte, 33)
  out[0] = 0x02
  if p.Y.Bit(0) == 1 {
    out[0] = 0x03
  }
  p.X.FillBytes(out[1:])
  return out, nil
}

var secpOrder, _ = new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141", 16)