// Package gnarl implements Schnorr signatures over the non-split torus of
// SL(2, Z_P) for a 216-bit prime P.
//
// Uses a 216-bit prime giving 27-byte field elements and ~107-bit
// Pollard-rho security.
//
// Key sizes:
//   - Secret key: 27 bytes (scalar mod Q)
//   - Public key: 27 bytes (torus y-coordinate, y < P/3)
//   - Signature:  54 bytes (27-byte challenge + 27-byte response)
//
// The challenge hash uses GnarlMid (27-byte trinary Hamadryad hash) from
// pkg/crypto, binding the signature to the trinary hash primitive.
//
// Prime P satisfies:
//   - P ≡ 2 mod 3 (unique cube roots, eliminates trinary x-only ambiguity)
//   - 5 is QNR mod P (non-split torus condition)
//   - N = P+1 = 6Q where Q is a ~213-bit prime
//
// The constraint P ≡ 2 mod 3 means there are 3 cosets of the index-6
// subgroup of order Q inside the torus, so we canonicalize the public key
// y-coordinate to [0, P/3) by trying s, 2s, Q-s, etc.
package gnarl

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

// P, Q, N as big.Int for subgroup membership checks and scalar reduction.
var (
	pBig *big.Int // field prime
	qBig *big.Int // subgroup order (prime), Q = (P+1)/6
	nBig *big.Int // torus order = P + 1 = 6Q
)

// SchnorrGen is the Schnorr signature generator on the non-split torus.
var schnorrGenTM tmat

func init() {
	initPrime()
	initGenerators()
}

func initPrime() {
	seed := sha256.Sum256([]byte("gnarl-bethe-sl2-dendrite-v1"))

	one := big.NewInt(1)
	two := big.NewInt(2)
	five := big.NewInt(5)
	six := big.NewInt(6)

	// Q starting point from seed.
	// Take first 27 bytes, mask to 213 bits, set bit 212, make odd.
	maxQ := new(big.Int).Div(new(big.Int).Lsh(one, 216), six)
	q := new(big.Int).SetBytes(seed[:27])
	qMask := new(big.Int).Sub(new(big.Int).Lsh(one, 213), one)
	q.And(q, qMask)
	q.SetBit(q, 212, 1)
	if q.Bit(0) == 0 {
		q.Add(q, one)
	}

	for {
		if q.Cmp(maxQ) > 0 {
			panic("gnarl: prime search exhausted")
		}

		// Q mod 5 must be 3 or 4 for P = 6Q-1 to have 5 as QNR.
		qm5 := new(big.Int).Mod(q, five).Int64()
		if qm5 == 3 || qm5 == 4 {
			if q.ProbablyPrime(1) {
				p := new(big.Int).Mul(six, q)
				p.Sub(p, one)
				if p.ProbablyPrime(1) {
					if q.ProbablyPrime(20) && p.ProbablyPrime(20) {
						// Verify 5 is QNR mod P.
						pm1 := new(big.Int).Sub(p, one)
						exp := new(big.Int).Rsh(pm1, 1)
						euler := new(big.Int).Exp(five, exp, p)
						if euler.Cmp(pm1) == 0 {
							pBig = p
							qBig = q
							nBig = new(big.Int).Add(p, one)

							// Verify the derived constants match our hardcoded values.
							verifyConstants(p, q)
							return
						}
					}
				}
			}
		}
		q.Add(q, two)
	}
}

