// MP12 gadget-trapdoor GPV signatures (Micciancio-Peikert 2012).
//
// This is a proper hash-then-sign lattice signature — fundamentally different
// from the Lyubashevsky Fiat-Shamir with Aborts in gpv.go (which is a sigma
// protocol turned into a signature via rejection sampling).
//
// The MP12 gadget replaces the simple `b = -a·r` key with a ℓ-component
// vector key where each B[i] embeds a gadget base power:
//
//   Public key:  (a, B[0], B[1], ..., B[ℓ-1])
//                where B[i] = -a·R[i] + base^i (mod q)
//   Secret key:  R[0], ..., R[ℓ-1] (short ternary trapdoor components)
//
// Verification:  a·E1 + Σ B[i]·E2[i] = H(m)  (mod q)
//
// The signing algorithm uses a Gaussian perturbation to statistically hide
// the trapdoor, then decomposes the (target - perturbation) into base-b digits:
//
//   1. t = H(m)
//   2. p ← D_σ  (zero-centered Gaussian perturbation)
//   3. t' = t - a·p  (coset shift — prevents trapdoor leakage)
//   4. E2 = GadgetDecompose(t', base)  — ℓ digit polynomials
//   5. E1 = p + Σ R[i]·E2[i]
//
// NOTE: The public key contains ℓ+1 polynomials. For Falcon-512 with base=2,
// ℓ = 14, making the public key ~13 KB. This is acceptable for enterprise/
// government deployments (10-100 counterparties) but impractical for
// consumer-scale use. The compact production target is the Falcon NTRU
// trapdoor (~700 B keys, ~700 B signatures), which is tracked as separate work.

package ring

import (
	"crypto/rand"
	"encoding/binary"
	"io"

	"golang.org/x/crypto/sha3"
)

// GPVGadgetPublicKey is an ℓ-component MP12 gadget public key.
// Verification: A·E1 + Σ B[i]·E2[i] = H(m) (mod q).
type GPVGadgetPublicKey struct {
	A *Poly       // uniform ring element (NTT form)
	B []*Poly     // B[i] = -A·R[i] + base^i (mod q), NTT form
	P GPVParams
}

// GPVGadgetSecretKey is the ℓ-component trapdoor for the gadget public key.
type GPVGadgetSecretKey struct {
	R  []*Poly           // ℓ short ternary polynomials
	PK *GPVGadgetPublicKey
}

// GPVGadgetSignature is a gadget-based lattice signature: short (E1, E2)
// satisfying A·E1 + Σ B[i]·E2[i] = H(m).
type GPVGadgetSignature struct {
	E1 *Poly      // short response polynomial
	E2 []*Poly    // ℓ gadget-digit polynomials, each with short coefficients
}

// GPVGadgetKeyGen generates an MP12 gadget key pair.
func GPVGadgetKeyGen(gp GPVParams) (*GPVGadgetPublicKey, *GPVGadgetSecretKey) {
	return GPVGadgetKeyGenFrom(gp, rand.Reader)
}

// GPVGadgetKeyGenFrom generates a key pair from the given randomness source.
func GPVGadgetKeyGenFrom(gp GPVParams, rng io.Reader) (*GPVGadgetPublicKey, *GPVGadgetSecretKey) {
	p := gp.Ring
	ℓ := gp.GadgetLevels
	base := gp.GadgetBase

	a := UniformPolyFrom(p, rng)
	NTT(a)

	r := make([]*Poly, ℓ)
	b := make([]*Poly, ℓ)

	for i := 0; i < ℓ; i++ {
		r[i] = TernaryPolyFrom(p, rng)

		rNTT := r[i].Clone()
		NTT(rNTT)
		ar := MulPointwise(a, rNTT)
		INTT(ar)

		b[i] = Neg(ar)
		gadget := uint32(1)
		for k := 0; k < i; k++ {
			gadget = mulMod(gadget, base, p.Q)
		}
		// base^i is a constant polynomial: only coefficient 0 is non-zero.
		b[i].Coeffs[0] = addMod(b[i].Coeffs[0], gadget, p.Q)
		NTT(b[i])
	}

	pk := &GPVGadgetPublicKey{A: a, B: b, P: gp}
	sk := &GPVGadgetSecretKey{R: r, PK: pk}
	return pk, sk
}

// GadgetDecompose splits a polynomial into ℓ digit polynomials.
// Each coefficient v is decomposed as v = Σ_{i=0}^{ℓ-1} digit_i · base^i,
// where each digit_i ∈ [0, base). Uses unsigned decomposition.
func GadgetDecompose(a *Poly, base uint32) []*Poly {
	p := a.params
	q := p.Q

	var levels int
	for tmp := q; tmp > 0; tmp /= base {
		levels++
	}

	digits := make([]*Poly, levels)
	for i := range digits {
		digits[i] = New(p)
	}

	for ci, coeff := range a.Coeffs {
		val := coeff
		for li := 0; li < levels; li++ {
			digits[li].Coeffs[ci] = val % base
			val /= base
		}
	}
	return digits
}

