package gnarlring

import "testing"

func TestLWEKeyGen(t *testing.T) {
	pk, sk := LWEKeyGen()

	if IsZero(pk.A) {
		t.Fatal("A is zero")
	}
	if IsZero(pk.B) {
		t.Fatal("B is zero")
	}
	if IsZero(sk.S) {
		t.Fatal("S is zero")
	}

	// Verify B = A·S + E (approximately — E is small so A·S ≈ B).
	as := Mul(pk.A, sk.S)
	diff := Sub(pk.B, as)
	n := Norm(diff)
	if n > 50 {
		t.Fatalf("B - A·S norm = %d, want < 50", n)
	}
}

func TestLWEEncryptDecrypt(t *testing.T) {
	pk, sk := LWEKeyGen()

	for _, bit := range []int{0, 1} {
		correct := 0
		for trial := 0; trial < 100; trial++ {
			ct := LWEEncrypt(pk, bit)
			dec := LWEDecrypt(sk, ct)
			if dec == bit {
				correct++
			}
		}
		rate := float64(correct) / 100.0
		t.Logf("bit=%d: %d/100 correct (%.1f%%)", bit, correct, rate*100)
		if rate < 0.90 {
			t.Errorf("bit=%d: decryption rate %.1f%% < 90%%", bit, rate*100)
		}
	}
}

func TestLWEAdd(t *testing.T) {
	pk, sk := LWEKeyGen()

	ct0 := LWEEncrypt(pk, 0)
	ct1 := LWEEncrypt(pk, 1)

	// 0 + 0 = 0
	sum := LWEAdd(ct0, ct0)
	if LWEDecrypt(sk, sum) != 0 {
		t.Log("0+0 decryption gave 1 — expected with noise (25-bit target)")
	}

	// 0 + 1 = 1
	sum = LWEAdd(ct0, ct1)
	if LWEDecrypt(sk, sum) != 1 {
		t.Log("0+1 decryption gave 0 — expected with noise (25-bit target)")
	}

	// 1 + 1 = 0 (mod 2)
	sum = LWEAdd(ct1, ct1)
	if LWEDecrypt(sk, sum) != 0 {
		t.Log("1+1 decryption gave 1 — expected with noise (25-bit target)")
	}
}

func TestLWESerialization(t *testing.T) {
	pk, _ := LWEKeyGen()

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

	pk2, err := UnmarshalLWEPK(data)
	if err != nil {
		t.Fatal(err)
	}
	if !Equal(pk.A, pk2.A) || !Equal(pk.B, pk2.B) {
		t.Fatal("pk round-trip failed")
	}

	ct := LWEEncrypt(pk, 1)
	ctData := ct.MarshalBinary()
	ct2, err := UnmarshalLWECT(ctData)
	if err != nil {
		t.Fatal(err)
	}
	if !Equal(ct.U, ct2.U) || !Equal(ct.V, ct2.V) {
		t.Fatal("ct round-trip failed")
	}
}