optimizer.go

package tok

import (
	"fmt"
	"sort"
	"strings"
	"sync"
)

// ContextOptimizer maximizes information density within a token budget.
type ContextOptimizer struct {
	Budget   int
	Strategy string // "greedy", "balanced", "priority"
	mu       sync.RWMutex
}

// ContentBlock represents a unit of content that can be optimized.
type ContentBlock struct {
	ID                string
	Content           string
	Tokens            int
	Priority          float64
	Category          string // "system", "memory", "conversation", "tool_output", "context"
	Compressible      bool
	CompressedContent string
}

// OptimizationResult holds the outcome of an optimization pass.
type OptimizationResult struct {
	Kept        []ContentBlock
	Dropped     []ContentBlock
	Compressed  []ContentBlock
	TotalTokens int
	BudgetUsed  float64
	Savings     int
}

// NewContextOptimizer creates an optimizer with the given token budget.
func NewContextOptimizer(budget int) *ContextOptimizer {
	return &ContextOptimizer{
		Budget:   budget,
		Strategy: "priority",
	}
}

// Optimize selects the best strategy and optimizes content blocks within budget.
func (co *ContextOptimizer) Optimize(blocks []ContentBlock) *OptimizationResult {
	co.mu.RLock()
	strategy := co.Strategy
	budget := co.Budget
	co.mu.RUnlock()

	switch strategy {
	case "greedy":
		return GreedyOptimize(blocks, budget)
	case "balanced":
		return BalancedOptimize(blocks, budget)
	default:
		return PriorityOptimize(blocks, budget)
	}
}

// CompressBlock applies compression to a block proportional to how much we need to save.
// targetTokens is the desired token count after compression.
func CompressBlock(block ContentBlock, targetTokens int) ContentBlock {
	if block.Tokens == 0 || targetTokens >= block.Tokens {
		return block
	}

	ratio := float64(targetTokens) / float64(block.Tokens)
	compressed := EstimateCompression(block.Content, ratio)

	result := block
	result.CompressedContent = compressed
	result.Tokens = targetTokens
	return result
}

// GreedyOptimize takes highest priority blocks until budget is full.
func GreedyOptimize(blocks []ContentBlock, budget int) *OptimizationResult {
	sorted := make([]ContentBlock, len(blocks))
	copy(sorted, blocks)
	sort.Slice(sorted, func(i, j int) bool {
		return sorted[i].Priority > sorted[j].Priority
	})

	result := &OptimizationResult{}
	used := 0
	originalTotal := 0
	for _, b := range sorted {
		originalTotal += b.Tokens
	}

	for _, block := range sorted {
		if used+block.Tokens <= budget {
			result.Kept = append(result.Kept, block)
			used += block.Tokens
		} else {
			result.Dropped = append(result.Dropped, block)
		}
	}

	result.TotalTokens = used
	if budget > 0 {
		result.BudgetUsed = float64(used) / float64(budget)
	}
	result.Savings = originalTotal - used
	return result
}

// BalancedOptimize ensures each category gets minimum representation, then fills by priority.
func BalancedOptimize(blocks []ContentBlock, budget int) *OptimizationResult {
	categories := map[string][]ContentBlock{}
	for _, b := range blocks {
		categories[b.Category] = append(categories[b.Category], b)
	}

	// Sort each category by priority descending
	for cat := range categories {
		sort.Slice(categories[cat], func(i, j int) bool {
			return categories[cat][i].Priority > categories[cat][j].Priority
		})
	}

	result := &OptimizationResult{}
	used := 0
	originalTotal := 0
	for _, b := range blocks {
		originalTotal += b.Tokens
	}

	included := map[string]bool{}

	// Phase 1: guarantee at least one block per category (highest priority in each)
	for cat, catBlocks := range categories {
		if len(catBlocks) == 0 {
			continue
		}
		best := catBlocks[0]
		if used+best.Tokens <= budget {
			result.Kept = append(result.Kept, best)
			used += best.Tokens
			included[best.ID] = true
		} else if best.Compressible {
			remaining := budget - used
			if remaining > 0 {
				compressed := CompressBlock(best, remaining)
				result.Compressed = append(result.Compressed, compressed)
				used += compressed.Tokens
				included[best.ID] = true
			}
		}
		_ = cat
	}

	// Phase 2: fill remaining budget by priority across all blocks
	allSorted := make([]ContentBlock, len(blocks))
	copy(allSorted, blocks)
	sort.Slice(allSorted, func(i, j int) bool {
		return allSorted[i].Priority > allSorted[j].Priority
	})

	for _, block := range allSorted {
		if included[block.ID] {
			continue
		}
		if used+block.Tokens <= budget {
			result.Kept = append(result.Kept, block)
			used += block.Tokens
			included[block.ID] = true
		} else if block.Compressible {
			remaining := budget - used
			if remaining > 0 {
				compressed := CompressBlock(block, remaining)
				result.Compressed = append(result.Compressed, compressed)
				used += compressed.Tokens
				included[block.ID] = true
			}
		}
	}

	// Anything not included is dropped
	for _, block := range blocks {
		if !included[block.ID] {
			result.Dropped = append(result.Dropped, block)
		}
	}

	result.TotalTokens = used
	if budget > 0 {
		result.BudgetUsed = float64(used) / float64(budget)
	}
	result.Savings = originalTotal - used
	return result
}