// hashMessageToTarget hashes a message to a sparse polynomial in R_q.
// The target has τ=40 non-zero coefficients, each ±1.
// Domain-separated from hashToChallenge (no commitment polynomial).
func hashMessageToTarget(p Params, message []byte) *Poly {
	h := sha3.NewShake256()
	h.Write([]byte("gpv-target-v1"))
	h.Write(message)

	tau := 40
	if tau > p.N {
		tau = p.N / 2
	}

	c := New(p)
	var buf [2]byte
	positions := make([]int, p.N)
	for i := range positions {
		positions[i] = i
	}

	for i := range tau {
		h.Read(buf[:])
		j := i + int(binary.LittleEndian.Uint16(buf[:]))%(p.N-i)
		positions[i], positions[j] = positions[j], positions[i]

		h.Read(buf[:1])
		if buf[0]&1 == 0 {
			c.Coeffs[positions[i]] = 1
		} else {
			c.Coeffs[positions[i]] = p.Q - 1
		}
	}
	return c
}

// GPVGadgetSign signs via MP12 hash-then-sign with perturbation.
//
// Algorithm:
//  1. t = H(m) — sparse polynomial via hashMessageToTarget
//  2. p ← D_σ (zero-centered Gaussian perturbation)
//  3. t' = t - a·p (coset shift — critical for trapdoor hiding)
//  4. E2 = GadgetDecompose(t', base) — ℓ digit polys, short coefficients
//  5. E1 = p + Σ R[i]·E2[i]
//
// (E1, E2) satisfy A·E1 + Σ B[i]·E2[i] = t.
func GPVGadgetSign(sk *GPVGadgetSecretKey, message []byte) *GPVGadgetSignature {
	return GPVGadgetSignFrom(sk, message, rand.Reader)
}

// GPVGadgetSignFrom signs with the given randomness source.
func GPVGadgetSignFrom(sk *GPVGadgetSecretKey, message []byte, rng io.Reader) *GPVGadgetSignature {
	pk := sk.PK
	gp := pk.P
	p := gp.Ring
	base := gp.GadgetBase
	sigma := gp.Sigma
	gs := NewGaussianSamplerFrom(sigma, rng)

	t := hashMessageToTarget(p, message)

	perturb := gs.SamplePoly(p)
	perturbNTT := perturb.Clone()
	NTT(perturbNTT)

	aNTT := pk.A
	ap := MulPointwise(aNTT, perturbNTT)
	INTT(ap)
	tPrime := Sub(t, ap)

	e2 := GadgetDecompose(tPrime, base)

	e1 := perturb.Clone()
	for i := range e2 {
		rNTT := sk.R[i].Clone()
		NTT(rNTT)
		e2NTT := e2[i].Clone()
		NTT(e2NTT)
		prod := MulPointwise(rNTT, e2NTT)
		INTT(prod)
		e1 = Add(e1, prod)
	}

	return &GPVGadgetSignature{
		E1: e1,
		E2: e2,
	}
}

// GPVGadgetVerify verifies a gadget-based GPV signature.
//
// Checks:
//  1. A·E1 + Σ B[i]·E2[i] = H(m) (mod q) — algebraic correctness
//  2. ||E2[i]||_∞ < base — each digit is a valid gadget digit
//  3. ||E1||_∞ ≤ bound — response is short enough for security
func GPVGadgetVerify(pk *GPVGadgetPublicKey, message []byte, sig *GPVGadgetSignature) bool {
	gp := pk.P
	p := gp.Ring
	base := gp.GadgetBase
	sigma := gp.Sigma

	e1 := sig.E1
	e2 := sig.E2

	if len(e2) != len(pk.B) {
		return false
	}

	for i := range e2 {
		for _, c := range e2[i].Coeffs {
			if c >= base {
				return false
			}
		}
	}

	ℓ := len(e2)
	tail := uint32(13 * sigma)
	conv := uint32(float64(p.N) * float64(base-1) / 3.0)
	e1Bound := tail + uint32(ℓ)*conv
	half := p.Q / 2
	if e1Bound > half {
		e1Bound = half
	}
	if Norm(e1) > e1Bound {
		return false
	}

	t := hashMessageToTarget(p, message)

	// Compute via coefficient-form Mul to avoid NTT pipeline issues.
	aCoeff := pk.A.Clone()
	INTT(aCoeff)
	ae1 := Mul(aCoeff, e1)

	lhs := ae1
	for i := range e2 {
		bCoeff := pk.B[i].Clone()
		INTT(bCoeff)
		term := Mul(bCoeff, e2[i])
		lhs = Add(lhs, term)
	}

	return Equal(lhs, t)
}