package crypto

import (
	"testing"

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

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

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

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

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

	// Challenge must round-trip exactly.
	if sig2.Challenge != sig.Challenge {
		t.Error("challenge mismatch after round-trip")
	}

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

	// Proof lengths must be consistent.
	if len(sig2.Proof.LockIns) != len(sig2.Proof.NeighborCounts) {
		t.Error("proof length mismatch in round-trip")
	}
	if len(sig2.Proof.LockIns) != len(sig2.Proof.HexTrace) {
		t.Error("proof/hex trace length mismatch in round-trip")
	}
}

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

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

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

	t.Logf("compact signature size: %d bytes", len(data))

	// For Security128 (N=256), compact should be well under 2 KB.
	if len(data) > 2048 {
		t.Errorf("compact signature too large: %d bytes (want <= 2048)", len(data))
	}
}

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

func TestUnmarshalTruncated(t *testing.T) {
	_, err := UnmarshalSignature([]byte{1, 2, 3})
	if err == nil {
		t.Error("expected error with truncated data")
	}
}

func TestQuantizeLockIn(t *testing.T) {
	tests := []struct {
		name string
		num  int64
		den  int64
		want uint8
	}{
		{"zero", 0, 1, 0},
		{"half", 1, 2, 127},
		{"one", 1, 1, 255},
		{"quarter", 1, 4, 63},
		{"over_one", 3, 2, 255}, // clamped
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			r := ratio.New(tt.num, tt.den)
			got := quantizeLockIn(r)
			if got != tt.want {
				t.Errorf("quantizeLockIn(%s) = %d, want %d", r, got, tt.want)
			}
		})
	}
}

func TestDequantizeRoundTrip(t *testing.T) {
	// Quantize then dequantize should be approximately equal.
	r := ratio.New(3, 7)
	q := quantizeLockIn(r)
	back := dequantizeLockIn(q)

	// The dequantized value should be within 1/255 of the original.
	diff := r.Sub(back).Abs()
	tolerance := ratio.New(2, 255)
	if diff.Greater(tolerance) {
		t.Errorf("round-trip error too large: %s → %d → %s (diff=%s)",
			r, q, back, diff)
	}
}

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

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

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

	if string(data1) != string(data2) {
		t.Error("marshaling the same signature twice should produce identical bytes")
	}
}