package crypto

import (
	"bytes"
	"testing"

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

// TestPhasePermsDeterministic verifies that PhasePerms produces the
// same result for the same epoch and index across multiple calls.
func TestPhasePermsDeterministic(t *testing.T) {
	ep := epoch.Colony
	for i := range 100 {
		b1, d1 := PhasePerms(ep, i)
		b2, d2 := PhasePerms(ep, i)
		if b1 != b2 || d1 != d2 {
			t.Errorf("index %d: PhasePerms not deterministic: (%d,%d) vs (%d,%d)",
				i, b1, d1, b2, d2)
		}
	}
}

// TestPhasePermsRange verifies that PhasePerms always returns valid
// S_3 permutation indices (0-5).
func TestPhasePermsRange(t *testing.T) {
	for _, ep := range []epoch.Epoch{epoch.Colony, epoch.CryptoWalk128, epoch.CryptoWalk256} {
		for i := range int(ep.Period) {
			b, d := PhasePerms(ep, i)
			if b >= permutation.Count || d >= permutation.Count {
				t.Errorf("epoch %s index %d: perm out of range: (%d,%d)", ep, i, b, d)
			}
		}
	}
}

// TestShadowDecomposeRecompose verifies that ShadowDecompose followed
// by ShadowRecompose is the identity for all hexagrams and permutation pairs.
func TestShadowDecomposeRecompose(t *testing.T) {
	for h := range 64 {
		hex := state.Hexagram(h)
		for bp := range permutation.Count {
			for dp := range permutation.Count {
				binP := permutation.Perm(bp)
				decP := permutation.Perm(dp)
				encrypted := ShadowDecompose(hex, binP, decP)
				recovered := ShadowRecompose(encrypted, binP, decP)
				if recovered != hex {
					t.Errorf("hex=%d binP=%d decP=%d: decompose/recompose failed: got %d",
						h, bp, dp, recovered)
				}
			}
		}
	}
}

// TestPhasePairIndexRoundTrip verifies that PhasePairIndex and
// PairFromIndex are exact inverses for all 36 pairs across all
// 36 possible home positions. This is the exhaustive check: every
// projection must have an inverse that recovers the original.
func TestPhasePairIndexRoundTrip(t *testing.T) {
	for hb := range permutation.Count {
		for hd := range permutation.Count {
			homeBin := permutation.Perm(hb)
			homeDec := permutation.Perm(hd)

			// Forward: every pair must map to a unique index 0-35.
			seen := make(map[uint8]bool)
			for bp := range permutation.Count {
				for dp := range permutation.Count {
					binP := permutation.Perm(bp)
					decP := permutation.Perm(dp)
					idx := PhasePairIndex(binP, decP, homeBin, homeDec)
					if idx > 35 {
						t.Errorf("home=(%d,%d) pair=(%d,%d): index %d out of range",
							hb, hd, bp, dp, idx)
					}
					if seen[idx] {
						t.Errorf("home=(%d,%d) pair=(%d,%d): duplicate index %d",
							hb, hd, bp, dp, idx)
					}
					seen[idx] = true
				}
			}

			// All 36 indices must be used.
			for i := range 36 {
				if !seen[uint8(i)] {
					t.Errorf("home=(%d,%d): index %d not produced by any pair", hb, hd, i)
				}
			}

			// Inverse: PairFromIndex must recover the original pair.
			for bp := range permutation.Count {
				for dp := range permutation.Count {
					binP := permutation.Perm(bp)
					decP := permutation.Perm(dp)
					idx := PhasePairIndex(binP, decP, homeBin, homeDec)
					rb, rd := PairFromIndex(idx, homeBin, homeDec)
					if rb != binP || rd != decP {
						t.Errorf("home=(%d,%d) pair=(%d,%d) -> idx=%d -> (%d,%d): round-trip failed",
							hb, hd, bp, dp, idx, rb, rd)
					}
				}
			}
		}
	}
}

// TestShadowEncryptDecryptRoundTrip tests full encrypt/decrypt cycle.
func TestShadowEncryptDecryptRoundTrip(t *testing.T) {
	params := DefaultParams(Security128)
	tags := []string{"alpha", "beta", "gamma"}

	// Use the existing Generate function to create a keypair.
	kp, err := Generate(params, tags, func(tag string) axiom.Constraint {
		return publicConstraint{tag: tag}
	})
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}

	ep := epoch.Colony
	nonce := []byte("test-nonce-unique-per-message")

	// Test cases must fit within the lattice capacity: N=256 nodes
	// shared across len(tags)=3 type layers. With hexagram encoding
	// (3 bytes → 4 tokens), the max plaintext is roughly N*3/4 bytes,
	// minus dissolution losses. Keep inputs small.
	testCases := [][]byte{
		{},
		{0x42},
		[]byte("hello"),
		bytes.Repeat([]byte("shadow "), 5),
	}

	for i, plaintext := range testCases {
		if len(plaintext) == 0 {
			continue // empty plaintext handled separately
		}

		ct, err := ShadowEncrypt(&kp.Public, plaintext, params, ep, nonce)
		if err != nil {
			t.Errorf("case %d: ShadowEncrypt: %v", i, err)
			continue
		}

		if ct.OrigLen != len(plaintext) {
			t.Errorf("case %d: OrigLen = %d, want %d", i, ct.OrigLen, len(plaintext))
		}
		if ct.EpochDec != ep.DecExp || ct.EpochBin != ep.BinExp {
			t.Errorf("case %d: epoch mismatch", i)
		}

		recovered, err := ShadowDecrypt(&kp.Private, ct, ep)
		if err != nil {
			t.Errorf("case %d: ShadowDecrypt: %v", i, err)
			continue
		}

		if !bytes.Equal(recovered, plaintext) {
			t.Errorf("case %d: recovered != plaintext\n  got:  %x\n  want: %x",
				i, recovered, plaintext)
		}
	}
}

