spatial_memory_test.go

package yaad

import (
	"fmt"
	"sync"
	"testing"
	"time"
)

// newTestSpatialMemory creates a SpatialMemory with tunable timing parameters,
// allowing tests to control promotion/demotion without real-world waits.
func newTestSpatialMemory(hotBudget, warmBudget, coldBudget int, opts ...func(*SpatialMemory)) *SpatialMemory {
	sm := NewSpatialMemory(hotBudget, warmBudget, coldBudget)
	for _, o := range opts {
		o(sm)
	}
	return sm
}

func withHotThreshold(n int) func(*SpatialMemory) {
	return func(sm *SpatialMemory) { sm.hotThreshold = n }
}

func withHotWindow(d time.Duration) func(*SpatialMemory) {
	return func(sm *SpatialMemory) { sm.hotWindow = d }
}

func withHotDemoteAfter(d time.Duration) func(*SpatialMemory) {
	return func(sm *SpatialMemory) { sm.hotDemoteAfter = d }
}

func withWarmDemoteAfter(d time.Duration) func(*SpatialMemory) {
	return func(sm *SpatialMemory) { sm.warmDemoteAfter = d }
}

// ---------------------------------------------------------------------------
// 1. NewSpatialMemory initialization
// ---------------------------------------------------------------------------

func TestNewSpatialMemoryInitialization(t *testing.T) {
	sm := NewSpatialMemory(100, 200, 300)

	if sm.hotBudget != 100 {
		t.Errorf("hotBudget: got %d, want 100", sm.hotBudget)
	}
	if sm.warmBudget != 200 {
		t.Errorf("warmBudget: got %d, want 200", sm.warmBudget)
	}
	if sm.coldBudget != 300 {
		t.Errorf("coldBudget: got %d, want 300", sm.coldBudget)
	}
	if sm.hotThreshold != defaultHotAccessThreshold {
		t.Errorf("hotThreshold: got %d, want %d", sm.hotThreshold, defaultHotAccessThreshold)
	}
	if sm.hotWindow != defaultHotAccessWindow {
		t.Errorf("hotWindow: got %v, want %v", sm.hotWindow, defaultHotAccessWindow)
	}
	if sm.hotDemoteAfter != defaultHotDemotionAfter {
		t.Errorf("hotDemoteAfter: got %v, want %v", sm.hotDemoteAfter, defaultHotDemotionAfter)
	}
	if sm.warmDemoteAfter != defaultWarmDemotionAfter {
		t.Errorf("warmDemoteAfter: got %v, want %v", sm.warmDemoteAfter, defaultWarmDemotionAfter)
	}

	// Entries map should be non-nil and empty.
	if sm.entries == nil {
		t.Fatal("entries map is nil")
	}
	if len(sm.entries) != 0 {
		t.Errorf("entries: got %d, want 0", len(sm.entries))
	}
}

// ---------------------------------------------------------------------------
// 2. Add entries to each tier (all new entries land in Cold)
// ---------------------------------------------------------------------------

func TestAddEntriesLandInCold(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)

	cases := []struct {
		name string
		id   string
	}{
		{"first entry", "a"},
		{"second entry", "b"},
		{"third entry", "c"},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			sm.Add(MemoryEntry{ID: tc.id, Content: "content-" + tc.id, TokenCount: 10})
		})
	}

	// Verify all entries are in Cold tier (use TierStats to avoid triggering promotion).
	stats := sm.TierStats()
	cold := stats[TierCold]
	if cold.Count != 3 {
		t.Errorf("cold count: got %d, want 3", cold.Count)
	}
	if cold.Tokens != 30 {
		t.Errorf("cold tokens: got %d, want 30", cold.Tokens)
	}

	// Additional verification with a fresh instance.
	sm2 := NewSpatialMemory(1000, 1000, 1000)
	sm2.Add(MemoryEntry{ID: "x", Content: "x", TokenCount: 5})
	sm2.Add(MemoryEntry{ID: "y", Content: "y", TokenCount: 5})

	stats2 := sm2.TierStats()
	cold2 := stats2[TierCold]
	if cold2.Count != 2 {
		t.Errorf("cold count: got %d, want 2", cold2.Count)
	}
	if cold2.Tokens != 10 {
		t.Errorf("cold tokens: got %d, want 10", cold2.Tokens)
	}
}

