package crypto

import (
	"testing"

	"git.mleku.dev/mleku/dendrite/pkg/spore"
)

func TestSubsetRoundTrip(t *testing.T) {
	tests := []struct {
		name    string
		indices []int
		n       int
	}{
		{"single", []int{0}, 256},
		{"first_last", []int{0, 255}, 256},
		{"consecutive", []int{10, 11, 12, 13, 14}, 256},
		{"spread", []int{0, 50, 100, 150, 200, 250}, 256},
		{"large_k", func() []int {
			idx := make([]int, 56)
			for i := range idx {
				idx[i] = i * 4 // every 4th site
			}
			return idx
		}(), 256},
		{"n512", []int{0, 100, 200, 300, 400, 500}, 512},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			encoded := encodeSubset(tt.indices, tt.n)
			decoded := decodeSubset(encoded, tt.n, len(tt.indices))

			if len(decoded) != len(tt.indices) {
				t.Fatalf("length mismatch: got %d, want %d", len(decoded), len(tt.indices))
			}
			for i := range tt.indices {
				if decoded[i] != tt.indices[i] {
					t.Errorf("index %d: got %d, want %d", i, decoded[i], tt.indices[i])
				}
			}
		})
	}
}

func TestSubsetEncodingSize(t *testing.T) {
	// C(256,56) has 190 bits → 24 bytes.
	nbytes := subsetBytes(256, 56)
	t.Logf("C(256,56) encoding: %d bytes (%d bits)", nbytes, bigBinom(256, 56).BitLen())
	if nbytes != 24 {
		t.Errorf("subset encoding unexpected: %d bytes (want 24)", nbytes)
	}

	// C(512,56) has 251 bits → 32 bytes.
	nbytes512 := subsetBytes(512, 56)
	t.Logf("C(512,56) encoding: %d bytes (%d bits)", nbytes512, bigBinom(512, 56).BitLen())
	if nbytes512 != 32 {
		t.Errorf("subset encoding unexpected: %d bytes (want 32)", nbytes512)
	}
}

func TestBitWriterReader(t *testing.T) {
	bw := newBitWriter(16)
	bw.writeBits(5, 3)   // 101
	bw.writeBits(200, 8) // 11001000
	bw.writeBits(3, 3)   // 011
	bw.writeBits(0, 2)   // 00

	data := bw.bytes()

	br := newBitReader(data)
	v1, _ := br.readBits(3)
	v2, _ := br.readBits(8)
	v3, _ := br.readBits(3)
	v4, _ := br.readBits(2)

	if v1 != 5 {
		t.Errorf("v1: got %d, want 5", v1)
	}
	if v2 != 200 {
		t.Errorf("v2: got %d, want 200", v2)
	}
	if v3 != 3 {
		t.Errorf("v3: got %d, want 3", v3)
	}
	if v4 != 0 {
		t.Errorf("v4: got %d, want 0", v4)
	}
}

func TestMarshalV2RoundTrip(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	msg := []byte("compact v2 round-trip")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	data, err := sig.MarshalV2(params.N)
	if err != nil {
		t.Fatalf("MarshalV2: %v", err)
	}

	sig2, consumed, err := UnmarshalSignatureV2(data)
	if err != nil {
		t.Fatalf("UnmarshalSignatureV2: %v", err)
	}
	if consumed != len(data) {
		t.Errorf("consumed %d bytes, data is %d bytes", consumed, len(data))
	}

	// Occupied count must match.
	origOcc := 0
	for _, s := range sig.Response {
		if s.Occupied {
			origOcc++
		}
	}
	if len(sig2.Response) != origOcc {
		t.Errorf("occupied count: orig=%d, round-trip=%d", origOcc, len(sig2.Response))
	}

	// Per-site data must match for occupied sites.
	origSites := make([]SiteMark, 0, origOcc)
	for _, s := range sig.Response {
		if s.Occupied {
			origSites = append(origSites, s)
		}
	}

	for i, s2 := range sig2.Response {
		if i >= len(origSites) {
			break
		}
		s1 := origSites[i]
		if s2.Index != s1.Index {
			t.Errorf("site %d: index %d != %d", i, s2.Index, s1.Index)
		}
		if s2.TypeTag != s1.TypeTag {
			t.Errorf("site %d: tag %q != %q", i, s2.TypeTag, s1.TypeTag)
		}
		if s2.Projection != s1.Projection {
			t.Errorf("site %d: projection %d != %d", i, s2.Projection, s1.Projection)
		}
		if s2.Perm != s1.Perm {
			t.Errorf("site %d: perm %d != %d", i, s2.Perm, s1.Perm)
		}
	}

	// Proof lengths must be consistent.
	if len(sig2.Proof.LockIns) != len(sig2.Response) {
		t.Error("proof/response length mismatch")
	}
}

func TestMarshalV2Size(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	msg := []byte("size test v2")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	data, err := sig.MarshalV2(params.N)
	if err != nil {
		t.Fatalf("MarshalV2: %v", err)
	}

	v1data, err := sig.Marshal()
	if err != nil {
		t.Fatalf("Marshal V1: %v", err)
	}

	t.Logf("V1 compact: %d bytes", len(v1data))
	t.Logf("V2 compact: %d bytes", len(data))
	t.Logf("compression ratio: %.1f%%", 100.0*(1.0-float64(len(data))/float64(len(v1data))))

	// V2 should be under 200 bytes for N=256.
	if len(data) > 200 {
		t.Errorf("V2 signature too large: %d bytes (want <= 200)", len(data))
	}
}

func TestMarshalV2Deterministic(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	msg := []byte("deterministic v2")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	data1, err := sig.MarshalV2(params.N)
	if err != nil {
		t.Fatalf("MarshalV2 1: %v", err)
	}
	data2, err := sig.MarshalV2(params.N)
	if err != nil {
		t.Fatalf("MarshalV2 2: %v", err)
	}

	if string(data1) != string(data2) {
		t.Error("MarshalV2 is not deterministic")
	}
}

func TestVerifyCompactV2(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	msg := []byte("verify compact v2")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	data, err := sig.MarshalV2(params.N)
	if err != nil {
		t.Fatalf("MarshalV2: %v", err)
	}

	sig2, _, err := UnmarshalSignatureV2(data)
	if err != nil {
		t.Fatalf("UnmarshalSignatureV2: %v", err)
	}

	fp := FingerprintFromSpore(spore.Extract(kp.Private.Lattice))
	if !VerifyCompactV2(fp, msg, sig2) {
		t.Error("VerifyCompactV2 failed on valid signature")
	}

	// Wrong message should fail.
	if VerifyCompactV2(fp, []byte("wrong message"), sig2) {
		t.Error("VerifyCompactV2 accepted wrong message")
	}
}

func TestMarshalV2NilSignature(t *testing.T) {
	var s *Signature
	_, err := s.MarshalV2(256)
	if err == nil {
		t.Error("expected error marshaling nil signature")
	}
}

func TestUnmarshalV2Truncated(t *testing.T) {
	_, _, err := UnmarshalSignatureV2([]byte{0xD2, 0x00})
	if err == nil {
		t.Error("expected error with truncated data")
	}
}

func TestUnmarshalV2BadMagic(t *testing.T) {
	_, _, err := UnmarshalSignatureV2([]byte{0xFF, 0x00, 0x01, 0x01, 0x01, 'x', 0x01, 0x00})
	if err == nil {
		t.Error("expected error with bad magic")
	}
}