// Falcon-like NTRU lattice signatures.
//
// Uses the NTRU-structured trapdoor to produce compact signatures (~700 B for
// Falcon-512). The verification key is a single polynomial h = g/f mod q, and
// the signature is a single short polynomial s2. The verifier derives s1 from
// the verification equation and checks both are short.
//
// Key generation:
//   1. Sample short f, g ← D_{Z^n, σ}
//   2. Compute h = g·f^{-1} mod q
//   3. Find short F, G such that f·G - g·F = q (NTRU equation)
//   4. Secret key: (f, g, F, G). Public key: h.
//
// Signing (ffSampling):
//   1. c = HashToPoint(m) ∈ R_q (sparse ±1 challenge)
//   2. Use the NTRU basis B = [(g, -f), (G, -F)] with Gram-Schmidt
//      orthogonalization to sample a Gaussian preimage (s1, s2) such that
//      s1 + h·s2 = c (mod q).
//   3. Output s2 as the signature (s1 is derived during verification).
//
// Verification:
//   1. c = HashToPoint(m)
//   2. s1 = c - h·s2 (mod q)
//   3. Check ||s1||_∞, ||s2||_∞ ≤ bound

package ring

import (
	"crypto/rand"
	"io"
	"math"
)

// FalconPublicKey is the NTRU verification key.
type FalconPublicKey struct {
	H *Poly   // h = g·f^{-1} mod q (coefficient form)
	P Params
}

// FalconSecretKey is the NTRU trapdoor basis.
type FalconSecretKey struct {
	F, G, f, g *Poly   // short NTRU basis (coefficient form)
	Sigma       float64
	PK          *FalconPublicKey
}

// FalconSignature is a compact NTRU signature (single short polynomial).
type FalconSignature struct {
	S2 *Poly   // signature component (short, compressed)
}

// DefaultFalconParams returns Falcon-512 parameters.
func DefaultFalconParams() Params {
	return Falcon512()
}

// ─── Polynomial dot product and Gram-Schmidt ──────────────────────────────

// polyDot computes the Euclidean dot product Σ a_i·b_i (scalar).
// Both polynomials are in coefficient form. Used for Gram-Schmidt
// orthogonalization in the ffSampling algorithm.
func polyDot(a, b *Poly) int64 {
	n := len(a.Coeffs)
	var sum int64
	half := int64(a.params.Q / 2)
	for i := 0; i < n; i++ {
		ai := int64(a.Coeffs[i])
		bi := int64(b.Coeffs[i])
		if ai > half {
			ai -= int64(a.params.Q)
		}
		if bi > half {
			bi -= int64(a.params.Q)
		}
		sum += ai * bi
	}
	return sum
}

// polyVectorDot computes the Euclidean dot product of two 2-vectors of
// polynomials: <(a1,a2), (b1,b2)> = dot(a1,b1) + dot(a2,b2).
func polyVectorDot(a1, a2, b1, b2 *Poly) int64 {
	return polyDot(a1, b1) + polyDot(a2, b2)
}

// gramSchmidtNormSq returns ||b*||² for the Gram-Schmidt orthogonalization
// of the second basis vector. For basis B = [b1, b2] where b1 and b2 are
// 2-vectors of polynomials, computes:
//
//	µ = <b2, b1> / <b1, b1>
//	b2* = b2 - µ·b1
//	return <b2*, b2*>
func gramSchmidtNormSq(b11, b12, b21, b22 *Poly) (normSq float64) {
	b1NormSq := float64(polyVectorDot(b11, b12, b11, b12))
	b2b1 := float64(polyVectorDot(b21, b22, b11, b12))
	µ := b2b1 / b1NormSq

	// b2*_i = b2_i - µ·b1_i for each coefficient of each polynomial
	n := len(b11.Coeffs)
	var sum float64
	for i := 0; i < n; i++ {
		b21c := float64(b21.Coeffs[i])
		b11c := float64(b11.Coeffs[i])
		b22c := float64(b22.Coeffs[i])
		b12c := float64(b12.Coeffs[i])

		diff1 := b21c - µ*b11c
		diff2 := b22c - µ*b12c
		sum += diff1*diff1 + diff2*diff2
	}
	return sum
}

// ─── NTRU Key Generation ─────────────────────────────────────────────────