func TestAddOverwritesExistingEntry(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)

	sm.Add(MemoryEntry{ID: "dup", Content: "original", TokenCount: 10})
	sm.Add(MemoryEntry{ID: "dup", Content: "updated", TokenCount: 20})

	stats := sm.TierStats()
	cold := stats[TierCold]
	if cold.Count != 1 {
		t.Errorf("after overwrite: count %d, want 1", cold.Count)
	}
	if cold.Tokens != 20 {
		t.Errorf("after overwrite: tokens %d, want 20", cold.Tokens)
	}
}

func TestAddSetsCreatedAtIfZero(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)
	before := time.Now()
	sm.Add(MemoryEntry{ID: "ts", Content: "ts", TokenCount: 1})

	// Access to retrieve with timestamps set.
	e, _ := sm.Access("ts")
	if e.CreatedAt.Before(before) {
		t.Error("CreatedAt should be set to now if zero")
	}
	if e.LastAccessed.Before(before) {
		t.Error("LastAccessed should be set to now")
	}
}

func TestAddPreservesCreatedAtIfSet(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)
	past := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)
	sm.Add(MemoryEntry{ID: "ts", Content: "ts", TokenCount: 1, CreatedAt: past})

	e, _ := sm.Access("ts")
	if !e.CreatedAt.Equal(past) {
		t.Errorf("CreatedAt should be preserved: got %v, want %v", e.CreatedAt, past)
	}
}

// ---------------------------------------------------------------------------
// 3. Access triggers promotion
// ---------------------------------------------------------------------------

func TestAccessPromotesColdToWarm(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// First access: Cold -> Warm.
	e, ok := sm.Access("e1")
	if !ok {
		t.Fatal("entry not found")
	}
	if e.Tier != TierWarm {
		t.Errorf("after 1st access: tier %d, want TierWarm(%d)", e.Tier, TierWarm)
	}
	if e.AccessCount != 1 {
		t.Errorf("AccessCount: got %d, want 1", e.AccessCount)
	}
}

func TestAccessPromotesWarmToHotAfterThreshold(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(3),
		withHotWindow(1*time.Hour),
	)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// Access 1: Cold -> Warm.
	sm.Access("e1")
	// Access 2: still Warm (threshold=3, count=2).
	sm.Access("e1")
	// Access 3: Warm -> Hot (count=3 >= threshold).
	e, _ := sm.Access("e1")
	if e.Tier != TierHot {
		t.Errorf("after 3rd access: tier %d, want TierHot(%d)", e.Tier, TierHot)
	}
	if e.AccessCount != 3 {
		t.Errorf("AccessCount: got %d, want 3", e.AccessCount)
	}
}

func TestAccessReturnsFalseForMissingID(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)
	_, ok := sm.Access("nonexistent")
	if ok {
		t.Error("expected false for missing ID")
	}
}

func TestAccessDoesNotPromoteWarmToHotBeforeThreshold(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(5),
		withHotWindow(1*time.Hour),
	)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// Access 1: Cold -> Warm.
	sm.Access("e1")
	// Access 2-4: still Warm.
	for i := 2; i < 5; i++ {
		e, _ := sm.Access("e1")
		if e.Tier != TierWarm {
			t.Errorf("access %d: tier %d, want TierWarm(%d)", i, e.Tier, TierWarm)
		}
	}
}

func TestAccessPromotesWarmToHotRespectsWindow(t *testing.T) {
	// Use a very short window so accesses "expire".
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(2),
		withHotWindow(1*time.Nanosecond),
	)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// Access 1: Cold -> Warm.
	sm.Access("e1")
	// Sleep to exceed the window.
	time.Sleep(2 * time.Millisecond)
	// Access 2: still Warm because window expired (recentAccessCount returns 0).
	e, _ := sm.Access("e1")
	if e.Tier != TierWarm {
		t.Errorf("after expired window: tier %d, want TierWarm(%d)", e.Tier, TierWarm)
	}
}

