package gnarlring

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

	"git.smesh.lol/gnarl-hamadryad/crypto"
)

type NTRUPublicKey struct {
	H *Poly27
}

type NTRUPrivateKey struct {
	BTop [2]*Poly27 // b_i_top for each of 2 basis vectors
	BBot [2]*Poly27 // b_i_bot
	Sigma float64
	PK    *NTRUPublicKey
}

type NTRUSignature struct {
	Salt [16]byte
	S2   *Poly27
}

func NTRUKeyGen() (*NTRUPublicKey, *NTRUPrivateKey) {
	return NTRUKeyGenFrom(rand.Reader)
}

func NTRUKeyGenFrom(rng io.Reader) (*NTRUPublicKey, *NTRUPrivateKey) {
	sigma := math.Sqrt(float64(N)) * 2.0
	gs := NewGaussSamplerFrom(sigma, rng)

	var f, g, h *Poly27
	for {
		f = gs.SamplePoly()
		g = gs.SamplePoly()
		if Inverse(f) != nil {
			h = Mul(g, Inverse(f))
			break
		}
	}

	// Build NTRU lattice, LLL-reduce, use first 2 reduced columns as basis.
	bTop0, bBot0, bTop1, bBot1 := ntruBasisFromLLL(h)

	pk := &NTRUPublicKey{H: h}
	sk := &NTRUPrivateKey{
		BTop:  [2]*Poly27{bTop0, bTop1},
		BBot:  [2]*Poly27{bBot0, bBot1},
		Sigma: sigma,
		PK:    pk,
	}
	return pk, sk
}

func ntruBasisFromLLL(h *Poly27) (*Poly27, *Poly27, *Poly27, *Poly27) {
	n := N
	q := int64(Q)
	d := 2 * n

	hMat := buildIntMatrix(h, n)

	// Build kernel lattice K = {(u,v) ∈ Z^{2n} : u + h*v ≡ 0 mod q}.
	// Basis columns:
	//   col j (0 ≤ j < n):  u = -H_mat[:,j], v = e_j
	//   col n+j (0 ≤ j < n): u = q*e_j,       v = 0
	B := make([][]int64, d)
	for i := range B {
		B[i] = make([]int64, d)
	}

	for j := 0; j < n; j++ {
		// Top half: u = -H_mat[:,j]
		for i := 0; i < n; i++ {
			B[j][i] = -hMat[i][j]
		}
		// Bottom half: v = e_j
		B[j][n+j] = 1
	}

	for j := 0; j < n; j++ {
		// Top half: u = q*e_j
		B[n+j][j] = q
		// Bottom half: v = 0 (already zero)
	}

	lllReduce(B, 0.99)

	bTop0 := NewPoly27()
	bBot0 := NewPoly27()
	bTop1 := NewPoly27()
	bBot1 := NewPoly27()
	for i := 0; i < n; i++ {
		bTop0.Coeffs[i] = modInt64(B[0][i])
		bBot0.Coeffs[i] = modInt64(B[0][n+i])
		bTop1.Coeffs[i] = modInt64(B[1][i])
		bBot1.Coeffs[i] = modInt64(B[1][n+i])
	}

	n0 := float64(NormSq(bTop0) + NormSq(bBot0))
	n1 := float64(NormSq(bTop1) + NormSq(bBot1))
	if n1 < n0 {
		bTop0, bTop1 = bTop1, bTop0
		bBot0, bBot1 = bBot1, bBot0
	}

	return bTop0, bBot0, bTop1, bBot1
}

func buildIntMatrix(a *Poly27, n int) [][]int64 {
	M := make([][]int64, n)
	for k := 0; k < n; k++ {
		M[k] = make([]int64, n)
		for j := 0; j <= k; j++ {
			M[k][j] = int64(a.Coeffs[k-j])
		}
		for j := k + 1; j < n; j++ {
			M[k][j] = -int64(a.Coeffs[k-j+n])
		}
	}
	return M
}

func modInt64(x int64) uint16 {
	for x < 0 {
		x += int64(Q)
	}
	return uint16(uint64(x) % uint64(Q))
}