func verifyConstants(p, q *big.Int) {
	// Verify pLimbs matches P.
	pCheck := new(big.Int)
	pCheck.SetUint64(pLimbs[3])
	pCheck.Lsh(pCheck, 64)
	pCheck.Or(pCheck, new(big.Int).SetUint64(pLimbs[2]))
	pCheck.Lsh(pCheck, 64)
	pCheck.Or(pCheck, new(big.Int).SetUint64(pLimbs[1]))
	pCheck.Lsh(pCheck, 64)
	pCheck.Or(pCheck, new(big.Int).SetUint64(pLimbs[0]))
	if pCheck.Cmp(p) != 0 {
		panic("gnarl: pLimbs mismatch")
	}

	// Verify qLimbs matches Q.
	qCheck := new(big.Int)
	qCheck.SetUint64(qLimbs[3])
	qCheck.Lsh(qCheck, 64)
	qCheck.Or(qCheck, new(big.Int).SetUint64(qLimbs[2]))
	qCheck.Lsh(qCheck, 64)
	qCheck.Or(qCheck, new(big.Int).SetUint64(qLimbs[1]))
	qCheck.Lsh(qCheck, 64)
	qCheck.Or(qCheck, new(big.Int).SetUint64(qLimbs[0]))
	if qCheck.Cmp(q) != 0 {
		panic("gnarl: qLimbs mismatch")
	}
}

func initGenerators() {
	// g0 = [[1,1],[0,1]], g1 = [[1,0],[1,1]]
	g0 := m4FromSmall(1, 1, 0, 1)
	g1 := m4FromSmall(1, 0, 1, 1)

	// g0*g1 = [[2,1],[1,1]], which is on the torus (B = C = 1, symmetric).
	// This element has order dividing N = P+1 = 6Q.
	// To get an element of exact order Q, raise to the 6th power: G = (g0*g1)^6.
	var g01 mat4
	m4Mul(&g01, &g0, &g1)

	// (g0*g1)^2
	var g01sq mat4
	m4Mul(&g01sq, &g01, &g01)

	// (g0*g1)^3
	var g01cu mat4
	m4Mul(&g01cu, &g01sq, &g01)

	// (g0*g1)^6
	var gen mat4
	m4Mul(&gen, &g01cu, &g01cu)

	// Store as torus matrix (B should equal C for this generator).
	tmFromMat4(&schnorrGenTM, &gen)

	// Verify it's not the identity (it should have order Q).
	if m4IsIdentity(&gen) {
		panic("gnarl: generator is identity")
	}

	// Initialize precomputed exponentiation table.
	initGenTable(&schnorrGenTM)
}

// PrivateKey is a scalar in [1, Q).
type PrivateKey struct {
	s scalar
}

// PublicKey is a torus y-coordinate (27 bytes, y < P/3).
type PublicKey struct {
	tm tmat
}

// Signature is (e, z) where e is the 27-byte challenge and z is the 27-byte response.
type Signature struct {
	e [27]byte // challenge (GnarlMid hash output)
	z scalar   // response scalar
}

// GenerateKey creates a new keypair.
func GenerateKey() (*PrivateKey, *PublicKey, error) {
	var s scalar
	if err := scRandom(&s); err != nil {
		return nil, nil, err
	}

	// PK = SchnorrGen^s
	var pkTM tmat
	tmFixedExp(&pkTM, &s)

	// Canonicalize y: find an equivalent s such that y < P/3.
	// The torus has order 6Q, and Q is the subgroup order.
	// If y >= P/3, try negating (s -> Q - s gives inverse: y -> -y).
	// If still >= P/3, try doubling (s -> 2s mod Q).
	// With 3 cosets in [0,P) split into [0,P/3), [P/3,2P/3), [2P/3,P),
	// exactly one of {s, Q-s, ...} will land in [0, P/3).
	//
	// Simpler approach: just try s and Q-s; one of them should give y < P/3
	// with high probability for a uniformly random s. If neither works,
	// we can multiply by a cube root of unity, but for P ≡ 2 mod 3 there
	// are no non-trivial cube roots in Z_P, so -y is the only other option.
	// Since P ≡ 2 mod 3, exactly (P-1)/3 values are in [0, P/3) and
	// (P-1)/3 more in [P/3, 2P/3). So y and P-y together cover 2 of the
	// 3 thirds. We need the third case: if y is in [2P/3, P), then P-y is
	// in [0, P/3). So at most one negation suffices.
	var bNorm fe
	feFromMont(&bNorm, &pkTM.b)
	if feIsGtThirdP(&bNorm) == 1 {
		// Try negating: s -> Q - s, which inverts the matrix.
		scNeg(&s, &s)
		tmInv(&pkTM, &pkTM)

		// Check again.
		feFromMont(&bNorm, &pkTM.b)
		if feIsGtThirdP(&bNorm) == 1 {
			// y is in the middle third [P/3, 2P/3). P-y is in (P/3, 2P/3).
			// Neither is < P/3. This shouldn't happen: if y is in [P/3, 2P/3)
			// then P-y is in (P/3, 2P/3] which is also > P/3.
			// But y = 0 maps to itself, and y in (P/3, 2P/3) maps to (P/3, 2P/3).
			// So we accept the middle third as canonical with the constraint y < P/2.
			feFromMont(&bNorm, &pkTM.b)
			if feIsGtHalfP(&bNorm) == 1 {
				scNeg(&s, &s)
				tmInv(&pkTM, &pkTM)
			}
		}
	}

	sk := &PrivateKey{s: s}
	pk := &PublicKey{tm: pkTM}
	return sk, pk, nil
}