// ---------------------------------------------------------------------------
// 4. Demotion after inactivity (Hot->Warm, Warm->Cold) via Compact
// ---------------------------------------------------------------------------

func TestCompactDemotesHotToWarmAfterInactivity(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(2),
		withHotWindow(1*time.Hour),
		withHotDemoteAfter(10*time.Millisecond),
	)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// Promote to Hot.
	sm.Access("e1") // Cold -> Warm
	sm.Access("e1") // Warm -> Hot

	e, _ := sm.Access("e1")
	if e.Tier != TierHot {
		t.Fatalf("should be Hot, got tier %d", e.Tier)
	}

	// Wait for idle threshold to pass.
	time.Sleep(15 * time.Millisecond)

	// Compact should demote Hot -> Warm.
	sm.Compact()

	stats := sm.TierStats()
	hot := stats[TierHot]
	warm := stats[TierWarm]
	if hot.Count != 0 {
		t.Errorf("hot count after compact: got %d, want 0", hot.Count)
	}
	if warm.Count != 1 {
		t.Errorf("warm count after compact: got %d, want 1", warm.Count)
	}
}

func TestCompactDemotesWarmToColdAfterInactivity(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withWarmDemoteAfter(10*time.Millisecond),
	)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// Access once: Cold -> Warm.
	sm.Access("e1")

	time.Sleep(15 * time.Millisecond)

	sm.Compact()

	stats := sm.TierStats()
	warm := stats[TierWarm]
	cold := stats[TierCold]
	if warm.Count != 0 {
		t.Errorf("warm count after compact: got %d, want 0", warm.Count)
	}
	if cold.Count != 1 {
		t.Errorf("cold count after compact: got %d, want 1", cold.Count)
	}
}

func TestCompactDoesNotDemoteRecentlyAccessed(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(2),
		withHotWindow(1*time.Hour),
		withHotDemoteAfter(1*time.Hour),
	)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	sm.Access("e1") // Cold -> Warm
	sm.Access("e1") // Warm -> Hot

	sm.Compact()

	stats := sm.TierStats()
	if stats[TierHot].Count != 1 {
		t.Errorf("recently accessed entry should stay Hot")
	}
}

// ---------------------------------------------------------------------------
// 5. Token budget enforcement
// ---------------------------------------------------------------------------

func TestBudgetEnforcementDemotesEntries(t *testing.T) {
	// Hot budget is 50 tokens. Each entry is 20 tokens.
	// 3 entries = 60 tokens > 50 budget. Oldest should be demoted.
	sm := NewSpatialMemory(50, 1000, 1000)

	// Add 3 entries and promote all to Hot.
	for i := 0; i < 3; i++ {
		id := fmt.Sprintf("e%d", i)
		sm.Add(MemoryEntry{ID: id, Content: id, TokenCount: 20})
	}

	// Promote all to Warm first.
	for i := 0; i < 3; i++ {
		sm.Access(fmt.Sprintf("e%d", i))
	}

	// Now access each enough times to promote to Hot.
	// Since threshold is 5, we need 4 more accesses per entry (total 5).
	for i := 0; i < 3; i++ {
		for j := 0; j < 4; j++ {
			sm.Access(fmt.Sprintf("e%d", i))
		}
	}

	stats := sm.TierStats()
	hot := stats[TierHot]
	// With budget 50 and 3 entries at 20 tokens each (60 total), one should be demoted.
	if hot.Count > 2 {
		t.Errorf("hot count: got %d, want at most 2 (budget=50, each=20)", hot.Count)
	}
}