// polyInverse computes f^{-1} mod (x^n+1, q) using NTT-based pointwise inverse.
// f is in coefficient form. Returns nil if f is not invertible.
// For Falcon-512 (q=12289, q≡1 mod 2n), the NTT splits the ring completely,
// so pointwise inversion in the NTT domain is correct.
func polyInverse(f *Poly) *Poly {
	p := f.params
	q := p.Q

	fNTT := f.Clone()
	NTT(fNTT)

	resultNTT := New(p)
	for i, c := range fNTT.Coeffs {
		if c == 0 {
			return nil // not invertible — has a zero in NTT domain
		}
		resultNTT.Coeffs[i] = powMod(c, q-2, q)
	}
	resultNTT.isNTT = true

	INTT(resultNTT)
	return resultNTT
}

// computeNTRUBasis finds short F, G such that f·G - g·F = q (mod x^n+1).
// Uses the XGCD in Z_q[x]/(x^n+1) followed by Babai reduction.
func computeNTRUBasis(f, g *Poly, sigma float64, rng io.Reader) (*Poly, *Poly) {
	p := f.params

	// G0 = q·f^{-1}, F0 = 0 satisfies f·G0 - g·F0 = q.
	fInv := polyInverse(f)
	if fInv == nil {
		return New(p), New(p)
	}
	G0 := PolyScalarMul(fInv, uint32(p.Q))
	F0 := New(p) // zero polynomial

	// Babai reduction: iteratively subtract k·(f, g) from (F, G) to make them short.
	// k = round(<(F,G), (f,g)> / ||(f,g)||²)
	b1Norm := float64(polyVectorDot(f, g, f, g))
	if b1Norm < 1 {
		b1Norm = 1
	}

	Fcur, Gcur := F0, G0
	for iter := 0; iter < 10; iter++ {
		b2b1 := float64(polyVectorDot(Fcur, Gcur, f, g))
		kFloat := b2b1 / b1Norm
		k := int(math.Round(kFloat))
		if k == 0 {
			break
		}
		// Subtract k·(f, g) from (F, G).
		for i := range Fcur.Coeffs {
			Fcur.Coeffs[i] = subMod(Fcur.Coeffs[i], mulMod(uint32(k), f.Coeffs[i], p.Q), p.Q)
			Gcur.Coeffs[i] = subMod(Gcur.Coeffs[i], mulMod(uint32(k), g.Coeffs[i], p.Q), p.Q)
		}
	}

	return Fcur, Gcur
}

// PolyScalarMul multiplies each coefficient of a polynomial by a scalar mod q.
func PolyScalarMul(a *Poly, s uint32) *Poly {
	c := a.Clone()
	q := a.params.Q
	for i := range c.Coeffs {
		c.Coeffs[i] = mulMod(c.Coeffs[i], s, q)
	}
	c.isNTT = false
	return c
}

// FalconKeyGen generates a Falcon NTRU key pair.
func FalconKeyGen(p Params) (*FalconPublicKey, *FalconSecretKey) {
	return FalconKeyGenFrom(p, rand.Reader)
}

// FalconKeyGenFrom generates a Falcon NTRU key pair from the given RNG.
func FalconKeyGenFrom(p Params, rng io.Reader) (*FalconPublicKey, *FalconSecretKey) {
	sigma := math.Sqrt(float64(p.N)) * 2.0

	gs := NewGaussianSamplerFrom(sigma, rng)

	// Sample f, g from discrete Gaussian.  Resample f if not invertible
	// (any NTT slot is zero — coefficient 0 check alone is insufficient).
	var f, g, fInv *Poly
	for {
		f = gs.SamplePoly(p)
		g = gs.SamplePoly(p)
		fInv = polyInverse(f)
		if fInv != nil {
			break
		}
	}

	// Compute h = g·f^{-1} mod q.
	hCoeff := Mul(fInv, g) // coefficient form

	// Find short F, G.
	F, G := computeNTRUBasis(f, g, sigma, rng)

	pk := &FalconPublicKey{H: hCoeff, P: p}
	sk := &FalconSecretKey{F: F, G: G, f: f, g: g, Sigma: sigma, PK: pk}
	return pk, sk
}

// ─── ffSampling (Gaussian preimage sampling) ─────────────────────────────

