package cayley

import (
	"math"
	"math/rand"
	"testing"
)

// TrapdoorTable stores precomputed paths from root to M vertices.
// The public key is the generator set. The trapdoor is the path table.
type TrapdoorTable struct {
	GS    *GeneratorSet
	Paths map[Mat2][]int8 // vertex → path from ID
	M     int              // precomputed table size
}

// BuildTrapdoor precomputes paths to M random vertices using BFS from root.
func BuildTrapdoor(gs *GeneratorSet, M int) *TrapdoorTable {
	pf := gs.BFS(-1) // BFS to find optimal paths to ALL vertices (one-time cost)

	// Randomly select M vertices to keep in the table.
	rng := rand.New(rand.NewSource(42))
	allKeys := make([]Mat2, 0, len(pf.Dist))
	for k := range pf.Dist {
		allKeys = append(allKeys, k)
	}
	rng.Shuffle(len(allKeys), func(i, j int) {
		allKeys[i], allKeys[j] = allKeys[j], allKeys[i]
	})

	tt := &TrapdoorTable{
		GS:    gs,
		Paths: make(map[Mat2][]int8, M),
		M:     M,
	}
	for i := 0; i < M && i < len(allKeys); i++ {
		path, ok := pf.PathTo(allKeys[i])
		if ok {
			tt.Paths[allKeys[i]] = path
		}
	}
	return tt
}

// Sign attempts to sign a target vertex using the trapdoor table.
// Returns (path, found) where found indicates whether the target was in the table.
func (tt *TrapdoorTable) Sign(target Mat2) ([]int8, bool) {
	path, ok := tt.Paths[target]
	return path, ok
}

// Coverage returns the fraction of |G| covered by the table.
func (tt *TrapdoorTable) Coverage() float64 {
	gs := tt.GS
	volume := gs.P * (gs.P*gs.P - 1)
	return float64(tt.M) / float64(volume)
}

func TestTrapdoorHitRate(t *testing.T) {
	primes := []int64{31, 101}
	rng := rand.New(rand.NewSource(77))

	for _, P := range primes {
		gs := randomGens(P, 4, rng)
		volume := P * (P*P - 1)

		// Vary table size M.
		for _, M := range []int{100, 500, 1000, 5000, 10000} {
			if M > int(volume)*9/10 {
				continue
			}
			tt := BuildTrapdoor(gs, M)
			coverage := float64(M) / float64(volume) * 100

			// Measure hit rate: generate random targets, count table matches.
			hits := 0
			trials := 5000
			if M < 1000 {
				trials = 1000
			}
			for i := 0; i < trials; i++ {
				target := randomSL2Fast(P, rng)
				if _, ok := tt.Sign(target); ok {
					hits++
				}
			}
			hitRate := float64(hits) / float64(trials) * 100

			if M <= 1000 || M == 10000 || (P == 31 && M == 5000) {
				t.Logf("P=%d M=%d/%d (%.2f%%): hit rate=%.2f%% (%d/%d)",
					P, M, volume, coverage, hitRate, hits, trials)
			}
		}
	}
}

// TestTrapdoorSigningCost measures the path length of trapdoor signatures
// vs BFS optimal. When the target is in the table, the trapdoor path IS
// the BFS optimal path — no degradation.
func TestTrapdoorSigningCost(t *testing.T) {
	P := int64(101)
	rng := rand.New(rand.NewSource(99))
	gs := randomGens(P, 4, rng)
	pf := gs.BFS(-1) // optimal reference

	M := 10000
	tt := BuildTrapdoor(gs, M)

	// Measure 100 random targets: BFS optimal length vs trapdoor length.
	var optTotal, trapTotal int
	n := 0
	for i := 0; i < 100; i++ {
		target := randomSL2Fast(P, rng)
		optPath, _ := pf.PathTo(target)
		trapPath, found := tt.Sign(target)

		if !found {
			continue // not in table → can't sign (measurement: hit rate is coverage × |G|)
		}
		optTotal += len(optPath)
		trapTotal += len(trapPath)
		n++
	}

	if n > 0 {
		optMean := float64(optTotal) / float64(n)
		trapMean := float64(trapTotal) / float64(n)
		t.Logf("P=%d M=%d n=%d:", P, M, n)
		t.Logf("  BFS optimal mean: %.1f steps", optMean)
		t.Logf("  Trapdoor mean:    %.1f steps", trapMean)
		t.Logf("  Ratio: %.2f (trapdoor = optimal when target in table)", trapMean/optMean)
	}
}

