// Sigma protocol layer for lattice-based zero-knowledge proofs.
//
// Builds on the existing GPV signature infrastructure to provide:
//
//   - Phase 2: Interactive sigma protocol for knowledge of a GPV secret key
//   - Phase 3: Non-interactive (Fiat-Shamir) proof of key knowledge
//   - Phase 4: Proof of knowledge of a valid GPV signature on a committed message
//              (uses ring-inversion approach — requires large modulus)
//
// All proofs use the verification equation a·e1 + b·e2 = H(m) from the existing
// GPV key structure (single polynomial B = -A·R). The gadget key
// (GPVGadgetPublicKey with vector B) is NOT supported — its Σ B[i]·E2[i]
// verification is algebraically different and would require a separate protocol.
//
// Security: all proofs are HVZK (honest-verifier zero knowledge) via the standard
// sigma protocol paradigm. The Fiat-Shamir transform gives non-interactive proofs
// in the random oracle model.

package ring

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

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

// ─── Phase 2: Interactive sigma protocol ─────────────────────────────────────

// Commit generates the commitment w = a·y for random Gaussian y.
// Returns (w, y). The prover sends w to the verifier and keeps y secret.
func Commit(pk *GPVPublicKey, rng io.Reader) (*Poly, *Poly) {
	p := pk.P.Ring
	sigma := pk.P.Sigma
	gs := NewGaussianSamplerFrom(sigma, rng)

	y := gs.SamplePoly(p)
	yNTT := y.Clone()
	NTT(yNTT)

	w := MulPointwise(pk.A, yNTT)
	INTT(w)

	return w, y
}

// Challenge generates a random short challenge polynomial.
// The challenge has τ=40 non-zero coefficients, each ±1.
func Challenge(p Params, rng io.Reader) *Poly {
	c := New(p)
	tau := 40
	if tau > p.N {
		tau = p.N / 2
	}

	positions := make([]int, p.N)
	for i := range positions {
		positions[i] = i
	}

	var buf [2]byte
	for i := range tau {
		io.ReadFull(rng, buf[:])
		j := i + int(binary.LittleEndian.Uint16(buf[:]))%(p.N-i)
		positions[i], positions[j] = positions[j], positions[i]

		io.ReadFull(rng, buf[:1])
		if buf[0]&1 == 0 {
			c.Coeffs[positions[i]] = 1
		} else {
			c.Coeffs[positions[i]] = p.Q - 1
		}
	}
	return c
}

// Respond computes the response z = y + r·c with rejection sampling.
// Returns (z, ok). If ok is false, the prover must retry with a fresh commitment.
func Respond(sk *GPVSecretKey, y, c *Poly) (*Poly, bool) {
	sigma := sk.PK.P.Sigma

	cNTT := c.Clone()
	NTT(cNTT)

	rNTT := sk.R.Clone()
	NTT(rNTT)

	rc := MulPointwise(rNTT, cNTT)
	INTT(rc)

	z := Add(y, rc)

	zNorm := Norm(z)
	bound := uint32(sigma * 1.5)
	if zNorm > bound {
		return z, false
	}
	return z, true
}

// VerifyInteractive checks a sigma protocol proof:
//
//	1. ||z||_∞ ≤ bound
//	2. a·z + b·c = w  (mod q)
func VerifyInteractive(pk *GPVPublicKey, w, c, z *Poly) bool {
	sigma := pk.P.Sigma

	zNorm := Norm(z)
	bound := uint32(sigma * 1.5)
	if zNorm > bound {
		return false
	}

	zNTT := z.Clone()
	NTT(zNTT)
	cNTT := c.Clone()
	NTT(cNTT)

	az := MulPointwise(pk.A, zNTT)
	bc := MulPointwise(pk.B, cNTT)
	wComputed := Add(az, bc)
	INTT(wComputed)

	return Equal(w, wComputed)
}

// ─── Phase 3: Non-interactive (Fiat-Shamir) proof of key knowledge ────────────

// SigmaProof is a non-interactive proof of GPV secret key knowledge.
type SigmaProof struct {
	Z *Poly // response (short, rejection-sampled)
	C *Poly // challenge (sparse, τ=40 ±1)
}