// PriorityOptimize uses strict priority ordering, compressing before dropping.
func PriorityOptimize(blocks []ContentBlock, budget int) *OptimizationResult {
	sorted := make([]ContentBlock, len(blocks))
	copy(sorted, blocks)
	sort.Slice(sorted, func(i, j int) bool {
		return sorted[i].Priority > sorted[j].Priority
	})

	result := &OptimizationResult{}
	used := 0
	originalTotal := 0
	for _, b := range sorted {
		originalTotal += b.Tokens
	}

	for _, block := range sorted {
		switch {
		case used+block.Tokens <= budget:
			result.Kept = append(result.Kept, block)
			used += block.Tokens
		case block.Compressible:
			remaining := budget - used
			if remaining > 0 {
				// Determine compression level based on how much we need to save
				compressed := CompressBlock(block, remaining)
				result.Compressed = append(result.Compressed, compressed)
				used += compressed.Tokens
			} else {
				result.Dropped = append(result.Dropped, block)
			}
		default:
			result.Dropped = append(result.Dropped, block)
		}
	}

	result.TotalTokens = used
	if budget > 0 {
		result.BudgetUsed = float64(used) / float64(budget)
	}
	result.Savings = originalTotal - used
	return result
}

// EstimateCompression applies ratio-based truncation intelligently, keeping first/last and dropping middle.
func EstimateCompression(content string, ratio float64) string {
	if ratio >= 1.0 {
		return content
	}
	if ratio <= 0 {
		return ""
	}

	// Determine compression level
	if ratio >= 0.8 {
		// Light compression: remove filler words and extra whitespace
		return lightCompress(content)
	} else if ratio >= 0.5 {
		// Medium compression: keep first and last portions
		return mediumCompress(content, ratio)
	}
	// Heavy compression: one-line summary (first sentence or truncated)
	return heavyCompress(content)
}

// lightCompress removes filler words and collapses whitespace.
func lightCompress(content string) string {
	fillers := []string{
		" actually", " basically", " just", " really", " very",
		" simply", " quite", " rather", " somewhat",
	}
	result := content
	for _, filler := range fillers {
		result = strings.ReplaceAll(result, filler, "")
	}
	// Collapse multiple spaces
	for strings.Contains(result, "  ") {
		result = strings.ReplaceAll(result, "  ", " ")
	}
	// Collapse multiple newlines
	for strings.Contains(result, "\n\n\n") {
		result = strings.ReplaceAll(result, "\n\n\n", "\n\n")
	}
	result = strings.TrimSpace(result)
	return result
}

// mediumCompress keeps the first and last portions of content.
func mediumCompress(content string, ratio float64) string {
	targetLen := int(float64(len(content)) * ratio)
	if targetLen >= len(content) {
		return content
	}
	if targetLen <= 0 {
		return ""
	}

	// Keep first 60% of target and last 40% of target
	firstPart := int(float64(targetLen) * 0.6)
	lastPart := targetLen - firstPart

	if firstPart+lastPart > len(content) {
		return content
	}

	first := content[:firstPart]
	last := content[len(content)-lastPart:]
	return first + " [...] " + last
}

