package secp256k1

import "crypto/sha256"

// BIP-340 protocol sizes.
const (
	PubKeyBytesLen = 32
	SignatureSize  = 64
)

// ValidateSecretKey returns true if sk is a valid secp256k1 scalar (32 bytes,
// in [1, n-1]). Callers wanting a fresh random key should loop reading bytes
// from their preferred source until ValidateSecretKey returns true; this
// package intentionally does not import crypto/rand so WASM consumers that
// only need verification do not link unused entropy machinery.
func ValidateSecretKey(sk []byte) bool {
	if len(sk) != 32 {
		return false
	}
	k := scalarFromBytes(sk)
	return !scalarIsZero(k) && feCmp(k, curveN()) < 0
}

// ValidatePubKey returns true if pubkey is a valid 32-byte BIP-340 x-only
// public key: correct length, x < p, and the x-coordinate lifts to a point on
// the curve.
func ValidatePubKey(pubkey []byte) bool {
	if len(pubkey) != PubKeyBytesLen {
		return false
	}
	x := feFromBytes(pubkey)
	if feCmp(x, _feP()) >= 0 {
		return false
	}
	_, ok := LiftX(x)
	return ok
}

// VerifySchnorr verifies a BIP-340 Schnorr signature over a 32-byte message hash.
func VerifySchnorr(pubkey, msg [32]byte, sig [64]byte) bool {
	return VerifySchnorrBytes(pubkey, msg[:], sig)
}

// VerifySchnorrBytes verifies a BIP-340 Schnorr signature over a variable-length
// message. BIP-340 (post-2022) defines signing/verification over arbitrary-length
// input; Nostr always passes a 32-byte sha256 hash but the spec is general.
func VerifySchnorrBytes(pubkey [32]byte, msg []byte, sig [64]byte) bool {
	px := feFromBytes(pubkey[:])
	P, ok := LiftX(px)
	if !ok {
		return false
	}

	rx := feFromBytes(sig[:32])
	s := scalarFromBytes(sig[32:])

	// BIP-340 explicit reject: sig[0:32] >= p or sig[32:64] >= n.
	if feCmp(rx, _feP()) >= 0 {
		return false
	}
	if feCmp(s, curveN()) >= 0 {
		return false
	}

	e := computeChallenge(sig[:32], pubkey[:], msg)

	sG := ScalarBaseMult(s)
	eP := ScalarMult(pointFromAffine(P), e).toAffine()

	negEP := AffinePoint{eP.X, feNeg(eP.Y)}
	R := pointAdd(pointFromAffine(sG), pointFromAffine(negEP)).toAffine()

	if feIsZero(R.X) && feIsZero(R.Y) {
		return false
	}
	if R.X != rx {
		return false
	}
	if !feIsEven(R.Y) {
		return false
	}
	return true
}

// SignSchnorr creates a BIP-340 Schnorr signature over a 32-byte message hash.
func SignSchnorr(seckey, msg, auxRand [32]byte) (sig [64]byte, ok bool) {
	return SignSchnorrBytes(seckey, msg[:], auxRand)
}

// SignSchnorrBytes signs a variable-length message per BIP-340.
func SignSchnorrBytes(seckey [32]byte, msg []byte, auxRand [32]byte) (sig [64]byte, ok bool) {
	sk := scalarFromBytes(seckey[:])
	if scalarIsZero(sk) || feCmp(sk, curveN()) >= 0 {
		return sig, false
	}

	P := ScalarBaseMult(sk)

	d := sk
	if !feIsEven(P.Y) {
		d = scalarNeg(d)
	}

	pkBytes := feToBytes(P.X)

	var t [32]byte
	aux := taggedHash("BIP0340/aux", auxRand[:])
	dBytes := feToBytes(d)
	for i := 0; i < 32; i++ {
		t[i] = dBytes[i] ^ aux[i]
	}

	// k' = tagged_hash("BIP0340/nonce", t || pkBytes || msg) mod n
	nonceInput := []byte{:0:64+len(msg)}
	nonceInput = append(nonceInput, t[:]...)
	nonceInput = append(nonceInput, pkBytes[:]...)
	nonceInput = append(nonceInput, msg...)
	kHash := taggedHash("BIP0340/nonce", nonceInput)
	k := scalarFromBytes(kHash[:])

	if feCmp(k, curveN()) >= 0 {
		k = scalarSub(k, curveN())
	}
	if scalarIsZero(k) {
		return sig, false
	}

	R := ScalarBaseMult(k)

	if !feIsEven(R.Y) {
		k = scalarNeg(k)
	}

	rxBytes := feToBytes(R.X)

	e := computeChallenge(rxBytes[:], pkBytes[:], msg)

	s := scalarAdd(k, scalarMul(e, d))
	sBytes := feToBytes(s)

	copy(sig[:32], rxBytes[:])
	copy(sig[32:], sBytes[:])
	return sig, true
}

