package secp256k1

import (
	"crypto/rand"
	"crypto/sha256"
	"math/big"
)

// taggedHash implements the BIP-340 tagged hash:
// SHA256(SHA256(tag) || SHA256(tag) || parts...).
func taggedHash(tag string, parts ...[]byte) [32]byte {
	th := sha256.Sum256([]byte(tag))
	buf := make([]byte, 0, 64+96)
	buf = append(buf, th[:]...)
	buf = append(buf, th[:]...)
	for _, p := range parts {
		buf = append(buf, p...)
	}
	return sha256.Sum256(buf)
}

// challenge computes e = int(taggedHash("BIP0340/challenge", R.x || pubkey ||
// msg)) mod N.
func challenge(rx *big.Int, pub, msg []byte) *big.Int {
	var rb [32]byte
	rx.FillBytes(rb[:])
	h := taggedHash("BIP0340/challenge", rb[:], pub, msg)
	return new(big.Int).Mod(new(big.Int).SetBytes(h[:]), N)
}

// Verify reports whether sig is a valid BIP-340 signature of the 32-byte msg
// under the x-only public key pub. Nostr event IDs are always 32 bytes, so
// this is the fast path.
func Verify(pub [32]byte, msg [32]byte, sig [64]byte) bool {
	return verifyFast(pub, msg, sig)
}

// VerifyMsg is Verify for an arbitrary-length message. BIP-340 hashes the
// message as-is. The 32-byte case, which is the only one Nostr needs, uses the
// fast implementation; other lengths use the math/big reference below.
func VerifyMsg(pub [32]byte, msg []byte, sig [64]byte) bool {
	if len(msg) == 32 {
		var m [32]byte
		copy(m[:], msg)
		return verifyFast(pub, m, sig)
	}
	return verifySlowMsg(pub, msg, sig)
}

// verifySlowMsg is the math/big BIP-340 verifier. It remains the reference for
// arbitrary-length messages and the independent cross-check for the fast path.
func verifySlowMsg(pub [32]byte, msg []byte, sig [64]byte) bool {
	px := new(big.Int).SetBytes(pub[:])
	pubPoint := liftX(px)
	if pubPoint == nil {
		return false
	}
	r := new(big.Int).SetBytes(sig[:32])
	s := new(big.Int).SetBytes(sig[32:])
	if r.Cmp(P) >= 0 || s.Cmp(N) >= 0 {
		return false
	}
	e := challenge(r, pub[:], msg[:])

	// R = s*G - e*P = s*G + (N-e)*P.
	sG := scalarMult(s, fromAffine(G))
	ne := new(big.Int).Sub(N, e)
	eP := scalarMult(ne, fromAffine(*pubPoint))
	R := add(sG, eP)
	if R.isInfinity() {
		return false
	}
	ra := toAffine(R)
	if ra.y.Bit(0) == 1 {
		return false
	}
	return ra.x.Cmp(r) == 0
}

// PubKey returns the x-only public key for a secret key.
func PubKey(sk [32]byte) ([32]byte, bool) {
	var out [32]byte
	d := new(big.Int).SetBytes(sk[:])
	if d.Sign() == 0 || d.Cmp(N) >= 0 {
		return out, false
	}
	P := scalarMult(d, fromAffine(G))
	if P.isInfinity() {
		return out, false
	}
	toAffine(P).x.FillBytes(out[:])
	return out, true
}

// Sign produces a BIP-340 signature using fresh auxiliary randomness.
func Sign(sk [32]byte, msg [32]byte) ([64]byte, bool) {
	return SignMsg(sk, msg[:])
}

// SignMsg is Sign for an arbitrary-length message.
func SignMsg(sk [32]byte, msg []byte) ([64]byte, bool) {
	var aux [32]byte
	if _, err := rand.Read(aux[:]); err != nil {
		// A zero aux_rand is valid BIP-340; it makes the nonce deterministic.
		aux = [32]byte{}
	}
	if len(msg) == 32 {
		var m [32]byte
		copy(m[:], msg)
		return signFastWithAux(sk, m, aux)
	}
	return signSlowAuxMsg(sk, msg, aux)
}

// SignAux produces a BIP-340 signature with an explicit aux_rand value. It is
// what the official test vectors exercise.
func SignAux(sk [32]byte, msg [32]byte, aux [32]byte) ([64]byte, bool) {
	return SignAuxMsg(sk, msg[:], aux)
}

// SignAuxMsg is SignAux for an arbitrary-length message.
func SignAuxMsg(sk [32]byte, msg []byte, aux [32]byte) ([64]byte, bool) {
	if len(msg) == 32 {
		var m [32]byte
		copy(m[:], msg)
		return signFastWithAux(sk, m, aux)
	}
	return signSlowAuxMsg(sk, msg, aux)
}

// signSlowAuxMsg is the math/big signer. It reproduces the official vectors for
// every message length and cross-checks the fast signer in tests.
func signSlowAuxMsg(sk [32]byte, msg []byte, aux [32]byte) ([64]byte, bool) {
	var sig [64]byte
	d0 := new(big.Int).SetBytes(sk[:])
	if d0.Sign() == 0 || d0.Cmp(N) >= 0 {
		return sig, false
	}

	Pp := scalarMult(d0, fromAffine(G))
	if Pp.isInfinity() {
		return sig, false
	}
	Pa := toAffine(Pp)

	var pub [32]byte
	Pa.x.FillBytes(pub[:])

	// The secret key is negated when P has odd y, so that the effective
	// public key is the even-y lift that BIP-340 specifies.
	d := d0
	if Pa.y.Bit(0) == 1 {
		d = new(big.Int).Sub(N, d0)
	}

	ah := taggedHash("BIP0340/aux", aux[:])
	var db [32]byte
	d.FillBytes(db[:])
	var t [32]byte
	for i := range t {
		t[i] = db[i] ^ ah[i]
	}

	nh := taggedHash("BIP0340/nonce", t[:], pub[:], msg)
	k0 := new(big.Int).Mod(new(big.Int).SetBytes(nh[:]), N)
	if k0.Sign() == 0 {
		return sig, false
	}
	R := scalarMult(k0, fromAffine(G))
	if R.isInfinity() {
		return sig, false
	}
	Ra := toAffine(R)
	k := k0
	if Ra.y.Bit(0) == 1 {
		k = new(big.Int).Sub(N, k0)
	}
	e := challenge(Ra.x, pub[:], msg[:])
	s := new(big.Int).Mod(new(big.Int).Add(k, new(big.Int).Mul(e, d)), N)
	Ra.x.FillBytes(sig[:32])
	s.FillBytes(sig[32:])

	// Self-check: never emit a signature that would be rejected.
	if !verifySlowMsg(pub, msg, sig) {
		return sig, false
	}
	return sig, true
}