func TestBudgetEnforcementDemotesOldestFirst(t *testing.T) {
	sm := NewSpatialMemory(30, 1000, 1000)

	// Add entry "old" first.
	sm.Add(MemoryEntry{ID: "old", Content: "old", TokenCount: 20})
	sm.Access("old") // Cold -> Warm
	// Promote to Hot: need 5 total accesses.
	for i := 0; i < 4; i++ {
		sm.Access("old")
	}

	// Small delay so "new" has a later LastAccessed.
	time.Sleep(2 * time.Millisecond)

	// Add entry "new".
	sm.Add(MemoryEntry{ID: "new", Content: "new", TokenCount: 20})
	sm.Access("new") // Cold -> Warm
	for i := 0; i < 4; i++ {
		sm.Access("new")
	}

	// Total hot tokens: 40 > budget 30. "old" should be demoted.
	stats := sm.TierStats()
	hot := stats[TierHot]

	if hot.Count == 2 {
		// Both are still Hot — budget not enforced. This shouldn't happen.
		// Let's check if "old" was demoted.
		e, _ := sm.Access("old")
		if e.Tier == TierHot {
			t.Error("expected 'old' to be demoted due to budget")
		}
	}

	// "new" should still be Hot (more recently accessed).
	e, ok := sm.Access("new")
	if !ok {
		t.Fatal("'new' entry missing")
	}
	if e.Tier != TierHot {
		t.Errorf("'new' tier: got %d, want TierHot(%d)", e.Tier, TierHot)
	}
}

func TestBudgetZeroDoesNotPanic(t *testing.T) {
	sm := NewSpatialMemory(0, 0, 0)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 100})
	// Should not panic.
	sm.Access("e1")
	sm.Compact()
}

func TestAddTriggersBudgetEnforcementOnColdTier(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 50)

	// Add 3 entries with 20 tokens each = 60 > cold budget 50.
	sm.Add(MemoryEntry{ID: "a", Content: "a", TokenCount: 20})
	sm.Add(MemoryEntry{ID: "b", Content: "b", TokenCount: 20})
	sm.Add(MemoryEntry{ID: "c", Content: "c", TokenCount: 20})

	// enforceBudget(TierCold) is called on Add. Cold budget is 50.
	// Total cold tokens = 60 > 50, so the oldest entry should be demoted.
	// But TierCold is the lowest tier — demotion has no lower tier.
	// Looking at enforceBudget: if tier < TierCold, demote. Cold is max, so nothing happens.
	// This is by design — cold has no lower tier.
	stats := sm.TierStats()
	cold := stats[TierCold]
	if cold.Count != 3 {
		t.Errorf("cold count: got %d, want 3 (cold is lowest tier, no demotion)", cold.Count)
	}
}

// ---------------------------------------------------------------------------
// 6. TierStats returns correct counts
// ---------------------------------------------------------------------------

func TestTierStatsEmptyMemory(t *testing.T) {
	sm := NewSpatialMemory(100, 100, 100)
	stats := sm.TierStats()
	if len(stats) != 0 {
		t.Errorf("empty memory stats: got %d tiers, want 0", len(stats))
	}
}

func TestTierStatsMixedTiers(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(2),
		withHotWindow(1*time.Hour),
	)

	// Add 3 entries (all start Cold).
	sm.Add(MemoryEntry{ID: "cold1", Content: "c", TokenCount: 10})
	sm.Add(MemoryEntry{ID: "warm1", Content: "c", TokenCount: 20})
	sm.Add(MemoryEntry{ID: "hot1", Content: "c", TokenCount: 30})

	// Promote warm1 to Warm.
	sm.Access("warm1")

	// Promote hot1 to Hot (access 2 times: Cold->Warm, then Warm->Hot).
	sm.Access("hot1")
	sm.Access("hot1")

	stats := sm.TierStats()

	if stats[TierCold].Count != 1 {
		t.Errorf("cold count: got %d, want 1", stats[TierCold].Count)
	}
	if stats[TierCold].Tokens != 10 {
		t.Errorf("cold tokens: got %d, want 10", stats[TierCold].Tokens)
	}
	if stats[TierWarm].Count != 1 {
		t.Errorf("warm count: got %d, want 1", stats[TierWarm].Count)
	}
	if stats[TierWarm].Tokens != 20 {
		t.Errorf("warm tokens: got %d, want 20", stats[TierWarm].Tokens)
	}
	if stats[TierHot].Count != 1 {
		t.Errorf("hot count: got %d, want 1", stats[TierHot].Count)
	}
	if stats[TierHot].Tokens != 30 {
		t.Errorf("hot tokens: got %d, want 30", stats[TierHot].Tokens)
	}
}