// PubKeyFromSecKey derives the x-only public key from a secret key.
// Returns zero [32]byte and false on failure.
func PubKeyFromSecKey(seckey [32]byte) (pubkey [32]byte, ok bool) {
	sk := scalarFromBytes(seckey[:])
	if scalarIsZero(sk) {
		return pubkey, false
	}
	P := ScalarBaseMult(sk)
	return feToBytes(P.X), true
}

// ECDH computes the shared secret x-coordinate: seckey * pubkey.
// pubkey is a 32-byte x-only public key.
// Returns the 32-byte x-coordinate and true on success.
func ECDH(seckey, pubkey [32]byte) ([32]byte, bool) {
	sk := scalarFromBytes(seckey[:])
	if scalarIsZero(sk) {
		return [32]byte{}, false
	}
	px := feFromBytes(pubkey[:])
	P, ok := LiftX(px)
	if !ok {
		return [32]byte{}, false
	}
	R := ScalarMult(pointFromAffine(P), sk).toAffine()
	if feIsZero(R.X) && feIsZero(R.Y) {
		return [32]byte{}, false
	}
	return feToBytes(R.X), true
}

// ScalarAddModN returns (a + b) mod n where n is the secp256k1 group order.
// a and b are 32-byte big-endian scalars.
func ScalarAddModN(a, b [32]byte) ([32]byte, bool) {
	sa := scalarFromBytes(a[:])
	sb := scalarFromBytes(b[:])
	r := scalarAdd(sa, sb)
	if feCmp(r, curveN()) >= 0 {
		r = scalarSub(r, curveN())
	}
	if scalarIsZero(r) {
		return [32]byte{}, false
	}
	return feToBytes(r), true
}

// CompressedPubKey returns the 33-byte SEC1 compressed public key for a secret key.
// Format: 0x02 or 0x03 prefix (even/odd Y) followed by 32-byte X coordinate.
func CompressedPubKey(seckey [32]byte) ([33]byte, bool) {
	sk := scalarFromBytes(seckey[:])
	if scalarIsZero(sk) {
		return [33]byte{}, false
	}
	P := ScalarBaseMult(sk)
	xb := feToBytes(P.X)
	yb := feToBytes(P.Y)
	var out [33]byte
	if yb[31]&1 == 0 {
		out[0] = 0x02
	} else {
		out[0] = 0x03
	}
	copy(out[1:], xb[:])
	return out, true
}

// tagged_hash(tag, msg) = SHA256(SHA256(tag) || SHA256(tag) || msg)
func taggedHash(tag string, msg []byte) [32]byte {
	tagHash := sha256.Sum([]byte(tag))
	data := []byte{:0:64+len(msg)}
	data = append(data, tagHash[:]...)
	data = append(data, tagHash[:]...)
	data = append(data, msg...)
	return sha256.Sum(data)
}

func computeChallenge(rx, px, msg []byte) Fe {
	input := []byte{:0:96}
	input = append(input, rx...)
	input = append(input, px...)
	input = append(input, msg...)
	hash := taggedHash("BIP0340/challenge", input)
	e := scalarFromBytes(hash[:])
	// Reduce mod n.
	if feCmp(e, curveN()) >= 0 {
		e = scalarSub(e, curveN())
	}
	return e
}