// TestShadowSemanticSecurity checks that encrypted hexagram tokens
// have roughly uniform distribution for plaintext with varied content.
//
// Note: all-zero plaintext maps to hexagram 0 = (0,0,0)/(0,0,0),
// and any S_3 permutation of identical elements is identity. So we
// test with data that produces varied hexagram values to verify the
// phase-dependent permutations actually disperse the distribution.
func TestShadowSemanticSecurity(t *testing.T) {
	ep := epoch.Colony

	// Sequential bytes produce varied hexagram values.
	plaintext := make([]byte, 300)
	for i := range plaintext {
		plaintext[i] = byte(i)
	}
	tokens := state.EncodeBytes(plaintext)

	// Original distribution.
	var origHist [64]int
	for _, tok := range tokens {
		origHist[tok&0x3F]++
	}

	// Encrypted distribution.
	encrypted := make([]state.Hexagram, len(tokens))
	var encHist [64]int
	for i, tok := range tokens {
		binP, decP := PhasePerms(ep, i)
		encrypted[i] = ShadowDecompose(tok, binP, decP)
		encHist[encrypted[i]&0x3F]++
	}

	// Check that encryption changes the distribution —
	// at least some hexagrams should differ from originals.
	changed := 0
	for i, tok := range tokens {
		if encrypted[i] != tok {
			changed++
		}
	}
	if changed == 0 {
		t.Errorf("semantic security: no tokens were changed by phase permutation")
	}

	// Check that we see reasonable diversity in encrypted output.
	distinct := 0
	for _, count := range encHist {
		if count > 0 {
			distinct++
		}
	}
	if distinct < 10 {
		t.Errorf("semantic security: only %d distinct encrypted hexagram values", distinct)
	}
}

// TestHexagramEncodingRoundTrip verifies state.EncodeBytes/DecodeHexagrams
// for all single-byte values.
func TestHexagramEncodingRoundTrip(t *testing.T) {
	for b := range 256 {
		data := []byte{byte(b)}
		tokens := state.EncodeBytes(data)
		recovered := state.DecodeHexagrams(tokens, 1)
		if len(recovered) != 1 || recovered[0] != byte(b) {
			t.Errorf("byte %d: round-trip failed: got %v", b, recovered)
		}
	}
}