// hashChallenge computes c = H(transcript) using SHAKE256.
// Domain separated from hashToChallenge and hashMessageToTarget.
func hashChallenge(p Params, components [][]byte) *Poly {
	h := sha3.NewShake256()
	h.Write([]byte("gpv-sigma-proof-v1"))
	for _, comp := range components {
		h.Write(comp)
	}

	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
}

// Prove generates a non-interactive proof of GPV secret key knowledge.
// Uses Fiat-Shamir: c = H(w, context).
func Prove(sk *GPVSecretKey, context []byte) *SigmaProof {
	return ProveFrom(sk, context, rand.Reader)
}

// ProveFrom generates a proof with the given randomness source.
func ProveFrom(sk *GPVSecretKey, context []byte, rng io.Reader) *SigmaProof {
	p := sk.PK.P.Ring

	for {
		w, y := Commit(sk.PK, rng)

		wSerialized := Serialize(w)
		var components [][]byte
		if len(context) > 0 {
			components = [][]byte{wSerialized, context}
		} else {
			components = [][]byte{wSerialized}
		}
		c := hashChallenge(p, components)

		z, ok := Respond(sk, y, c)
		if ok {
			return &SigmaProof{Z: z, C: c}
		}
	}
}

// VerifySigma checks a non-interactive proof of key knowledge.
func VerifySigma(pk *GPVPublicKey, context []byte, proof *SigmaProof) bool {
	p := pk.P.Ring
	z := proof.Z
	c := proof.C

	zNTT := z.Clone()
	NTT(zNTT)
	cNTT := c.Clone()
	NTT(cNTT)

	az := MulPointwise(pk.A, zNTT)
	bc := MulPointwise(pk.B, cNTT)
	w := Add(az, bc)
	INTT(w)

	wSerialized := Serialize(w)
	var components [][]byte
	if len(context) > 0 {
		components = [][]byte{wSerialized, context}
	} else {
		components = [][]byte{wSerialized}
	}
	cExpected := hashChallenge(p, components)

	if !Equal(c, cExpected) {
		return false
	}

	sigma := pk.P.Sigma
	bound := uint32(sigma * 1.5)
	if Norm(z) > bound {
		return false
	}

	return true
}

// MarshalSigmaProof serializes a SigmaProof to bytes.
// Challenge c is compressed as (position, sign) pairs.
// Response z is compressed via SerializeBounded at 12 bits/coeff.
func MarshalSigmaProof(proof *SigmaProof) []byte {
	q := proof.C.params.Q
	half := q / 2

	cPairs := make([]byte, 0, 80)
	for i, coeff := range proof.C.Coeffs {
		if coeff != 0 {
			absV := coeff
			sign := byte(0)
			if absV > half {
				sign = 1
			}
			cPairs = append(cPairs, byte(i&0xFF), byte(i>>8), sign)
		}
	}

	cData := cPairs
	zData := SerializeBounded(proof.Z, 12, true)

	out := make([]byte, 8+len(cData)+len(zData))
	binary.LittleEndian.PutUint32(out[0:4], uint32(len(cData)))
	copy(out[4:], cData)
	binary.LittleEndian.PutUint32(out[4+len(cData):4+len(cData)+4], uint32(len(zData)))
	copy(out[8+len(cData):], zData)
	return out
}

// UnmarshalSigmaProof deserializes a SigmaProof from bytes.
func UnmarshalSigmaProof(p Params, data []byte) *SigmaProof {
	if len(data) < 8 {
		return nil
	}
	cLen := int(binary.LittleEndian.Uint32(data[0:4]))
	if len(data) < 8+cLen {
		return nil
	}

	cData := data[4 : 4+cLen]
	zLen := int(binary.LittleEndian.Uint32(data[4+cLen : 4+cLen+4]))
	if len(data) < 8+cLen+zLen {
		return nil
	}
	zData := data[8+cLen : 8+cLen+zLen]

	c := New(p)
	for i := 0; i+2 < len(cData); i += 3 {
		pos := uint16(cData[i]) | uint16(cData[i+1])<<8
		sign := cData[i+2]
		if pos < uint16(p.N) {
			if sign == 0 {
				c.Coeffs[pos] = 1
			} else {
				c.Coeffs[pos] = p.Q - 1
			}
		}
	}

	z := DeserializeBounded(p, zData, 12, true)

	return &SigmaProof{Z: z, C: c}
}

