// Package secp256k1 implements the secp256k1 elliptic curve arithmetic and
// recoverable ECDSA signing used by TRON. The Go standard library does not
// provide this curve, so it is built directly on math/big following SEC 1.
package secp256k1

import (
  "crypto/hmac"
  "crypto/sha256"
  "errors"
  "math/big"
)

// Curve parameters for secp256k1.
var (
  // P is the field prime.
  P, _ = new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F", 16)
  // N is the group order.
  N, _ = new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141", 16)
  // Gx, Gy are the coordinates of the base point.
  Gx, _ = new(big.Int).SetString("79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798", 16)
  Gy, _ = new(big.Int).SetString("483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8", 16)

  halfN = new(big.Int).Rsh(N, 1)
)

// Errors returned by the package.
var (
  ErrInvalidKey = errors.New("secp256k1: invalid private key")
  ErrInvalidSig = errors.New("secp256k1: invalid signature")
  ErrNoRecover  = errors.New("secp256k1: public key recovery failed")
)

// Point is an affine point on the curve. A nil point is the point at infinity.
type Point struct {
  X, Y *big.Int
}

func mod(x, m *big.Int) *big.Int {
  r := new(big.Int).Mod(x, m)
  if r.Sign() < 0 {
    r.Add(r, m)
  }
  return r
}

// Add returns p + q on the curve.
func Add(p, q *Point) *Point {
  if p == nil {
    return q
  }
  if q == nil {
    return p
  }
  if p.X.Cmp(q.X) == 0 {
    if p.Y.Cmp(q.Y) != 0 || p.Y.Sign() == 0 {
      return nil
    }
    return double(p)
  }
  num := mod(new(big.Int).Sub(q.Y, p.Y), P)
  den := mod(new(big.Int).Sub(q.X, p.X), P)
  den.ModInverse(den, P)
  lam := mod(new(big.Int).Mul(num, den), P)
  return applyLambda(p, q, lam)
}

func double(p *Point) *Point {
  num := new(big.Int).Mul(p.X, p.X)
  num.Mul(num, big.NewInt(3))
  num = mod(num, P)
  den := new(big.Int).Lsh(p.Y, 1)
  den = mod(den, P)
  den.ModInverse(den, P)
  lam := mod(new(big.Int).Mul(num, den), P)
  return applyLambda(p, p, lam)
}

func applyLambda(p, q *Point, lam *big.Int) *Point {
  x3 := new(big.Int).Mul(lam, lam)
  x3.Sub(x3, p.X)
  x3.Sub(x3, q.X)
  x3 = mod(x3, P)
  y3 := new(big.Int).Sub(p.X, x3)
  y3.Mul(lam, y3)
  y3.Sub(y3, p.Y)
  y3 = mod(y3, P)
  return &Point{X: x3, Y: y3}
}

// ScalarBaseMult returns k*G.
func ScalarBaseMult(k *big.Int) *Point {
  return ScalarMult(k, &Point{X: Gx, Y: Gy})
}

// ScalarMult returns k*p using a double-and-add ladder.
func ScalarMult(k *big.Int, p *Point) *Point {
  var r *Point
  if p == nil || k.Sign() == 0 {
    return nil
  }
  kk := mod(k, N)
  for i := kk.BitLen() - 1; i >= 0; i-- {
    r = Add(r, r)
    if kk.Bit(i) == 1 {
      r = Add(r, p)
    }
  }
  return r
}

// IsOnCurve reports whether p satisfies the curve equation.
func IsOnCurve(p *Point) bool {
  if p == nil || p.X == nil || p.Y == nil {
    return false
  }
  if p.X.Sign() < 0 || p.X.Cmp(P) >= 0 || p.Y.Sign() < 0 || p.Y.Cmp(P) >= 0 {
    return false
  }
  left := mod(new(big.Int).Mul(p.Y, p.Y), P)
  right := new(big.Int).Mul(p.X, p.X)
  right.Mul(right, p.X)
  right.Add(right, big.NewInt(7))
  right = mod(right, P)
  return left.Cmp(right) == 0
}

// PublicKey derives the public key point from a 32-byte private key.
func PublicKey(priv []byte) (*Point, error) {
  if len(priv) != 32 {
    return nil, ErrInvalidKey
  }
  k := new(big.Int).SetBytes(priv)
  if k.Sign() == 0 || k.Cmp(N) >= 0 {
    return nil, ErrInvalidKey
  }
  return ScalarBaseMult(k), nil
}

// Signature is a recoverable ECDSA signature (r, s) with recovery id V.
type Signature struct {
  R, S *big.Int
  V    byte
}

// Bytes returns the 65-byte compact representation r || s || v.
func (sig *Signature) Bytes() []byte {
  out := make([]byte, 65)
  sig.R.FillBytes(out[0:32])
  sig.S.FillBytes(out[32:64])
  out[64] = sig.V
  return out
}

