// GPV lattice signatures with Lyubashevsky Fiat-Shamir with Aborts.
package ring

import (
	"crypto/rand"
	"io"

	"crypto/sha3"
)

type GPVParams struct {
	Ring         Params
	Sigma        float64
	GadgetBase   uint32
	GadgetLevels int32
}

type GPVPublicKey struct {
	A *Poly
	B *Poly
	P GPVParams
}

type GPVSecretKey struct {
	R  *Poly
	PK *GPVPublicKey
}

type GPVSignature struct {
	E1 *Poly
	E2 *Poly
}

func DefaultGPVParams() (gp GPVParams) {
	p := Falcon512()
	base := uint32(2)
	levels := int32(0)
	for v := p.Q; v > 0; v /= base {
		levels++
	}
	return GPVParams{
		Ring:         p,
		Sigma:        RingGaussianSigma(p.N),
		GadgetBase:   base,
		GadgetLevels: levels,
	}
}

func SmallGPVParams() (gp GPVParams) {
	p := Params{
		N:           64,
		Q:           257,
		RootOfUnity: 9,
		MontR:       1 << 16,
		QInv:        qinv(257, 16),
	}
	base := uint32(2)
	levels := int32(0)
	for v := p.Q; v > 0; v /= base {
		levels++
	}
	return GPVParams{
		Ring:         p,
		Sigma:        RingGaussianSigma(p.N),
		GadgetBase:   base,
		GadgetLevels: levels,
	}
}

func GPVKeyGen(gp GPVParams) (pk *GPVPublicKey, sk *GPVSecretKey) {
	return GPVKeyGenFrom(gp, rand.Reader)
}

func GPVKeyGenFrom(gp GPVParams, rng io.Reader) (pk *GPVPublicKey, sk *GPVSecretKey) {
	p := gp.Ring

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

	r := TernaryPolyFrom(p, rng)

	rNTT := r.Clone()
	NTT(rNTT)

	ar := MulPointwise(a, rNTT)
	b := Neg(ar)
	INTT(b)
	bNTT := b.Clone()
	NTT(bNTT)

	pk = &GPVPublicKey{A: a, B: bNTT, P: gp}
	sk = &GPVSecretKey{R: r, PK: pk}
	return pk, sk
}

func GPVSign(sk *GPVSecretKey, message []byte) (sig *GPVSignature) {
	return GPVSignFrom(sk, message, rand.Reader)
}

func GPVSignFrom(sk *GPVSecretKey, message []byte, rng io.Reader) (sig *GPVSignature) {
	p := sk.PK.P.Ring
	sigma := sk.PK.P.Sigma
	gs := NewGaussianSamplerFrom(sigma, rng)

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

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

		w := MulPointwise(sk.PK.A, yNTT)
		INTT(w)

		c := hashToChallenge(p, w, message)
		cNTT := c.Clone()
		NTT(cNTT)

		rc := MulPointwise(rNTT, cNTT)
		INTT(rc)
		z := Add(y, rc)

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

		return &GPVSignature{
			E1: z,
			E2: c,
		}
	}
}

func GPVVerify(pk *GPVPublicKey, message []byte, sig *GPVSignature) (ok bool) {
	p := pk.P.Ring
	sigma := pk.P.Sigma

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

	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)
	wNTT := Add(az, bc)
	w := wNTT.Clone()
	INTT(w)

	cPrime := hashToChallenge(p, w, message)
	return Equal(c, cPrime)
}

func hashToChallenge(p Params, w *Poly, message []byte) (c *Poly) {
	h := sha3.NewSHAKE256()
	h.Write([]byte("gpv-challenge-v1"))

	wBytes := Serialize(w)
	h.Write(wBytes)
	h.Write(message)

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

	c = New(p)

	var buf [2]byte
	positions := []int32{:p.N}
	for i := int32(0); i < p.N; i++ {
		positions[i] = i
	}

	for i := int32(0); i < tau; i++ {
		h.Read(buf[:])
		j := i + int32(leU16(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
}