// PrivateKeyFromBytes deserializes a 27-byte secret key.
func PrivateKeyFromBytes(data []byte) (*PrivateKey, error) {
	if len(data) < 27 {
		return nil, errors.New("gnarl: short private key")
	}
	var s scalar
	scFromBytes27(&s, data)
	if scIsZero(&s) {
		return nil, errors.New("gnarl: zero private key")
	}
	return &PrivateKey{s: s}, nil
}

// Bytes serializes the private key as 27 bytes.
func (sk *PrivateKey) Bytes() []byte {
	buf := make([]byte, 27)
	scToBytes27(buf, &sk.s)
	return buf
}

// PublicKeyFromPrivate derives the public key from a private key.
func PublicKeyFromPrivate(sk *PrivateKey) *PublicKey {
	var pkTM tmat
	tmFixedExp(&pkTM, &sk.s)
	return &PublicKey{tm: pkTM}
}

// YBytes serializes the public key as 27 bytes (y-coordinate only).
func (pk *PublicKey) YBytes() []byte {
	buf := make([]byte, 27)
	feToBytes27(buf, &pk.tm.b)
	return buf
}

// FullBytes serializes the torus matrix (a, b, d) as 81 bytes.
func (pk *PublicKey) FullBytes() []byte {
	buf := make([]byte, 81)
	tmToBytes(buf, &pk.tm)
	return buf
}

// PublicKeyFromYBytes reconstructs a public key from a 27-byte y-coordinate.
func PublicKeyFromYBytes(data []byte) (*PublicKey, error) {
	if len(data) < 27 {
		return nil, errors.New("gnarl: short public key")
	}

	var yfe fe
	if !feFromBytes27(&yfe, data) {
		return nil, errors.New("gnarl: y >= P")
	}

	// x^2 = 1 + 5*y^2/4 mod P
	var y2, five, inv4, fiveY2o4, x2, xfe fe
	montSquare(&y2, &yfe)
	feFromSmall(&five, 5)
	feFromSmall(&inv4, 4)
	feInv(&inv4, &inv4)
	montMul(&fiveY2o4, &five, &y2)
	montMul(&fiveY2o4, &fiveY2o4, &inv4)
	feAdd(&x2, &feOne, &fiveY2o4)

	if !feSqrt(&xfe, &x2) {
		return nil, errors.New("gnarl: no square root (y not on torus)")
	}

	// M = [[x + y/2, y], [y, x - y/2]]
	var inv2, yHalf, afe, dfe fe
	feFromSmall(&inv2, 2)
	feInv(&inv2, &inv2)
	montMul(&yHalf, &yfe, &inv2)
	feAdd(&afe, &xfe, &yHalf)
	feSub(&dfe, &xfe, &yHalf)

	// Check subgroup membership: M^Q = I.
	tm1 := tmat{a: afe, b: yfe, d: dfe}
	var m4check, m4cand mat4
	tmToMat4(&m4cand, &tm1)
	m4PowBig(&m4check, &m4cand, qBig)
	if m4IsIdentity(&m4check) {
		return &PublicKey{tm: tm1}, nil
	}

	// Try the other root (-x).
	var xNeg fe
	feNeg(&xNeg, &xfe)
	feAdd(&afe, &xNeg, &yHalf)
	feSub(&dfe, &xNeg, &yHalf)

	tm2 := tmat{a: afe, b: yfe, d: dfe}
	tmToMat4(&m4cand, &tm2)
	m4PowBig(&m4check, &m4cand, qBig)
	if m4IsIdentity(&m4check) {
		return &PublicKey{tm: tm2}, nil
	}

	return nil, errors.New("gnarl: neither root in subgroup")
}