// TestTrapdoorKeygenCost measures the one-time cost of building the trapdoor.
func TestTrapdoorKeygenCost(t *testing.T) {
	primes := []int64{31, 101}

	for _, P := range primes {
		volume := P * (P*P - 1)

		// Keygen cost: BFS over the ENTIRE Cayley graph.
		// This is O(|G|) — one-time at keygen. The trapdoor is a subset of the BFS results.
		t.Logf("P=%d |G|=%d: BFS to build path table (one-time keygen cost)...", P, volume)
		// BFS already done in BuildTrapdoor — time is captured by the BFS enumeration test.
	}
}

// TestTrapdoorScaling estimates production parameters by extrapolation.
func TestTrapdoorScaling(t *testing.T) {
	t.Log("=== TRAPDOOR PATH TABLE SCALING ===")
	t.Log("")
	t.Log("The trapdoor is a precomputed path table. Keygen runs BFS to find")
	t.Log("optimal paths to all vertices, then keeps M entries.")
	t.Log("")
	t.Log("At production scale (P~2^216):")
	t.Log("  |SL(2,Z_P)| = P^3 ≈ 2^648 — BFS infeasible (exponential)")
	t.Log("  Keygen must use a DIFFERENT method:")
	t.Log("    - Signer samples M random walks of length D from root")
	t.Log("    - Records endpoints and paths")
	t.Log("    - Table size M determines hit rate")
	t.Log("")
	t.Log("  For b=8 generators, diameter ≈ 1.3×log_b(|G|) ≈ 281 steps")
	t.Log("  Each random walk of length D covers 1 vertex")
	t.Log("  To cover 0.1% of |G|: M = 0.001×2^648 ≈ 2^638 — infeasible")
	t.Log("")
	t.Log("  The naive path table is NOT a viable trapdoor at production scale.")
	t.Log("  The trapdoor must give the signer a way to reach MANY vertices")
	t.Log("  from FEW precomputed paths — not a 1:1 table.")
	t.Log("")

	// Demonstrate the combinatorial explosion.
	b := 8
	t.Logf("With b=%d generators, each table entry enables reaching ALL vertices", b)
	t.Logf("reachable via PRODUCTS of that entry with other entries.")
	t.Logf("")

	// How many unique vertices can be reached by products of k table entries?
	for k := 1; k <= 6; k++ {
		products := math.Pow(float64(b), float64(k))
		t.Logf("  k=%d: b^k = %.0f unique products", k, products)
	}
	t.Logf("")
	t.Logf("With M=1000 table entries, the signer can reach b^D vertices precisely")
	t.Logf("(all walks of length D). The table is a BASIS of SHORT PATHS, not a")
	t.Logf("lookup of endpoints. The signer decomposes the target into combinations")
	t.Logf("of basis paths. This IS tree-SIS but with precomputed short walks.")
	t.Logf("")
	t.Logf("The trapdoor reduces tree-SIS to a CLOSER-TO-TARGET decomposition:")
	t.Logf("  1. Given target T, find the closest table entry e_i (by graph distance)")
	t.Logf("  2. Compose: signature = path_to(e_i) + short_correction_jump(e_i, T)")
	t.Logf("  3. The correction jump is short because e_i is close to T")
	t.Logf("")
	t.Logf("At production scale: the table entries are carefully chosen ANCHOR")
	t.Logf("points that cover the Cayley graph uniformly. The signer's advantage")
	t.Logf("comes from knowing these anchor positions and their shortest paths.")
	t.Logf("The attacker must discover them from scratch.")
}