func ffSampling(bTop [2]*Poly27, bBot [2]*Poly27, target *Poly27, sigma float64, rng io.Reader) (*Poly27, *Poly27) {
	tTop := NewPoly27()
	tBot := target.Clone()

	b1Top := bTop[0]
	b1Bot := bBot[0]
	b2Top := bTop[1]
	b2Bot := bBot[1]

	b1NormSq := float64(Dot(b1Top, b1Top) + Dot(b1Bot, b1Bot))
	b1b2 := float64(Dot(b2Top, b1Top) + Dot(b2Bot, b1Bot))
	mu := b1b2 / b1NormSq

	b2StarTop := Sub(b2Top, scalarMul(b1Top, mu))
	b2StarBot := Sub(b2Bot, scalarMul(b1Bot, mu))
	b2StarNormSq := float64(Dot(b2StarTop, b2StarTop) + Dot(b2StarBot, b2StarBot))

	if b1NormSq < 1 {
		b1NormSq = 1
	}
	if b2StarNormSq < 1 {
		b2StarNormSq = 1
	}

	tb2 := float64(Dot(tTop, b2Top) + Dot(tBot, b2Bot))
	tb1 := float64(Dot(tTop, b1Top) + Dot(tBot, b1Bot))
	c2 := (tb2 - mu*tb1) / b2StarNormSq
	sigma2 := sigma / math.Sqrt(b2StarNormSq)
	gs2 := NewGaussSamplerFrom(sigma2, rng)
	z2 := gs2.SampleZ(c2)

	for i := 0; i < N; i++ {
		tTop.Coeffs[i] = subMod(tTop.Coeffs[i], mulModInt(b2Top.Coeffs[i], z2))
		tBot.Coeffs[i] = subMod(tBot.Coeffs[i], mulModInt(b2Bot.Coeffs[i], z2))
	}

	tb1r := float64(Dot(tTop, b1Top) + Dot(tBot, b1Bot))
	c1 := tb1r / b1NormSq
	sigma1 := sigma / math.Sqrt(b1NormSq)
	gs1 := NewGaussSamplerFrom(sigma1, rng)
	z1 := gs1.SampleZ(c1)

	for i := 0; i < N; i++ {
		tTop.Coeffs[i] = subMod(tTop.Coeffs[i], mulModInt(b1Top.Coeffs[i], z1))
		tBot.Coeffs[i] = subMod(tBot.Coeffs[i], mulModInt(b1Bot.Coeffs[i], z1))
	}

	return tTop, tBot
}

func scalarMul(a *Poly27, s float64) *Poly27 {
	r := NewPoly27()
	si := uint16(math.Round(math.Abs(s)))
	if si == 0 {
		return r
	}
	for i, v := range a.Coeffs {
		m := crypto.Mod271(uint32(v) * uint32(si))
		if s < 0 {
			r.Coeffs[i] = subMod(0, m)
		} else {
			r.Coeffs[i] = m
		}
	}
	return r
}

func mulModInt(a uint16, z int64) uint16 {
	if z >= 0 {
		return crypto.Mod271(uint32(z) * uint32(a))
	}
	neg := crypto.Mod271(uint32(-z) * uint32(a))
	return subMod(0, neg)
}

// NTRUSignTarget signs a pre-hashed target polynomial directly.
func NTRUSignTarget(sk *NTRUPrivateKey, target *Poly27, rng io.Reader) *NTRUSignature {
	if rng == nil {
		rng = rand.Reader
	}
	var salt [16]byte
	io.ReadFull(rng, salt[:])
	_, s2 := ffSampling(sk.BTop, sk.BBot, target, sk.Sigma, rng)
	return &NTRUSignature{Salt: salt, S2: s2}
}

// NTRUVerifyTarget verifies a signature against a pre-hashed target.
func NTRUVerifyTarget(pk *NTRUPublicKey, target *Poly27, sig *NTRUSignature) bool {
	sigma := math.Sqrt(float64(N)) * 2.0
	hs2 := Mul(pk.H, sig.S2)
	s1 := Sub(target, hs2)
	tail := 13
	bound := uint16(sigma * (1.5 + float64(tail)))
	return Norm(s1) <= bound && Norm(sig.S2) <= bound
}

