package gnarlring

import "testing"

func TestShareSecret(t *testing.T) {
	// Generate a master secret and split into 5 shares.
	_, sk := LWEKeyGen()
	master := sk.S

	shares := ShareSecret(master, 5)
	if len(shares) != 5 {
		t.Fatalf("got %d shares, want 5", len(shares))
	}

	// Verify Σ shares = master.
	sum := NewPoly27()
	for _, s := range shares {
		sum = Add(sum, s)
	}
	if !Equal(sum, master) {
		t.Fatal("Σ shares ≠ master secret")
	}
}

func TestDistributedDecryption(t *testing.T) {
	seed := []byte("common-reference-string")
	a := GenerateSharedA(seed)

	// 5 parties generate keys with shared A.
	k := 5
	pks := make([]*LWEPublicKey, k)
	sks := make([]*LWESecretKey, k)
	for i := 0; i < k; i++ {
		pks[i], sks[i] = LWEKeyGenWithA(a)
	}

	// Aggregate public key.
	aggPK := LWEKeyAgg(pks)

	// Encrypt under aggregate key.
	for _, bit := range []int{0, 1} {
		correct := 0
		secrets := make([]*Poly27, k)
		for i := 0; i < k; i++ {
			secrets[i] = sks[i].S
		}
		for trial := 0; trial < 50; trial++ {
			ct := EncryptGroup(aggPK, bit)

			// Each party computes partial decryption.
			partials := make([]*Poly27, k)
			for i := 0; i < k; i++ {
				partials[i] = PartialDecryption(secrets[i], ct.U)
			}

			dec := DistributedDecrypt(ct, partials)
			if dec == bit {
				correct++
			}
		}
		rate := float64(correct) / 50.0
		t.Logf("bit=%d: %d/50 correct (%.1f%%)", bit, correct, rate*100)
		if rate < 0.85 {
			t.Errorf("distributed decryption rate %.1f%% < 85%% for bit=%d", rate*100, bit)
		}
	}
}

func TestDistributedDecryptionSubset(t *testing.T) {
	seed := []byte("threshold-crs")
	a := GenerateSharedA(seed)

	k := 7
	pks := make([]*LWEPublicKey, k)
	sks := make([]*LWESecretKey, k)
	for i := 0; i < k; i++ {
		pks[i], sks[i] = LWEKeyGenWithA(a)
	}
	aggPK := LWEKeyAgg(pks)

	// Encrypt under aggregate key.
	ct := EncryptGroup(aggPK, 1)

	// 3 of 7 parties decrypt (subsets work with additive shares).
	partials := make([]*Poly27, 3)
	for i := 0; i < 3; i++ {
		partials[i] = PartialDecryption(sks[i].S, ct.U)
	}

	// Incomplete share set — decryption should fail (the sum of 3 shares
	// doesn't decode correctly because the encrypted aggregate uses
	// all 7 secrets).
	dec := DistributedDecrypt(ct, partials)
	t.Logf("subset decryption (3/7): %d (expected failure)", dec)
	// This will usually NOT decode correctly with only 3 of 7 shares.
}

func TestSharedADeterministic(t *testing.T) {
	seed := []byte("deterministic-a")
	a1 := GenerateSharedA(seed)
	a2 := GenerateSharedA(seed)
	if !Equal(a1, a2) {
		t.Fatal("shared A is not deterministic")
	}

	a3 := GenerateSharedA([]byte("different-seed"))
	if Equal(a1, a3) {
		t.Fatal("different seeds produced same A")
	}
}