package gnarlring

import (
	"testing"

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

func TestPoly27New(t *testing.T) {
	p := NewPoly27()
	if !IsZero(p) {
		t.Fatal("NewPoly27 should be zero")
	}
	if p.isNTT {
		t.Fatal("NewPoly27 should be coefficient form")
	}
}

func TestPoly27Clone(t *testing.T) {
	p := NewPoly27()
	p.Coeffs[0] = 42
	p.isNTT = true

	c := p.Clone()
	if !Equal(p, c) {
		t.Fatal("Clone not equal")
	}
	c.Coeffs[0] = 99
	if p.Coeffs[0] != 42 {
		t.Fatal("Clone shares backing array")
	}
}

func TestAddSub(t *testing.T) {
	a := NewPoly27()
	b := NewPoly27()
	a.Coeffs[0] = 100
	b.Coeffs[0] = 200

	sum := Add(a, b)
	if sum.Coeffs[0] != 29 { // 100+200=300 ≡ 29 mod 271
		t.Fatalf("Add: got %d, want 29", sum.Coeffs[0])
	}

	diff := Sub(b, a)
	if diff.Coeffs[0] != 100 {
		t.Fatalf("Sub: got %d, want 100", diff.Coeffs[0])
	}
}

func TestNeg(t *testing.T) {
	a := NewPoly27()
	a.Coeffs[0] = 42

	n := Neg(a)
	if n.Coeffs[0] != Q-42 {
		t.Fatalf("Neg: got %d, want %d", n.Coeffs[0], Q-42)
	}

	z := NewPoly27()
	nz := Neg(z)
	if nz.Coeffs[0] != 0 {
		t.Fatal("Neg of zero should be zero")
	}
}

func TestScalarMul(t *testing.T) {
	a := NewPoly27()
	a.Coeffs[0] = 50

	s := ScalarMul(a, 10)
	if s.Coeffs[0] != crypto.Mod271(500) {
		t.Fatalf("ScalarMul: got %d, want %d", s.Coeffs[0], crypto.Mod271(500))
	}
}

func TestNTTINTTRoundTrip(t *testing.T) {
	p := NewPoly27()
	for i := range N {
		p.Coeffs[i] = uint16(i * 10 % Q)
	}
	orig := p.Clone()

	p.NTT()
	if !p.isNTT {
		t.Fatal("NTT did not set isNTT")
	}

	p.INTT()
	if p.isNTT {
		t.Fatal("INTT did not clear isNTT")
	}

	if !Equal(p, orig) {
		t.Fatal("NTT/INTT round-trip failed")
	}
}

func TestNTTIdempotent(t *testing.T) {
	p := NewPoly27()
	p.Coeffs[0] = 1
	p.NTT()
	c := p.Clone()
	p.NTT() // should be no-op
	if !Equal(p, c) {
		t.Fatal("NTT not idempotent")
	}
}

func TestMulViaNTTvsSchoolbook(t *testing.T) {
	a := NewPoly27()
	b := NewPoly27()
	for i := range N {
		a.Coeffs[i] = uint16(i + 1)
		b.Coeffs[i] = uint16(N - i)
	}

	result := Mul(a, b)

	// Schoolbook multiplication mod (x^27 + 1).
	var expected [N]uint16
	for i := 0; i < N; i++ {
		for j := 0; j < N; j++ {
			k := i + j
			if k < N {
				expected[k] = addMod(expected[k], crypto.Mod271(uint32(a.Coeffs[i])*uint32(b.Coeffs[j])))
			} else {
				// x^27 = -1, so x^k = -x^{k-27}
				k -= N
				prod := crypto.Mod271(uint32(a.Coeffs[i]) * uint32(b.Coeffs[j]))
				expected[k] = subMod(expected[k], prod)
			}
		}
	}

	for i, v := range result.Coeffs {
		if v != expected[i] {
			t.Fatalf("mul mismatch at %d: got %d, want %d", i, v, expected[i])
		}
	}
}

func TestInverse(t *testing.T) {
	// 1^{-1} = 1.
	a := NewPoly27()
	a.Coeffs[0] = 1
	inv := Inverse(a)
	if inv == nil {
		t.Fatal("Inverse(1) returned nil")
	}
	prod := Mul(a, inv)
	if prod.Coeffs[0] != 1 {
		t.Fatalf("1 * 1^{-1}[0] = %d, want 1", prod.Coeffs[0])
	}
	for i := 1; i < N; i++ {
		if prod.Coeffs[i] != 0 {
			t.Fatalf("1 * 1^{-1}[%d] = %d, want 0", i, prod.Coeffs[i])
		}
	}

	// Random invertible polynomial: non-zero constant.
	a2 := NewPoly27()
	a2.Coeffs[0] = 42
	inv2 := Inverse(a2)
	if inv2 == nil {
		t.Fatal("Inverse(42) returned nil (constant in NTT domain should be [42,...,42])")
	}
	prod2 := Mul(a2, inv2)
	if prod2.Coeffs[0] != 1 {
		t.Fatalf("42 * 42^{-1}[0] = %d, want 1", prod2.Coeffs[0])
	}

	// Complex polynomial: sparse with known non-zero NTT values.
	// f(x) = 1 + 2x + 3x^2. After NTT, each slot is f(-psi^(2j+1)).
	a3 := NewPoly27()
	a3.Coeffs[0] = 1
	a3.Coeffs[1] = 2
	a3.Coeffs[2] = 3
	inv3 := Inverse(a3)
	if inv3 == nil {
		t.Fatal("Inverse(1+2x+3x^2) returned nil")
	}
	prod3 := Mul(a3, inv3)
	if prod3.Coeffs[0] != 1 {
		t.Fatalf("round-trip failed for 1+2x+3x^2: got %d", prod3.Coeffs[0])
	}
}

func TestInverseNonInvertible(t *testing.T) {
	// Zero polynomial is not invertible.
	z := NewPoly27()
	if Inverse(z) != nil {
		t.Fatal("zero should not be invertible")
	}

	// A polynomial with a zero in the NTT domain is not invertible.
	// x + 1 has x=-1 as a root... but we need to check.
	// Simpler: any polynomial where all coefficients are Q is zero = not invertible.
	qAll := NewPoly27()
	for i := range N {
		qAll.Coeffs[i] = Q // same as 0 mod Q
	}
	if Inverse(qAll) != nil {
		t.Fatal("Q-vector should not be invertible")
	}
}

func TestNorm(t *testing.T) {
	// 200 > 135: centered value = 200-271 = -71, abs = 71.
	p := NewPoly27()
	p.Coeffs[0] = 200
	if Norm(p) != 71 {
		t.Fatalf("Norm of 200 = %d, want 71", Norm(p))
	}

	// 42 ≤ 135: centered value = 42, abs = 42.
	p.Coeffs[0] = 42
	if Norm(p) != 42 {
		t.Fatalf("Norm of 42 = %d, want 42", Norm(p))
	}

	// Q-5 = 266 > 135: centered value = -5, abs = 5.
	p.Coeffs[0] = Q - 5
	if Norm(p) != 5 {
		t.Fatalf("Norm of Q-5 = %d, want 5", Norm(p))
	}
}

func TestNormSq(t *testing.T) {
	p := NewPoly27()
	p.Coeffs[0] = 3
	p.Coeffs[1] = 4

	nsq := NormSq(p)
	if nsq != 25 { // 3² + 4² = 25
		t.Fatalf("NormSq = %d, want 25", nsq)
	}
}

func TestDot(t *testing.T) {
	a := NewPoly27()
	b := NewPoly27()
	a.Coeffs[0] = 3
	a.Coeffs[1] = 4
	b.Coeffs[0] = 5
	b.Coeffs[1] = 6

	d := Dot(a, b)
	if d != 39 { // 3*5 + 4*6 = 39
		t.Fatalf("Dot = %d, want 39", d)
	}
}

func TestIsZero(t *testing.T) {
	z := NewPoly27()
	if !IsZero(z) {
		t.Fatal("zero not detected")
	}
	z.Coeffs[0] = 1
	if IsZero(z) {
		t.Fatal("non-zero detected as zero")
	}
}

func TestMarshalUnmarshalRoundTrip(t *testing.T) {
	a := NewPoly27()
	for i := range N {
		a.Coeffs[i] = uint16((i*37 + 1) % Q)
	}

	data := a.MarshalBinary()
	if len(data) != PolyBytes {
		t.Fatalf("MarshalBinary len = %d, want %d", len(data), PolyBytes)
	}

	b, err := UnmarshalBinary(data)
	if err != nil {
		t.Fatal(err)
	}
	if !Equal(a, b) {
		t.Fatal("marshal round-trip failed")
	}
}

func TestUnmarshalBinaryShort(t *testing.T) {
	_, err := UnmarshalBinary([]byte{0, 0})
	if err == nil {
		t.Fatal("expected error for short data")
	}
}

func TestSerializeDeserialize(t *testing.T) {
	a := NewPoly27()
	for i := range N {
		a.Coeffs[i] = uint16((i * 13) % Q)
	}

	// 9-bit unsigned.
	data := Serialize(a, 9, false)
	b := Deserialize(data, 9, false)
	if b == nil || !Equal(a, b) {
		t.Fatal("9-bit unsigned round-trip failed")
	}

	// 8-bit signed (only works if coefficients are in range).
	a2 := NewPoly27()
	for i := range N {
		a2.Coeffs[i] = uint16(i % 60) // < 128, fits in signed 8-bit
	}
	data2 := Serialize(a2, 8, true)
	b2 := Deserialize(data2, 8, true)
	if b2 == nil {
		t.Fatal("8-bit signed Deserialize returned nil")
	}
	for i := range N {
		if b2.Coeffs[i] != a2.Coeffs[i] {
			t.Fatalf("8-bit signed mismatch at %d: got %d, want %d", i, b2.Coeffs[i], a2.Coeffs[i])
		}
	}

	// Signed negative test: Q-1 is centered as -1 → abs=1.
	a3 := NewPoly27()
	a3.Coeffs[0] = Q - 1 // centered = -1
	data3 := Serialize(a3, 8, true)
	b3 := Deserialize(data3, 8, true)
	if b3 == nil || b3.Coeffs[0] != Q-1 {
		t.Fatalf("signed negative round-trip: got %d, want %d", b3.Coeffs[0], Q-1)
	}
}

func TestSerializeDeserializeShort(t *testing.T) {
	b := Deserialize([]byte{0}, 9, false)
	if b != nil {
		t.Fatal("expected nil for short data")
	}
}