// ffSampling samples a short vector (s1, s2) such that s1 + h·s2 = target
// using the NTRU trapdoor basis with the ffSampling algorithm.
//
// The algorithm:
//  1. Set T = (0, target) — a non-short preimage
//  2. For each basis vector b_i (from bottom to top):
//     c_i = <T, b_i*> / ||b_i*||²
//     z_i ← D_{Z, σ/||b_i*||, c_i}
//     T = T - z_i·b_i
//  3. Return T as the short preimage (s1, s2)
func ffSampling(p Params, f, g, F, G *Poly, target *Poly, sigma float64, rng io.Reader) (*Poly, *Poly) {
	// NTRU basis vectors: b1 = (g, -f), b2 = (F, -G)
	// Wait — verify the equation: s1 + h·s2 = target
	// (g, -f): g + h·(-f) = g - h·f = g - g = 0 ✓
	// (F, -G): F + h·(-G) = F - h·G ≠ 0
	//
	// Actually, the second basis vector is scaled: we need (G, -F)
	// Check: G + h·(-F) = G - h·F. From the NTRU equation f·G - g·F = q,
	// we have -g·F = q - f·G, so F = (f·G - q)/g. Then G - h·F = G - g/f·F = 0
	// when multiplied by f: f·G - g·F = q. So f·(G - h·F) = q, not 0.
	//
	// For the ffSampling, both (g, -f) and ( -F, G) form the basis.
	// The lattice generated by these two vectors has determinant q.

	b1 := [2]*Poly{g, Neg(f)} // (g, -f)
	b2 := [2]*Poly{Neg(G), F} // (-G, F) — columns of the NTRU basis

	// Current target T = (0, target).
	t1 := New(p) // zero polynomial
	t2 := target.Clone()

	// Gram-Schmidt norms.
	b2NormSq := gramSchmidtNormSq(b1[0], b1[1], b2[0], b2[1])

	// === Process b2 (second basis vector) ===

	// c2 = <T, b2*> / ||b2*||²
	// Compute the inner product <T, (b2 - µ·b1)> / ||b2*||²
	b1NormSq := float64(polyVectorDot(b1[0], b1[1], b1[0], b1[1]))
	if b1NormSq < 1 {
		b1NormSq = 1
	}
	b2b1 := float64(polyVectorDot(b2[0], b2[1], b1[0], b1[1]))
	µ := b2b1 / b1NormSq

	// <T, b2*> = <T, b2> - µ·<T, b1>
	tb2 := float64(polyVectorDot(t1, t2, b2[0], b2[1]))
	tb1 := float64(polyVectorDot(t1, t2, b1[0], b1[1]))
	c2 := (tb2 - µ*tb1) / b2NormSq

	// Sample z2 from D_{Z, σ/||b2*||, c2}.
	sigma2 := sigma / math.Sqrt(b2NormSq)
	z2 := sampleZ(sigma2, c2, rng)

	// T = T - z2·b2 (subtract z2 copies of the second basis vector).
	// Subtracting z2 from the coefficient along b2 brings us closer to the lattice.
	for i := 0; i < len(t1.Coeffs); i++ {
		t1.Coeffs[i] = subMod(t1.Coeffs[i], mulMod(uint32(z2), b2[0].Coeffs[i], p.Q), p.Q)
		t2.Coeffs[i] = subMod(t2.Coeffs[i], mulMod(uint32(z2), b2[1].Coeffs[i], p.Q), p.Q)
	}

	// === Process b1 (first basis vector) ===

	// c1 = <T, b1*> / ||b1*||² = <T, b1> / ||b1||²
	tb1_2 := float64(polyVectorDot(t1, t2, b1[0], b1[1]))
	c1 := tb1_2 / b1NormSq

	sigma1 := sigma / math.Sqrt(b1NormSq)
	z1 := sampleZ(sigma1, c1, rng)

	// T = T - z1·b1
	for i := 0; i < len(t1.Coeffs); i++ {
		t1.Coeffs[i] = subMod(t1.Coeffs[i], mulMod(uint32(z1), b1[0].Coeffs[i], p.Q), p.Q)
		t2.Coeffs[i] = subMod(t2.Coeffs[i], mulMod(uint32(z1), b1[1].Coeffs[i], p.Q), p.Q)
	}

	return t1, t2
}