// heavyCompress reduces content to a one-line summary.
func heavyCompress(content string) string {
	// Take the first sentence or first 80 characters
	lines := strings.SplitN(content, "\n", 2)
	firstLine := strings.TrimSpace(lines[0])

	// Find first sentence boundary
	for _, sep := range []string{". ", "! ", "? "} {
		if idx := strings.Index(firstLine, sep); idx > 0 && idx < 80 {
			return firstLine[:idx+1]
		}
	}

	if len(firstLine) > 80 {
		return firstLine[:77] + "..."
	}
	return firstLine
}

// FormatResult produces a human-readable summary of the optimization result.
func FormatResult(result *OptimizationResult) string {
	if result == nil {
		return ""
	}

	// Calculate budget from usage
	budget := 0
	if result.BudgetUsed > 0 {
		budget = int(float64(result.TotalTokens) / result.BudgetUsed)
	}

	var sb strings.Builder
	sb.WriteString("Context Optimization:\n")
	sb.WriteString(fmt.Sprintf("Budget: %s tokens\n", fmtTokenCount(budget)))
	sb.WriteString("─────────────────────────\n")

	keptTokens := 0
	for _, b := range result.Kept {
		keptTokens += b.Tokens
	}
	sb.WriteString(fmt.Sprintf("Kept (%d blocks):      %s tokens\n", len(result.Kept), fmtTokenCount(keptTokens)))

	compressedTokens := 0
	savedFromCompression := 0
	for _, b := range result.Compressed {
		compressedTokens += b.Tokens
	}
	savedFromCompression = result.Savings - droppedTokens(result)
	if savedFromCompression < 0 {
		savedFromCompression = 0
	}
	sb.WriteString(fmt.Sprintf("Compressed (%d blocks): %s tokens (saved %s)\n",
		len(result.Compressed), fmtTokenCount(compressedTokens), fmtTokenCount(savedFromCompression)))

	droppedTok := droppedTokens(result)
	sb.WriteString(fmt.Sprintf("Dropped (%d blocks):    -%s tokens\n", len(result.Dropped), fmtTokenCount(droppedTok)))

	sb.WriteString("─────────────────────────\n")

	pct := result.BudgetUsed * 100
	sb.WriteString(fmt.Sprintf("Total: %s/%s (%.1f%% utilized)\n",
		fmtTokenCount(result.TotalTokens), fmtTokenCount(budget), pct))
	sb.WriteString(fmt.Sprintf("Savings: %s tokens from compression + drops\n", fmtTokenCount(result.Savings)))

	return sb.String()
}

// droppedTokens calculates total tokens in dropped blocks.
func droppedTokens(result *OptimizationResult) int {
	total := 0
	for _, b := range result.Dropped {
		total += b.Tokens
	}
	return total
}

// SuggestBudget recommends an optimal budget based on content blocks.
func SuggestBudget(blocks []ContentBlock) int {
	if len(blocks) == 0 {
		return 0
	}

	totalTokens := 0
	highPriorityTokens := 0
	for _, b := range blocks {
		totalTokens += b.Tokens
		if b.Priority >= 0.7 {
			highPriorityTokens += b.Tokens
		}
	}

	// Suggested budget: enough for all high priority content + 30% buffer for lower priority
	lowPriorityTokens := totalTokens - highPriorityTokens
	suggested := highPriorityTokens + int(float64(lowPriorityTokens)*0.3)

	// Ensure minimum 50% of total if suggestion is too low
	minBudget := totalTokens / 2
	if suggested < minBudget {
		suggested = minBudget
	}

	// Cap at total tokens (no need to suggest more than needed)
	if suggested > totalTokens {
		suggested = totalTokens
	}

	return suggested
}

// fmtTokenCount adds comma separators to integers for display.
func fmtTokenCount(n int) string {
	if n < 0 {
		return "-" + fmtTokenCount(-n)
	}
	s := fmt.Sprintf("%d", n)
	if len(s) <= 3 {
		return s
	}

	var result strings.Builder
	remainder := len(s) % 3
	if remainder > 0 {
		result.WriteString(s[:remainder])
	}
	for i := remainder; i < len(s); i += 3 {
		if result.Len() > 0 {
			result.WriteByte(',')
		}
		result.WriteString(s[i : i+3])
	}
	return result.String()
}