optimizer_test.go

package tok

import (
	"strings"
	"testing"
)

func sampleBlocks() []ContentBlock {
	return []ContentBlock{
		{ID: "sys1", Content: "You are a helpful assistant.", Tokens: 500, Priority: 1.0, Category: "system", Compressible: false},
		{ID: "mem1", Content: "User prefers concise answers. User lives in NYC. User works in finance.", Tokens: 800, Priority: 0.8, Category: "memory", Compressible: true},
		{ID: "conv1", Content: "What is the capital of France? The capital of France is Paris.", Tokens: 1200, Priority: 0.6, Category: "conversation", Compressible: true},
		{ID: "conv2", Content: "Tell me about quantum computing. Quantum computing uses qubits...", Tokens: 2000, Priority: 0.5, Category: "conversation", Compressible: true},
		{ID: "tool1", Content: "File contents: package main\nfunc main() {\n\tfmt.Println(\"hello\")\n}", Tokens: 1500, Priority: 0.7, Category: "tool_output", Compressible: true},
		{ID: "ctx1", Content: "Project uses Go 1.21. Main package in cmd/. Tests in *_test.go files.", Tokens: 600, Priority: 0.9, Category: "context", Compressible: true},
	}
}

func TestNewContextOptimizer(t *testing.T) {
	opt := NewContextOptimizer(8000)
	if opt.Budget != 8000 {
		t.Errorf("expected budget 8000, got %d", opt.Budget)
	}
	if opt.Strategy != "priority" {
		t.Errorf("expected strategy 'priority', got %q", opt.Strategy)
	}
}

func TestOptimize_Priority(t *testing.T) {
	opt := NewContextOptimizer(4000)
	opt.Strategy = "priority"
	blocks := sampleBlocks()

	result := opt.Optimize(blocks)
	if result == nil {
		t.Fatal("result should not be nil")
	}
	if result.TotalTokens > 4000 {
		t.Errorf("total tokens %d exceeds budget 4000", result.TotalTokens)
	}
	if result.BudgetUsed < 0 || result.BudgetUsed > 1.0 {
		t.Errorf("budget used %f should be between 0 and 1", result.BudgetUsed)
	}
	// High priority blocks should be kept
	found := false
	for _, b := range result.Kept {
		if b.ID == "sys1" {
			found = true
			break
		}
	}
	if !found {
		t.Error("highest priority block sys1 should be kept")
	}
}

func TestOptimize_Greedy(t *testing.T) {
	opt := NewContextOptimizer(3000)
	opt.Strategy = "greedy"
	blocks := sampleBlocks()

	result := opt.Optimize(blocks)
	if result == nil {
		t.Fatal("result should not be nil")
	}
	if result.TotalTokens > 3000 {
		t.Errorf("total tokens %d exceeds budget 3000", result.TotalTokens)
	}
	if len(result.Dropped) == 0 {
		t.Error("with 3000 budget and ~6600 total tokens, some blocks should be dropped")
	}
}

func TestOptimize_Balanced(t *testing.T) {
	opt := NewContextOptimizer(4000)
	opt.Strategy = "balanced"
	blocks := sampleBlocks()

	result := opt.Optimize(blocks)
	if result == nil {
		t.Fatal("result should not be nil")
	}
	if result.TotalTokens > 4000 {
		t.Errorf("total tokens %d exceeds budget 4000", result.TotalTokens)
	}

	// Balanced should attempt to represent multiple categories
	categoriesSeen := map[string]bool{}
	for _, b := range result.Kept {
		categoriesSeen[b.Category] = true
	}
	for _, b := range result.Compressed {
		categoriesSeen[b.Category] = true
	}
	if len(categoriesSeen) < 3 {
		t.Errorf("balanced should include at least 3 categories, got %d", len(categoriesSeen))
	}
}

func TestGreedyOptimize(t *testing.T) {
	blocks := []ContentBlock{
		{ID: "a", Tokens: 100, Priority: 0.9},
		{ID: "b", Tokens: 200, Priority: 0.5},
		{ID: "c", Tokens: 150, Priority: 0.7},
	}

	result := GreedyOptimize(blocks, 300)
	if result.TotalTokens > 300 {
		t.Errorf("greedy exceeded budget: %d > 300", result.TotalTokens)
	}
	// Should include a (100) and c (150) = 250 tokens
	if len(result.Kept) != 2 {
		t.Errorf("expected 2 kept blocks, got %d", len(result.Kept))
	}
	if len(result.Dropped) != 1 {
		t.Errorf("expected 1 dropped block, got %d", len(result.Dropped))
	}
}