// ─── Phase 4: Randomized GPV signature (signature-hiding, message-revealing) ──
//
// Randomizes the GPV signature response z by adding a Gaussian perturbation,
// and sends a commitment to the perturbation. This is coefficient-wise addition
// (not ring convolution), so no modulus wrapping occurs for any ring.
//
// Protocol:
//
//	Prover has GPV signature (z, c) valid on m:
//	  w = A·z + B·c   (GPV commitment)
//	  c = H(w, m)     (Fiat-Shamir challenge)
//
//	1. Sample r_m ← D_σ (additive Gaussian randomizer)
//	2. Compute z' = z + r_m (randomized response)
//	3. Compute C = A·r_m    (commitment to randomizer)
//	4. Proof: (z', c, C, m)
//
//	Verifier:
//	  1. w' = A·z' + B·c - C = A·(z+r_m) + B·c - A·r_m = A·z + B·c = w
//	  2. c == H(w', m)     ← original Fiat-Shamir binds proof to m
//	  3. Norm(z') ≤ bound'  ← increased bound accounts for r_m

// RandomizedSigProof is a randomized GPV signature.
// Hides (z) but reveals (c, m) to the verifier.
type RandomizedSigProof struct {
	ZPrime *Poly // randomized z' = z + r_m
	Chal   *Poly // original challenge c (binds to m via Fiat-Shamir)
	Commit *Poly // commitment C = A·r_m
}

// RandomizeGPVSignature randomizes a GPV signature to produce a hiding proof.
// The message m is required for the Fiat-Shamir challenge check.
func RandomizeGPVSignature(pk *GPVPublicKey, sig *GPVSignature, message []byte, rng io.Reader) *RandomizedSigProof {
	p := pk.P.Ring
	sigma := pk.P.Sigma
	gs := NewGaussianSamplerFrom(sigma, rng)

	z := sig.E1
	c := sig.E2

	rM := gs.SamplePoly(p)
	zPrime := Add(z, rM)

	rMNTT := rM.Clone()
	NTT(rMNTT)
	CNTT := MulPointwise(pk.A, rMNTT)
	INTT(CNTT)

	return &RandomizedSigProof{
		ZPrime: zPrime,
		Chal:   c,
		Commit: CNTT,
	}
}

// VerifyRandomizedGPVSignature verifies a randomized GPV signature.
//
// Checks:
//  1. Norm(z') ≤ σ·1.5 + 13·σ (accounts for r_m)
//  2. w' = A·z' + B·c - C
//  3. c == H(w', m) (original Fiat-Shamir, binds to m)
func VerifyRandomizedGPVSignature(pk *GPVPublicKey, message []byte, proof *RandomizedSigProof) bool {
	sigma := pk.P.Sigma
	zPrime := proof.ZPrime
	c := proof.Chal
	commit := proof.Commit

	// Norm check: ||z'|| ≤ ||z|| + ||r_m|| ≤ σ·1.5 + 13·σ.
	bound := uint32(sigma*1.5 + 13*sigma)
	if Norm(zPrime) > bound {
		return false
	}

	// w' = A·z' + B·c - C.
	zNTT := zPrime.Clone()
	NTT(zNTT)
	cNTT := c.Clone()
	NTT(cNTT)
	commitNTT := commit.Clone()
	NTT(commitNTT)

	az := MulPointwise(pk.A, zNTT)
	bc := MulPointwise(pk.B, cNTT)
	wNTT := Add(az, bc)
	wNTT = Sub(wNTT, commitNTT)
	INTT(wNTT)

	// c must equal H(w', m).
	cExpected := hashToChallenge(pk.P.Ring, wNTT, message)

	return Equal(c, cExpected)
}