func TestTierStatsAfterRemove(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)
	sm.Add(MemoryEntry{ID: "a", Content: "a", TokenCount: 10})
	sm.Add(MemoryEntry{ID: "b", Content: "b", TokenCount: 20})

	sm.Remove("a")

	stats := sm.TierStats()
	if stats[TierCold].Count != 1 {
		t.Errorf("after remove: cold count %d, want 1", stats[TierCold].Count)
	}
	if stats[TierCold].Tokens != 20 {
		t.Errorf("after remove: cold tokens %d, want 20", stats[TierCold].Tokens)
	}
}

// ---------------------------------------------------------------------------
// 7. Compact demotes stale entries and enforces budgets
// ---------------------------------------------------------------------------

func TestCompactEmptyMemory(t *testing.T) {
	sm := NewSpatialMemory(100, 100, 100)
	// Should not panic.
	sm.Compact()
}

func TestCompactEnforcesBudgets(t *testing.T) {
	// Hot budget 30, entries are 20 each.
	sm := newTestSpatialMemory(
		30, 1000, 1000,
		withHotThreshold(2),
		withHotWindow(1*time.Hour),
		withHotDemoteAfter(1*time.Hour), // no idle demotion
	)

	sm.Add(MemoryEntry{ID: "a", Content: "a", TokenCount: 20})
	sm.Add(MemoryEntry{ID: "b", Content: "b", TokenCount: 20})

	// Promote both to Hot.
	sm.Access("a")
	sm.Access("a")
	sm.Access("b")
	sm.Access("b")

	// Total hot tokens = 40 > budget 30.
	sm.Compact()

	stats := sm.TierStats()
	if stats[TierHot].Count > 1 {
		t.Errorf("after compact: hot count %d, want at most 1 (budget=30, each=20)", stats[TierHot].Count)
	}
}

func TestCompactDemotesWarmAndEnforcesBudget(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 30, 1000,
		withWarmDemoteAfter(10*time.Millisecond),
	)

	sm.Add(MemoryEntry{ID: "a", Content: "a", TokenCount: 20})
	sm.Add(MemoryEntry{ID: "b", Content: "b", TokenCount: 20})

	// Promote both to Warm.
	sm.Access("a")
	sm.Access("b")

	// Wait for idle demotion.
	time.Sleep(15 * time.Millisecond)

	sm.Compact()

	stats := sm.TierStats()
	// Both should have been demoted to Cold due to inactivity.
	if stats[TierWarm].Count != 0 {
		t.Errorf("warm count after compact: got %d, want 0", stats[TierWarm].Count)
	}
	if stats[TierCold].Count != 2 {
		t.Errorf("cold count after compact: got %d, want 2", stats[TierCold].Count)
	}
}

// ---------------------------------------------------------------------------
// 8. Thread safety (concurrent access)
// ---------------------------------------------------------------------------