// Sign produces a Gnarl signature on msg.
// The challenge hash is GMid(msg || R.FullBytes()), using the trinary
// Hamadryad hash. The challengeFunc parameter allows the caller to
// provide the hash function (avoiding a circular import with pkg/crypto).
func Sign(sk *PrivateKey, msg []byte, challengeFunc func([]byte) [27]byte) (*Signature, error) {
	// k <- random in [1, Q)
	var k scalar
	if err := scRandom(&k); err != nil {
		return nil, err
	}

	// R = SchnorrGen^k
	var rTM tmat
	tmFixedExp(&rTM, &k)

	// Serialize R for challenge hash.
	rBytes := make([]byte, 81)
	tmToBytes(rBytes, &rTM)

	// e = GMid(msg || R.FullBytes())
	hashInput := make([]byte, len(msg)+81)
	copy(hashInput, msg)
	copy(hashInput[len(msg):], rBytes)
	eBytes := challengeFunc(hashInput)

	// z = (k - e*s) mod Q
	var eSc, es, z scalar
	scFromBytes27(&eSc, eBytes[:])
	scMul(&es, &eSc, &sk.s)
	scSub(&z, &k, &es)

	return &Signature{e: eBytes, z: z}, nil
}

// Verify checks a Gnarl signature.
func Verify(pk *PublicKey, msg []byte, sig *Signature, challengeFunc func([]byte) [27]byte) bool {
	// R' = SchnorrGen^z * PK^e
	var eSc scalar
	scFromBytes27(&eSc, sig.e[:])

	var rPrime tmat
	tmShamirExp(&rPrime, &sig.z, &pk.tm, &eSc)

	// Serialize R' and compute challenge.
	rBytes := make([]byte, 81)
	tmToBytes(rBytes, &rPrime)

	hashInput := make([]byte, len(msg)+81)
	copy(hashInput, msg)
	copy(hashInput[len(msg):], rBytes)
	ePrime := challengeFunc(hashInput)

	return sig.e == ePrime
}

// SignatureBytes serializes a signature as 54 bytes (27 challenge + 27 response).
func (sig *Signature) Bytes() []byte {
	buf := make([]byte, 54)
	copy(buf[:27], sig.e[:])
	scToBytes27(buf[27:], &sig.z)
	return buf
}

// SignatureFromBytes deserializes a 54-byte signature.
func SignatureFromBytes(data []byte) (*Signature, error) {
	if len(data) < 54 {
		return nil, errors.New("gnarl: short signature")
	}
	sig := &Signature{}
	copy(sig.e[:], data[:27])
	scFromBytes27(&sig.z, data[27:54])
	return sig, nil
}

// Params returns the scheme parameters.
func Params() (prime, subgroupOrder, torusOrder *big.Int) {
	return new(big.Int).Set(pBig), new(big.Int).Set(qBig), new(big.Int).Set(nBig)
}