// Sign signs a 32-byte message hash with the given private key, producing a
// deterministic recoverable signature (libsecp256k1-compatible nonce, low-s
// form). This matches coincurve/libsecp256k1's default recoverable signing.
func Sign(priv, hash []byte) (*Signature, error) {
  if len(priv) != 32 {
    return nil, ErrInvalidKey
  }
  if len(hash) != 32 {
    return nil, ErrInvalidSig
  }
  d := new(big.Int).SetBytes(priv)
  if d.Sign() == 0 || d.Cmp(N) >= 0 {
    return nil, ErrInvalidKey
  }
  e := mod(new(big.Int).SetBytes(hash), N)

  for {
    k := nonceRFC6979(priv, hash)
    kp := ScalarBaseMult(k)
    if kp == nil {
      continue
    }
    r := mod(kp.X, N)
    if r.Sign() == 0 {
      continue
    }
    kinv := new(big.Int).ModInverse(k, N)
    s := new(big.Int).Mul(r, d)
    s.Add(s, e)
    s.Mul(s, kinv)
    s.Mod(s, N)
    if s.Sign() == 0 {
      continue
    }
    v := byte(kp.Y.Bit(0))
    if s.Cmp(halfN) > 0 {
      s.Sub(N, s)
      v ^= 1
    }
    return &Signature{R: r, S: s, V: v}, nil
  }
}

// nonceRFC6979 reproduces libsecp256k1's default RFC 6979 nonce function. It
// feeds HMAC-SHA256 with the private key, the message and 48 zero bytes of
// optional data, exactly as secp256k1_nonce_function_rfc6979 does when no
// extra entropy is supplied.
func nonceRFC6979(priv, hash []byte) *big.Int {
  // libsecp256k1's rfc6979 default nonce feeds exactly key32||msg32 (no
  // extra data, no algo16), i.e. a 64-byte HMAC key.
  keydata := make([]byte, 64)
  copy(keydata[0:32], priv)
  copy(keydata[32:64], hash)

  v := make([]byte, 32)
  for i := range v {
    v[i] = 0x01
  }
  k := make([]byte, 32)

  k = hmacSHA256(k, v, []byte{0x00}, keydata[:])
  v = hmacSHA256(k, v)
  k = hmacSHA256(k, v, []byte{0x01}, keydata[:])
  v = hmacSHA256(k, v)

  for {
    v = hmacSHA256(k, v)
    candidate := new(big.Int).SetBytes(v)
    if candidate.Sign() > 0 && candidate.Cmp(N) < 0 {
      return candidate
    }
    k = hmacSHA256(k, v, []byte{0x00})
    v = hmacSHA256(k, v)
  }
}

func hmacSHA256(key []byte, parts ...[]byte) []byte {
  mac := hmac.New(sha256.New, key)
  for _, p := range parts {
    mac.Write(p)
  }
  return mac.Sum(nil)
}

// RecoverPublicKey recovers the public key point from a message hash and a
// recoverable signature.
func RecoverPublicKey(hash []byte, sig *Signature) (*Point, error) {
  if sig == nil || sig.R == nil || sig.S == nil || len(hash) != 32 {
    return nil, ErrInvalidSig
  }
  if sig.R.Sign() == 0 || sig.S.Sign() == 0 || sig.R.Cmp(N) >= 0 || sig.S.Cmp(N) >= 0 {
    return nil, ErrInvalidSig
  }
  x := new(big.Int).Set(sig.R)
  if sig.V&2 != 0 {
    x.Add(x, N)
  }
  if x.Cmp(P) >= 0 {
    return nil, ErrNoRecover
  }
  alpha := new(big.Int).Mul(x, x)
  alpha.Mul(alpha, x)
  alpha.Add(alpha, big.NewInt(7))
  alpha = mod(alpha, P)
  beta := sqrtMod(alpha)
  if beta == nil {
    return nil, ErrNoRecover
  }
  var y *big.Int
  if byte(beta.Bit(0)) == sig.V&1 {
    y = beta
  } else {
    y = new(big.Int).Sub(P, beta)
  }
  rp := &Point{X: x, Y: y}

  e := mod(new(big.Int).SetBytes(hash), N)
  rinv := new(big.Int).ModInverse(sig.R, N)
  u1 := mod(new(big.Int).Neg(e), N)
  u1.Mul(u1, rinv)
  u1.Mod(u1, N)
  u2 := mod(new(big.Int).Mul(sig.S, rinv), N)
  q := Add(ScalarBaseMult(u1), ScalarMult(u2, rp))
  if q == nil || !IsOnCurve(q) {
    return nil, ErrNoRecover
  }
  return q, nil
}

// sqrtMod returns a square root of a modulo P (P = 3 mod 4), or nil.
func sqrtMod(a *big.Int) *big.Int {
  exp := new(big.Int).Add(P, big.NewInt(1))
  exp.Rsh(exp, 2)
  r := new(big.Int).Exp(a, exp, P)
  if mod(new(big.Int).Mul(r, r), P).Cmp(mod(a, P)) != 0 {
    return nil
  }
  return r
}