func TestConcurrentAddAccessCompact(t *testing.T) {
	sm := NewSpatialMemory(10000, 10000, 10000)

	const numGoroutines = 50
	const opsPerGoroutine = 100

	var wg sync.WaitGroup

	// Concurrent Adds.
	wg.Add(numGoroutines)
	for i := 0; i < numGoroutines; i++ {
		go func(id int) {
			defer wg.Done()
			for j := 0; j < opsPerGoroutine; j++ {
				sm.Add(MemoryEntry{
					ID:         fmt.Sprintf("g%d-e%d", id, j),
					Content:    "content",
					TokenCount: 1,
				})
			}
		}(i)
	}
	wg.Wait()

	stats := sm.TierStats()
	total := 0
	for _, s := range stats {
		total += s.Count
	}
	if total != numGoroutines*opsPerGoroutine {
		t.Errorf("total entries: got %d, want %d", total, numGoroutines*opsPerGoroutine)
	}

	// Concurrent Accesses.
	wg.Add(numGoroutines)
	for i := 0; i < numGoroutines; i++ {
		go func(id int) {
			defer wg.Done()
			for j := 0; j < opsPerGoroutine; j++ {
				sm.Access(fmt.Sprintf("g%d-e%d", id, j))
			}
		}(i)
	}
	wg.Wait()

	// Concurrent Compact + Access + Remove.
	wg.Add(numGoroutines * 3)
	for i := 0; i < numGoroutines; i++ {
		go func() {
			defer wg.Done()
			for j := 0; j < 10; j++ {
				sm.Compact()
			}
		}()
		go func(id int) {
			defer wg.Done()
			for j := 0; j < opsPerGoroutine; j++ {
				sm.Access(fmt.Sprintf("g%d-e%d", id, j))
			}
		}(i)
		go func(id int) {
			defer wg.Done()
			for j := 0; j < opsPerGoroutine; j++ {
				sm.Remove(fmt.Sprintf("g%d-e%d", id, j))
			}
		}(i)
	}
	wg.Wait()

	// Should not panic or deadlock. Verify memory is empty after all removes.
	stats = sm.TierStats()
	total = 0
	for _, s := range stats {
		total += s.Count
	}
	if total != 0 {
		t.Errorf("after removing all: got %d entries, want 0", total)
	}
}

func TestConcurrentTierStats(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)

	// Pre-populate.
	for i := 0; i < 100; i++ {
		sm.Add(MemoryEntry{ID: fmt.Sprintf("e%d", i), Content: "c", TokenCount: 1})
	}

	var wg sync.WaitGroup
	wg.Add(100)
	for i := 0; i < 100; i++ {
		go func() {
			defer wg.Done()
			stats := sm.TierStats()
			total := 0
			for _, s := range stats {
				total += s.Count
			}
			// total may vary due to concurrent adds; just verify no panic.
			_ = total
		}()
	}
	wg.Wait()
}

func TestConcurrentAddSameID(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)

	var wg sync.WaitGroup
	wg.Add(100)
	for i := 0; i < 100; i++ {
		go func(version int) {
			defer wg.Done()
			sm.Add(MemoryEntry{
				ID:         "shared",
				Content:    fmt.Sprintf("v%d", version),
				TokenCount: version,
			})
		}(i)
	}
	wg.Wait()

	// Exactly one entry should exist.
	stats := sm.TierStats()
	if stats[TierCold].Count != 1 {
		t.Errorf("concurrent add same ID: count %d, want 1", stats[TierCold].Count)
	}
}

// ---------------------------------------------------------------------------
// 9. Edge cases
// ---------------------------------------------------------------------------

func TestEmptyMemoryAccess(t *testing.T) {
	sm := NewSpatialMemory(100, 100, 100)
	_, ok := sm.Access("anything")
	if ok {
		t.Error("expected false when accessing empty memory")
	}
}

func TestEmptyMemoryRemove(t *testing.T) {
	sm := NewSpatialMemory(100, 100, 100)
	ok := sm.Remove("anything")
	if ok {
		t.Error("expected false when removing from empty memory")
	}
}

func TestEmptyMemoryCompact(t *testing.T) {
	sm := NewSpatialMemory(100, 100, 100)
	sm.Compact()
	// Should not panic.
}