func NTRUSign(sk *NTRUPrivateKey, message []byte) *NTRUSignature {
	return NTRUSignFrom(sk, message, rand.Reader)
}

func NTRUSignFrom(sk *NTRUPrivateKey, message []byte, rng io.Reader) *NTRUSignature {
	if rng == nil {
		rng = rand.Reader
	}
	var salt [16]byte
	io.ReadFull(rng, salt[:])
	c := hashToTarget(salt[:], sk.PK.H, message)
	return NTRUSignTarget(sk, c, rng)
}

func NTRUVerify(pk *NTRUPublicKey, message []byte, sig *NTRUSignature) bool {
	c := hashToTarget(sig.Salt[:], pk.H, message)
	return NTRUVerifyTarget(pk, c, sig)
}

func hashToTarget(salt []byte, h *Poly27, msg []byte) *Poly27 {
	input := make([]byte, 0, len(salt)+PolyBytes+len(msg)+32)
	input = append(input, []byte("gnarl-ntru-v1")...)
	input = append(input, salt...)
	input = append(input, h.MarshalBinary()...)
	input = append(input, msg...)

	var state uint64 = 14695981039346656037
	c := NewPoly27()
	for i := 0; i < N; i++ {
		for _, b := range input {
			state ^= uint64(b)
			state *= 1099511628211
		}
		state ^= uint64(i)
		state *= 1099511628211
		c.Coeffs[i] = uint16(state % uint64(Q))
	}
	return c
}

func (pk *NTRUPublicKey) MarshalBinary() []byte {
	return pk.H.MarshalBinary()
}

func UnmarshalNTRUPK(data []byte) (*NTRUPublicKey, error) {
	h, err := UnmarshalBinary(data)
	if err != nil {
		return nil, err
	}
	return &NTRUPublicKey{H: h}, nil
}

func (sk *NTRUPrivateKey) MarshalBinary() []byte {
	buf := make([]byte, 0, 4*PolyBytes)
	buf = append(buf, sk.BTop[0].MarshalBinary()...)
	buf = append(buf, sk.BBot[0].MarshalBinary()...)
	buf = append(buf, sk.BTop[1].MarshalBinary()...)
	buf = append(buf, sk.BBot[1].MarshalBinary()...)
	return buf
}

func UnmarshalNTRUSK(data []byte, pk *NTRUPublicKey) (*NTRUPrivateKey, error) {
	if len(data) < 4*PolyBytes {
		return nil, errShortData
	}
	sigma := math.Sqrt(float64(N)) * 2.0
	var err error
	sk := &NTRUPrivateKey{Sigma: sigma, PK: pk}
	sk.BTop[0], err = UnmarshalBinary(data[0*PolyBytes : 1*PolyBytes])
	if err != nil {
		return nil, err
	}
	sk.BBot[0], err = UnmarshalBinary(data[1*PolyBytes : 2*PolyBytes])
	if err != nil {
		return nil, err
	}
	sk.BTop[1], err = UnmarshalBinary(data[2*PolyBytes : 3*PolyBytes])
	if err != nil {
		return nil, err
	}
	sk.BBot[1], err = UnmarshalBinary(data[3*PolyBytes : 4*PolyBytes])
	if err != nil {
		return nil, err
	}
	return sk, nil
}

const sigBytes = 16 + 31

func (sig *NTRUSignature) MarshalBinary() []byte {
	buf := make([]byte, 0, sigBytes)
	buf = append(buf, sig.Salt[:]...)
	buf = append(buf, Serialize(sig.S2, 9, true)...)
	return buf
}

func UnmarshalNTRUSig(data []byte) (*NTRUSignature, error) {
	if len(data) < sigBytes {
		return nil, errShortData
	}
	sig := &NTRUSignature{}
	copy(sig.Salt[:], data[0:16])
	s2 := Deserialize(data[16:sigBytes], 9, true)
	if s2 == nil {
		return nil, errShortData
	}
	sig.S2 = s2
	return sig, nil
}