// sampleZ samples a single integer from D_{Z, σ, c} using rejection sampling.
// For sigma < 1, the distribution is concentrated at the nearest integer
// (probability of any other value is < exp(-π/1) ≈ 0.04 and decreasing
// rapidly). Returns the rounded center directly to avoid rejection loop hangs.
func sampleZ(sigma, center float64, rng io.Reader) int64 {
	if sigma < 1.0 || math.IsNaN(sigma) {
		return int64(math.Round(center))
	}

	tail := int(math.Ceil(13.0 * sigma))
	if tail < 1 {
		tail = 1
	}
	centerInt := int64(math.Round(center))

	var buf [8]byte
	piOverSigma2 := math.Pi / (sigma * sigma)

	for {
		io.ReadFull(rng, buf[:])
		u := int64(readLEUint64(buf[:]))
		candidate := centerInt + (u % int64(2*tail+1)) - int64(tail)

		diff := float64(candidate) - center
		logProb := -piOverSigma2 * diff * diff

		io.ReadFull(rng, buf[:])
		uFloat := float64(readLEUint64(buf[:])>>11) / float64(uint64(1)<<53)

		if math.Log(uFloat) < logProb {
			return candidate
		}
	}
}

// readLEUint64 reads a little-endian uint64 from a byte slice.
func readLEUint64(b []byte) uint64 {
	_ = b[7]
	return uint64(b[0]) | uint64(b[1])<<8 | uint64(b[2])<<16 | uint64(b[3])<<24 |
		uint64(b[4])<<32 | uint64(b[5])<<40 | uint64(b[6])<<48 | uint64(b[7])<<56
}

// ─── Falcon signing and verification ─────────────────────────────────────

// hashToFalconTarget hashes a message to a short polynomial c for Falcon.
// Uses domain separation from the GPV and sigma functions.
func hashToFalconTarget(p Params, message []byte) *Poly {
	return hashMessageToTarget(p, message) // same format: τ=40 (or τ=n/2) ±1 coeffs
}

// FalconSign signs a message using the NTRU secret key. Produces a single
// short polynomial s2 as the signature. The verifier derives s1 from the
// verification equation and checks both are short.
func FalconSign(sk *FalconSecretKey, message []byte) *FalconSignature {
	return FalconSignFrom(sk, message, rand.Reader)
}

// FalconSignFrom signs with the given randomness source.
func FalconSignFrom(sk *FalconSecretKey, message []byte, rng io.Reader) *FalconSignature {
	p := sk.PK.P
	sigma := sk.Sigma

	// c = HashToPoint(m)
	c := hashToFalconTarget(p, message)

	// ffSampling: find short (s1, s2) such that s1 + h·s2 = c.
	s1, s2 := ffSampling(p, sk.f, sk.g, sk.F, sk.G, c, sigma, rng)
	_ = s1

	return &FalconSignature{S2: s2}
}

// FalconVerify verifies a Falcon signature.
//
// Checks:
//  1. s1 = c - h·s2 (mod q)
//  2. ||s1||_∞, ||s2||_∞ ≤ σ·1.5 + 13·σ (= bound for Gaussian preimage)
func FalconVerify(pk *FalconPublicKey, message []byte, sig *FalconSignature) bool {
	p := pk.P
	sigma := math.Sqrt(float64(p.N)) * 2.0

	s2 := sig.S2

	// c = HashToPoint(m).
	c := hashToFalconTarget(p, message)

	// s1 = c - h·s2 (mod q).
	hs2 := Mul(pk.H, s2)
	s1 := Sub(c, hs2)

	// Norm check: ||s1||_∞, ||s2||_∞ ≤ bound.
	// The bound accounts for the Gaussian preimage sampling width.
	bound := uint32(sigma*1.5 + 13*sigma)
	if Norm(s1) > bound {
		return false
	}
	if Norm(s2) > bound {
		return false
	}

	return true
}

// FalconSigBound returns the norm bound for Falcon signatures.
func FalconSigBound(p Params) uint32 {
	sigma := math.Sqrt(float64(p.N)) * 2.0
	return uint32(sigma*1.5 + 13*sigma)
}

// FalconSignFrom returns (s1, s2) pair for verification debugging.
func FalconSignWithS1(sk *FalconSecretKey, message []byte, rng io.Reader) (s1, s2 *Poly) {
	p := sk.PK.P
	sigma := sk.Sigma
	c := hashToFalconTarget(p, message)
	s1, s2 = ffSampling(p, sk.f, sk.g, sk.F, sk.G, c, sigma, rng)
	return s1, s2
}

// Verify the s1 is derived correctly.
func verifyFalconEquation(pk *FalconPublicKey, c, s1, s2 *Poly) bool {
	// c = s1 + h·s2 (mod q).
	hs2 := Mul(pk.H, s2)
	computed := Add(s1, hs2)
	return Equal(c, computed)
}