func TestSingleEntryFullLifecycle(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(3),
		withHotWindow(1*time.Hour),
		withHotDemoteAfter(10*time.Millisecond),
		withWarmDemoteAfter(10*time.Millisecond),
	)

	sm.Add(MemoryEntry{ID: "solo", Content: "solo", TokenCount: 10})

	// Verify in Cold.
	stats := sm.TierStats()
	if stats[TierCold].Count != 1 {
		t.Fatalf("initial: cold count %d, want 1", stats[TierCold].Count)
	}

	// Access: Cold -> Warm.
	sm.Access("solo")
	stats = sm.TierStats()
	if stats[TierWarm].Count != 1 {
		t.Errorf("after 1 access: warm count %d, want 1", stats[TierWarm].Count)
	}

	// Access 2 more times: Warm -> Hot (threshold=3).
	sm.Access("solo")
	sm.Access("solo")
	stats = sm.TierStats()
	if stats[TierHot].Count != 1 {
		t.Errorf("after 3 accesses: hot count %d, want 1", stats[TierHot].Count)
	}

	// Wait for idle demotion.
	time.Sleep(15 * time.Millisecond)

	// Compact: Hot -> Warm.
	sm.Compact()
	stats = sm.TierStats()
	if stats[TierWarm].Count != 1 {
		t.Errorf("after first compact: warm count %d, want 1", stats[TierWarm].Count)
	}

	// Compact again: Warm -> Cold.
	sm.Compact()
	stats = sm.TierStats()
	if stats[TierCold].Count != 1 {
		t.Errorf("after second compact: cold count %d, want 1", stats[TierCold].Count)
	}
}

func TestAllEntriesSameTier(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)

	// All entries stay in Cold (no access).
	for i := 0; i < 10; i++ {
		sm.Add(MemoryEntry{ID: fmt.Sprintf("e%d", i), Content: "c", TokenCount: 5})
	}

	stats := sm.TierStats()
	if stats[TierCold].Count != 10 {
		t.Errorf("all cold: count %d, want 10", stats[TierCold].Count)
	}
	if stats[TierCold].Tokens != 50 {
		t.Errorf("all cold: tokens %d, want 50", stats[TierCold].Tokens)
	}
	if _, ok := stats[TierWarm]; ok {
		t.Error("warm tier should not exist when all entries are cold")
	}
	if _, ok := stats[TierHot]; ok {
		t.Error("hot tier should not exist when all entries are cold")
	}
}

func TestRemoveExistingAndNonExisting(t *testing.T) {
	sm := NewSpatialMemory(100, 100, 100)
	sm.Add(MemoryEntry{ID: "exists", Content: "c", TokenCount: 1})

	if !sm.Remove("exists") {
		t.Error("Remove('exists') should return true")
	}
	if sm.Remove("exists") {
		t.Error("Remove('exists') again should return false")
	}
	if sm.Remove("never") {
		t.Error("Remove('never') should return false")
	}
}

func TestAddWithZeroTokenCount(t *testing.T) {
	sm := NewSpatialMemory(100, 100, 100)
	sm.Add(MemoryEntry{ID: "zero", Content: "zero", TokenCount: 0})

	stats := sm.TierStats()
	if stats[TierCold].Count != 1 {
		t.Errorf("zero-token entry: count %d, want 1", stats[TierCold].Count)
	}
	if stats[TierCold].Tokens != 0 {
		t.Errorf("zero-token entry: tokens %d, want 0", stats[TierCold].Tokens)
	}
}

func TestLargeBudgetDoesNotDemote(t *testing.T) {
	sm := NewSpatialMemory(1000000, 1000000, 1000000)

	for i := 0; i < 100; i++ {
		sm.Add(MemoryEntry{ID: fmt.Sprintf("e%d", i), Content: "c", TokenCount: 100})
	}

	// Promote all to Hot.
	for i := 0; i < 100; i++ {
		for j := 0; j < 6; j++ { // 6 accesses > threshold(5)
			sm.Access(fmt.Sprintf("e%d", i))
		}
	}

	stats := sm.TierStats()
	if stats[TierHot].Count != 100 {
		t.Errorf("large budget: hot count %d, want 100", stats[TierHot].Count)
	}
}

