package crypto

import (
	"bytes"
	"testing"

	"git.mleku.dev/mleku/dendrite/pkg/epoch"
	"git.mleku.dev/mleku/dendrite/pkg/state"
)

// TestShadowCompressDecompressRoundTrip tests full compress/decompress
// cycle with and without auxiliary data.
func TestShadowCompressDecompressRoundTrip(t *testing.T) {
	ep := epoch.Colony

	testCases := []struct {
		name string
		data []byte
		aux  []byte
	}{
		{"empty", nil, nil},
		{"single byte", []byte{42}, nil},
		{"short text", []byte("hello shadow compression"), nil},
		{"with auxiliary", []byte("primary data stream"), []byte("aux")},
		{"aux 1 byte", []byte("data for one byte aux"), []byte{0xFF}},
		{"aux 2 bytes", []byte("data for two byte aux"), []byte{0x42, 0x37}},
		{"aux 7 bytes", []byte("seven bytes of auxiliary payload here"), []byte("7bytes!")},
		{"long text", bytes.Repeat([]byte("the shadow between binary and decimal "), 20), nil},
		{"long with aux", bytes.Repeat([]byte("compress me "), 50), bytes.Repeat([]byte{0xAB}, 10)},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			sc, err := ShadowCompress(tc.data, tc.aux, ep)
			if err != nil {
				t.Fatalf("ShadowCompress: %v", err)
			}

			if len(tc.data) == 0 {
				if sc.OrigLen != 0 {
					t.Errorf("empty data: OrigLen = %d, want 0", sc.OrigLen)
				}
				return
			}

			if sc.OrigLen != len(tc.data) {
				t.Errorf("OrigLen = %d, want %d", sc.OrigLen, len(tc.data))
			}

			data, aux, err := ShadowDecompress(sc, ep)
			if err != nil {
				t.Fatalf("ShadowDecompress: %v", err)
			}

			if !bytes.Equal(data, tc.data) {
				t.Errorf("data mismatch:\n  got:  %x\n  want: %x", data, tc.data)
			}

			if tc.aux != nil {
				if !bytes.Equal(aux, tc.aux) {
					t.Errorf("auxiliary mismatch:\n  got:  %x\n  want: %x", aux, tc.aux)
				}
			}
		})
	}
}

// TestShadowCompressWireRoundTrip tests marshal/unmarshal cycle.
func TestShadowCompressWireRoundTrip(t *testing.T) {
	ep := epoch.Colony
	data := []byte("wire format round-trip test data")

	sc, err := ShadowCompress(data, []byte("aux"), ep)
	if err != nil {
		t.Fatalf("ShadowCompress: %v", err)
	}

	wire, err := MarshalShadowCompressed(sc)
	if err != nil {
		t.Fatalf("Marshal: %v", err)
	}

	sc2, err := UnmarshalShadowCompressed(wire)
	if err != nil {
		t.Fatalf("Unmarshal: %v", err)
	}

	// Verify all fields match.
	if sc2.OrigLen != sc.OrigLen {
		t.Errorf("OrigLen: %d vs %d", sc2.OrigLen, sc.OrigLen)
	}
	if sc2.TokenCount != sc.TokenCount {
		t.Errorf("TokenCount: %d vs %d", sc2.TokenCount, sc.TokenCount)
	}
	if sc2.EpochDec != sc.EpochDec || sc2.EpochBin != sc.EpochBin {
		t.Errorf("Epoch: %d/%d vs %d/%d", sc2.EpochDec, sc2.EpochBin, sc.EpochDec, sc.EpochBin)
	}
	if sc2.ShadowPayloadLen != sc.ShadowPayloadLen {
		t.Errorf("ShadowPayloadLen: %d vs %d", sc2.ShadowPayloadLen, sc.ShadowPayloadLen)
	}
	if !bytes.Equal(sc2.Primary, sc.Primary) {
		t.Errorf("Primary stream mismatch")
	}
	if !bytes.Equal(sc2.Shadow, sc.Shadow) {
		t.Errorf("Shadow stream mismatch")
	}
	if sc2.FreqTable != sc.FreqTable {
		t.Errorf("FreqTable mismatch")
	}
	if sc2.ContentHash != sc.ContentHash {
		t.Errorf("ContentHash mismatch")
	}

	// Verify full round-trip: unmarshal then decompress.
	recovered, aux, err := ShadowDecompress(sc2, ep)
	if err != nil {
		t.Fatalf("ShadowDecompress after unmarshal: %v", err)
	}
	if !bytes.Equal(recovered, data) {
		t.Errorf("data after wire round-trip: %x != %x", recovered, data)
	}
	if !bytes.Equal(aux, []byte("aux")) {
		t.Errorf("aux after wire round-trip: %x", aux)
	}
}

