package tron

import (
  "encoding/hex"

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

// PrivateKey is a validated 32-byte secp256k1 private key.
type PrivateKey struct {
  raw []byte
  Pub *PublicKey
}

// PublicKey is an uncompressed secp256k1 public key (64 bytes, x || y).
type PublicKey struct {
  raw []byte
}

// NewPrivateKey validates and wraps a 32-byte private key.
func NewPrivateKey(b []byte) (*PrivateKey, error) {
  if _, err := secp256k1.PublicKey(b); err != nil {
    return nil, ErrBadKey
  }
  return newPrivateKeyUnchecked(b), nil
}

func newPrivateKeyUnchecked(b []byte) *PrivateKey {
  p, _ := secp256k1.PublicKey(b)
  raw := make([]byte, 64)
  p.X.FillBytes(raw[0:32])
  p.Y.FillBytes(raw[32:64])
  return &PrivateKey{raw: append([]byte{}, b...), Pub: &PublicKey{raw: raw}}
}

// FromHex parses a hex-encoded private key.
func PrivateKeyFromHex(s string) (*PrivateKey, error) {
  b, err := hex.DecodeString(s)
  if err != nil {
    return nil, ErrBadKey
  }
  return NewPrivateKey(b)
}

// Bytes returns the 32-byte private key.
func (k *PrivateKey) Bytes() []byte {
  return append([]byte{}, k.raw...)
}

// Hex returns the private key as a hex string.
func (k *PrivateKey) Hex() string {
  return hex.EncodeToString(k.raw)
}

// Bytes returns the uncompressed public key (64 bytes).
func (p *PublicKey) Bytes() []byte {
  return append([]byte{}, p.raw...)
}

// Hex returns the uncompressed public key as a hex string.
func (p *PublicKey) Hex() string {
  return hex.EncodeToString(p.raw)
}

// ToBase58CheckAddress returns the TRON base58check address of the key.
func (p *PublicKey) ToBase58CheckAddress() string {
  h := keccak.Sum256(p.raw)
  payload := append([]byte{0x41}, h[12:]...)
  return base58.CheckEncode(payload)
}

// ToHexAddress returns the 41-prefixed hex address of the key.
func (p *PublicKey) ToHexAddress() string {
  h := keccak.Sum256(p.raw)
  return "41" + hex.EncodeToString(h[12:])
}

// Sign signs a 32-byte message hash (e.g. a transaction id) with the key.
func (k *PrivateKey) Sign(hash []byte) ([]byte, error) {
  sig, err := secp256k1.Sign(k.raw, hash)
  if err != nil {
    return nil, err
  }
  return sig.Bytes(), nil
}

// GenerateAddressInfo collects all display fields for a private key.
func GenerateAddressInfo(k *PrivateKey) map[string]string {
  return map[string]string{
    "base58check_address": k.Pub.ToBase58CheckAddress(),
    "hex_address":         k.Pub.ToHexAddress(),
    "private_key":         k.Hex(),
    "public_key":          k.Pub.Hex(),
  }
}