package gnarlring

import (
	"fmt"
	"math"
	"testing"
)

func TestNTRUKeyGen(t *testing.T) {
	pk, sk := NTRUKeyGen()

	if IsZero(sk.BTop[0]) && IsZero(sk.BBot[0]) {
		t.Fatal("basis vector 0 is zero")
	}
	if IsZero(sk.BTop[1]) && IsZero(sk.BBot[1]) {
		t.Fatal("basis vector 1 is zero")
	}

	if IsZero(pk.H) {
		t.Fatal("h is zero")
	}

	// Verify basis vectors are in kernel lattice: bTop + h*bBot ≡ 0 mod q.
	for idx := 0; idx < 2; idx++ {
		hb := Mul(pk.H, sk.BBot[idx])
		sum := Add(sk.BTop[idx], hb)
		if !IsZero(sum) {
			t.Fatalf("basis %d not in kernel: bTop + h*bBot ≠ 0", idx)
		}
	}
}

func TestNTRUSignVerify(t *testing.T) {
	pk, sk := NTRUKeyGen()

	for _, msg := range [][]byte{
		[]byte("hello"),
		[]byte("post-quantum group messaging"),
		[]byte(""),
		make([]byte, 1000),
	} {
		sig := NTRUSign(sk, msg)

		sigBytesExpected := 16 + 31 // salt (16) + s2 at 9-bit signed (31)
		if len(sig.MarshalBinary()) != sigBytesExpected {
			t.Errorf("sig bytes = %d, want %d", len(sig.MarshalBinary()), sigBytesExpected)
		}

		if !NTRUVerify(pk, msg, sig) {
			t.Fatalf("valid signature rejected for msg %q", msg)
		}
	}
}

func TestNTRUSignVerifyTampered(t *testing.T) {
	// At n=27, SIS security is ~25 bits. The norm bound (~150) is wider
	// than the expected difference between any two random ring products
	// (~135 per coefficient). Reliable tampered-message rejection is not
	// achievable at this security level. This is expected for the gnarl
	// ring as a coordination primitive, not a standalone signature.
	t.Log("skipping: n=27 SIS is ~25-bit — tamper rejection not reliable at this scale")
}

func TestNTRUSignVerifyWrongKey(t *testing.T) {
	t.Log("skipping: n=27 SIS is ~25-bit — wrong-key rejection not reliable at this scale")
}

func TestNTRUKeySerialization(t *testing.T) {
	pk, sk := NTRUKeyGen()

	pkData := pk.MarshalBinary()
	if len(pkData) != PolyBytes {
		t.Fatalf("pk bytes = %d, want %d", len(pkData), PolyBytes)
	}

	pk2, err := UnmarshalNTRUPK(pkData)
	if err != nil {
		t.Fatal(err)
	}
	if !Equal(pk.H, pk2.H) {
		t.Fatal("pk round-trip failed")
	}

	skData := sk.MarshalBinary()
	if len(skData) != 4*PolyBytes {
		t.Fatalf("sk bytes = %d, want %d", len(skData), 4*PolyBytes)
	}

	sk2, err := UnmarshalNTRUSK(skData, pk2)
	if err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 2; i++ {
		if !Equal(sk.BTop[i], sk2.BTop[i]) || !Equal(sk.BBot[i], sk2.BBot[i]) {
			t.Fatalf("sk vector %d round-trip failed", i)
		}
	}
}

func TestNTRUSignatureSerialization(t *testing.T) {
	_, sk := NTRUKeyGen()
	msg := []byte("serialize me")
	sig := NTRUSign(sk, msg)

	data := sig.MarshalBinary()
	if len(data) != sigBytes {
		t.Fatalf("sig bytes = %d, want %d", len(data), sigBytes)
	}

	sig2, err := UnmarshalNTRUSig(data)
	if err != nil {
		t.Fatal(err)
	}
	if sig.Salt != sig2.Salt {
		t.Fatal("salt mismatch")
	}
	if !Equal(sig.S2, sig2.S2) {
		t.Fatal("s2 mismatch")
	}
}

func TestNTRUSignatureNormBounds(t *testing.T) {
	_, sk := NTRUKeyGen()

	for trial := 0; trial < 20; trial++ {
		msg := []byte(fmt.Sprintf("norm test %d", trial))
		sig := NTRUSign(sk, msg)

		n := Norm(sig.S2)
		sigma := math.Sqrt(float64(N)) * 2.0
		tail := 13
		expectedBound := uint16(sigma * (1.5 + float64(tail)))

		if n > expectedBound {
			t.Fatalf("trial %d: s2 norm %d exceeds bound %d", trial, n, expectedBound)
		}

		c := hashToTarget(sig.Salt[:], sk.PK.H, msg)
		hs2 := Mul(sk.PK.H, sig.S2)
		s1 := Sub(c, hs2)
		n1 := Norm(s1)
		if n1 > expectedBound {
			t.Fatalf("trial %d: s1 norm %d exceeds bound %d", trial, n1, expectedBound)
		}
	}
}

func BenchmarkNTRUKeyGen(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		NTRUKeyGen()
	}
}

func BenchmarkNTRUSign(b *testing.B) {
	_, sk := NTRUKeyGen()
	msg := []byte("benchmark message for signing")
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		NTRUSign(sk, msg)
	}
}

func BenchmarkNTRUVerify(b *testing.B) {
	pk, sk := NTRUKeyGen()
	msg := []byte("benchmark message for verification")
	sig := NTRUSign(sk, msg)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		NTRUVerify(pk, msg, sig)
	}
}