// TestHuffmanTreeDeterminism verifies that building the Huffman tree
// from the same frequency table always produces identical codebooks.
func TestHuffmanTreeDeterminism(t *testing.T) {
	var freq [64]uint32
	freq[0] = 100
	freq[1] = 50
	freq[2] = 25
	freq[10] = 10
	freq[63] = 5

	total := int64(190)
	var allCodes []map[uint8]int // symbol → code length

	for range 100 {
		tree := buildHuffmanTree(freq, total)
		codes := make(map[uint8]int)
		codeMap := make(map[state.Hexagram]huffCode)
		buildCodes(tree, nil, codeMap)
		for sym, code := range codeMap {
			codes[uint8(sym)] = code.length
		}
		allCodes = append(allCodes, codes)
	}

	// All 100 codebooks should be identical.
	for i := 1; i < len(allCodes); i++ {
		for sym, length := range allCodes[0] {
			if allCodes[i][sym] != length {
				t.Errorf("iteration %d: symbol %d has length %d, want %d",
					i, sym, allCodes[i][sym], length)
			}
		}
	}
}

// TestShadowCompressRatio verifies that compression produces output
// smaller than or comparable to input for repetitive data.
func TestShadowCompressRatio(t *testing.T) {
	ep := epoch.Colony

	// Highly repetitive data should compress well.
	data := bytes.Repeat([]byte("AAAA"), 1000)
	sc, err := ShadowCompress(data, nil, ep)
	if err != nil {
		t.Fatalf("ShadowCompress: %v", err)
	}

	compressedSize := len(sc.Primary) + len(sc.Shadow) + shadowCompressedHeaderSize
	ratio := float64(compressedSize) / float64(len(data))
	t.Logf("repetitive data: %d bytes → %d bytes (ratio %.2f)",
		len(data), compressedSize, ratio)

	// The shadow stream adds 6 bits per token regardless of auxiliary
	// data, so the total output includes structural overhead. For
	// highly repetitive data the Huffman stream compresses well, but
	// the shadow stream is proportional to token count. Ratio under
	// 2.0 indicates the Huffman layer is contributing.
	if ratio > 2.0 {
		t.Errorf("compression ratio %.2f too high for repetitive data", ratio)
	}
}

// TestShadowCompressEpochVariation verifies that different epochs
// produce different shadow streams but recover the same data.
func TestShadowCompressEpochVariation(t *testing.T) {
	data := []byte("epoch variation test")

	epochs := []epoch.Epoch{epoch.Colony, epoch.CryptoWalk128, epoch.CryptoWalk256}

	for _, ep := range epochs {
		sc, err := ShadowCompress(data, nil, ep)
		if err != nil {
			t.Fatalf("epoch %s: ShadowCompress: %v", ep, err)
		}

		recovered, _, err := ShadowDecompress(sc, ep)
		if err != nil {
			t.Fatalf("epoch %s: ShadowDecompress: %v", ep, err)
		}

		if !bytes.Equal(recovered, data) {
			t.Errorf("epoch %s: recovered != original", ep)
		}
	}
}