func TestPriorityOptimize_CompressBeforeDrop(t *testing.T) {
	blocks := []ContentBlock{
		{ID: "a", Content: "Important system prompt content here.", Tokens: 500, Priority: 1.0, Compressible: false},
		{ID: "b", Content: "This is actually just a really very basic conversation about things that are somewhat relevant to the overall context of what we discussed.", Tokens: 600, Priority: 0.8, Compressible: true},
		{ID: "c", Content: "Low priority filler content.", Tokens: 400, Priority: 0.3, Compressible: true},
	}

	result := PriorityOptimize(blocks, 900)
	// Block a (500) fits, block b (600) won't fit fully so should compress
	if len(result.Compressed) == 0 {
		t.Error("expected at least one compressed block")
	}
	if result.TotalTokens > 900 {
		t.Errorf("exceeded budget: %d > 900", result.TotalTokens)
	}
}

func TestBalancedOptimize_MinRepresentation(t *testing.T) {
	blocks := []ContentBlock{
		{ID: "s1", Tokens: 100, Priority: 0.5, Category: "system"},
		{ID: "m1", Tokens: 100, Priority: 0.5, Category: "memory"},
		{ID: "c1", Tokens: 100, Priority: 0.5, Category: "conversation"},
		{ID: "t1", Tokens: 100, Priority: 0.5, Category: "tool_output"},
		{ID: "x1", Tokens: 100, Priority: 0.5, Category: "context"},
	}

	result := BalancedOptimize(blocks, 500)
	if len(result.Kept) != 5 {
		t.Errorf("with budget=500 and 5 blocks of 100 each, expected all 5 kept, got %d", len(result.Kept))
	}
}

func TestCompressBlock(t *testing.T) {
	block := ContentBlock{
		ID:      "test",
		Content: "This is a test of the compression system that should reduce content.",
		Tokens:  100,
	}

	compressed := CompressBlock(block, 50)
	if compressed.Tokens != 50 {
		t.Errorf("expected 50 tokens, got %d", compressed.Tokens)
	}
	if compressed.CompressedContent == "" {
		t.Error("compressed content should not be empty")
	}
}

func TestCompressBlock_NoCompression(t *testing.T) {
	block := ContentBlock{
		ID:      "test",
		Content: "Short content.",
		Tokens:  100,
	}

	// Target >= current tokens should not compress
	compressed := CompressBlock(block, 200)
	if compressed.CompressedContent != "" {
		t.Error("should not compress when target >= current tokens")
	}
	if compressed.Tokens != 100 {
		t.Errorf("tokens should remain 100, got %d", compressed.Tokens)
	}
}

func TestEstimateCompression_LightRatio(t *testing.T) {
	content := "This is actually just a really very basic test of light compression."
	result := EstimateCompression(content, 0.85)
	if strings.Contains(result, " actually") {
		t.Error("light compression should remove filler words like 'actually'")
	}
	if strings.Contains(result, " really") {
		t.Error("light compression should remove filler words like 'really'")
	}
}

func TestEstimateCompression_MediumRatio(t *testing.T) {
	content := "Start of the content. Middle section with lots of details. End of the content section."
	result := EstimateCompression(content, 0.6)
	if !strings.Contains(result, "[...]") {
		t.Error("medium compression should contain [...] marker")
	}
	if len(result) >= len(content) {
		t.Error("medium compression should reduce length")
	}
}

func TestEstimateCompression_HeavyRatio(t *testing.T) {
	content := "This is the first sentence. Then there is a lot more content that follows with many details and explanations that go on and on."
	result := EstimateCompression(content, 0.2)
	if len(result) > 80 {
		t.Errorf("heavy compression should produce short output, got length %d", len(result))
	}
}

func TestEstimateCompression_EdgeCases(t *testing.T) {
	// ratio >= 1.0 should return content unchanged
	content := "unchanged content"
	if EstimateCompression(content, 1.0) != content {
		t.Error("ratio 1.0 should return unchanged content")
	}
	if EstimateCompression(content, 1.5) != content {
		t.Error("ratio > 1.0 should return unchanged content")
	}
	// ratio <= 0 should return empty
	if EstimateCompression(content, 0) != "" {
		t.Error("ratio 0 should return empty string")
	}
	if EstimateCompression(content, -0.5) != "" {
		t.Error("negative ratio should return empty string")
	}
}