func TestTableDrivenPromotionSequence(t *testing.T) {
	tests := []struct {
		name          string
		accesses      int
		hotThreshold  int
		wantFinalTier TierType
	}{
		{"0 accesses stays cold", 0, 5, TierCold},
		{"1 access promotes to warm", 1, 5, TierWarm},
		{"2 accesses stays warm", 2, 5, TierWarm},
		{"4 accesses stays warm", 4, 5, TierWarm},
		{"5 accesses promotes to hot", 5, 5, TierHot},
		{"10 accesses stays hot", 10, 5, TierHot},
		{"1 access with threshold 1", 1, 1, TierWarm}, // Cold->Warm on first, then threshold check
		{"2 accesses with threshold 1", 2, 1, TierHot},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			sm := newTestSpatialMemory(
				1000, 1000, 1000,
				withHotThreshold(tc.hotThreshold),
				withHotWindow(1*time.Hour),
			)
			sm.Add(MemoryEntry{ID: "e", Content: "c", TokenCount: 1})

			var entry MemoryEntry
			for i := 0; i < tc.accesses; i++ {
				e, ok := sm.Access("e")
				if !ok {
					t.Fatal("entry not found")
				}
				entry = e
			}

			if tc.accesses == 0 {
				// No accesses; check via TierStats.
				stats := sm.TierStats()
				if stats[TierCold].Count != 1 {
					t.Errorf("0 accesses: expected 1 cold entry")
				}
				return
			}

			if entry.Tier != tc.wantFinalTier {
				t.Errorf("after %d accesses with threshold %d: tier %d, want %d",
					tc.accesses, tc.hotThreshold, entry.Tier, tc.wantFinalTier)
			}
		})
	}
}

func TestTableDrivenBudgetEnforcement(t *testing.T) {
	tests := []struct {
		name       string
		hotBudget  int
		numEntries int
		tokenEach  int
		wantMaxHot int
	}{
		{"under budget", 1000, 3, 100, 3},
		{"exactly at budget", 300, 3, 100, 3},
		{"over budget by 1 entry", 250, 3, 100, 2},
		{"zero budget means no limit", 0, 3, 100, 3},
		{"single entry over budget", 50, 1, 100, 0},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			sm := newTestSpatialMemory(
				tc.hotBudget, 10000, 10000,
				withHotThreshold(2),
				withHotWindow(1*time.Hour),
			)

			for i := 0; i < tc.numEntries; i++ {
				id := fmt.Sprintf("e%d", i)
				sm.Add(MemoryEntry{ID: id, Content: id, TokenCount: tc.tokenEach})
			}

			// Promote all to Hot.
			for i := 0; i < tc.numEntries; i++ {
				for j := 0; j < 5; j++ {
					sm.Access(fmt.Sprintf("e%d", i))
				}
			}

			stats := sm.TierStats()
			hotCount := stats[TierHot].Count
			if hotCount > tc.wantMaxHot {
				t.Errorf("hot count: got %d, want at most %d", hotCount, tc.wantMaxHot)
			}
		})
	}
}

func TestAccessUpdatesLastAccessedTime(t *testing.T) {
	sm := NewSpatialMemory(1000, 1000, 1000)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// First access.
	e1, _ := sm.Access("e1")
	time.Sleep(2 * time.Millisecond)

	// Second access.
	e2, _ := sm.Access("e1")

	if !e2.LastAccessed.After(e1.LastAccessed) {
		t.Error("LastAccessed should advance with each access")
	}
}

func TestPromoteAlreadyHotIsNoop(t *testing.T) {
	sm := newTestSpatialMemory(
		1000, 1000, 1000,
		withHotThreshold(2),
		withHotWindow(1*time.Hour),
	)
	sm.Add(MemoryEntry{ID: "e1", Content: "c", TokenCount: 5})

	// Promote to Hot.
	sm.Access("e1")
	sm.Access("e1")

	// Additional accesses should keep it Hot.
	for i := 0; i < 10; i++ {
		e, _ := sm.Access("e1")
		if e.Tier != TierHot {
			t.Errorf("access %d: tier %d, want TierHot(%d)", i+3, e.Tier, TierHot)
		}
	}
}