func TestFormatResult(t *testing.T) {
	result := &OptimizationResult{
		Kept:        []ContentBlock{{Tokens: 3000}, {Tokens: 2200}},
		Compressed:  []ContentBlock{{Tokens: 900}, {Tokens: 900}},
		Dropped:     []ContentBlock{{Tokens: 600}},
		TotalTokens: 7000,
		BudgetUsed:  0.875,
		Savings:     2000,
	}

	formatted := FormatResult(result)
	if !strings.Contains(formatted, "Context Optimization:") {
		t.Error("should contain header")
	}
	if !strings.Contains(formatted, "8,000") {
		t.Errorf("should contain budget 8,000, got:\n%s", formatted)
	}
	if !strings.Contains(formatted, "87.5%") {
		t.Error("should contain utilization percentage")
	}
	if !strings.Contains(formatted, "Kept (2 blocks)") {
		t.Error("should show kept block count")
	}
	if !strings.Contains(formatted, "Compressed (2 blocks)") {
		t.Error("should show compressed block count")
	}
	if !strings.Contains(formatted, "Dropped (1 blocks)") {
		t.Error("should show dropped block count")
	}
}

func TestFormatResult_Nil(t *testing.T) {
	if FormatResult(nil) != "" {
		t.Error("nil result should return empty string")
	}
}

func TestSuggestBudget(t *testing.T) {
	blocks := sampleBlocks()
	suggested := SuggestBudget(blocks)

	totalTokens := 0
	for _, b := range blocks {
		totalTokens += b.Tokens
	}

	if suggested <= 0 {
		t.Error("suggested budget should be positive")
	}
	if suggested > totalTokens {
		t.Errorf("suggested %d should not exceed total %d", suggested, totalTokens)
	}
	// Should be at least 50% of total
	if suggested < totalTokens/2 {
		t.Errorf("suggested %d should be at least %d (50%% of total)", suggested, totalTokens/2)
	}
}

func TestSuggestBudget_Empty(t *testing.T) {
	if SuggestBudget(nil) != 0 {
		t.Error("empty blocks should suggest 0 budget")
	}
	if SuggestBudget([]ContentBlock{}) != 0 {
		t.Error("empty blocks should suggest 0 budget")
	}
}

func TestSuggestBudget_AllHighPriority(t *testing.T) {
	blocks := []ContentBlock{
		{Tokens: 500, Priority: 0.9},
		{Tokens: 500, Priority: 0.8},
		{Tokens: 500, Priority: 0.7},
	}
	suggested := SuggestBudget(blocks)
	// All are high priority so should suggest close to total
	if suggested < 1500 {
		t.Errorf("all high priority: suggested %d should be at least 1500", suggested)
	}
}

func TestFmtTokenCount(t *testing.T) {
	tests := []struct {
		input    int
		expected string
	}{
		{0, "0"},
		{999, "999"},
		{1000, "1,000"},
		{8000, "8,000"},
		{12345, "12,345"},
		{1000000, "1,000,000"},
	}
	for _, tt := range tests {
		got := fmtTokenCount(tt.input)
		if got != tt.expected {
			t.Errorf("fmtTokenCount(%d) = %q, want %q", tt.input, got, tt.expected)
		}
	}
}

func TestOptimizer_ConcurrentSafety(t *testing.T) {
	opt := NewContextOptimizer(5000)
	blocks := sampleBlocks()

	done := make(chan bool, 10)
	for i := 0; i < 10; i++ {
		go func() {
			_ = opt.Optimize(blocks)
			done <- true
		}()
	}
	for i := 0; i < 10; i++ {
		<-done
	}
}

func TestOptimizationResult_Savings(t *testing.T) {
	blocks := []ContentBlock{
		{ID: "a", Tokens: 500, Priority: 1.0, Compressible: false},
		{ID: "b", Tokens: 500, Priority: 0.5, Compressible: false},
	}

	result := GreedyOptimize(blocks, 600)
	if result.Savings < 0 {
		t.Error("savings should not be negative")
	}
	// Total is 1000, budget is 600, so we can only fit 500, savings = 500
	if result.Savings != 500 {
		t.Errorf("expected savings of 500, got %d", result.Savings)
	}
}

func TestPriorityOptimize_AllFit(t *testing.T) {
	blocks := []ContentBlock{
		{ID: "a", Tokens: 100, Priority: 1.0},
		{ID: "b", Tokens: 100, Priority: 0.5},
	}

	result := PriorityOptimize(blocks, 1000)
	if len(result.Kept) != 2 {
		t.Errorf("all blocks should fit, got %d kept", len(result.Kept))
	}
	if len(result.Dropped) != 0 {
		t.Error("no blocks should be dropped when all fit")
	}
	if len(result.Compressed) != 0 {
		t.Error("no blocks should be compressed when all fit")
	}
	if result.TotalTokens != 200 {
		t.Errorf("total tokens should be 200, got %d", result.TotalTokens)
	}
}

func TestGreedyOptimize_EmptyBlocks(t *testing.T) {
	result := GreedyOptimize(nil, 1000)
	if result.TotalTokens != 0 {
		t.Error("empty blocks should yield 0 tokens")
	}
	if len(result.Kept) != 0 {
		t.Error("empty blocks should yield no kept blocks")
	}
}