optimizer.go

package tok

import (
	"fmt"
	"regexp"
	"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 _, 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
			}
		}
	}

	// 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.
// Uses word-boundary matching to avoid destroying words that merely contain
// a filler as a substring (e.g., "actually" must not be touched by the
// " all" filler, and "justified" must survive the " just" filler).
func lightCompress(content string) string {
	// Filler words to remove. Each is preceded by a space so we can match
	// " <word>" at a word boundary. The regex ensures the character after
	// the filler is not a word character (letter/digit/underscore).
	fillers := []string{
		"actually", "basically", "just", "really", "very",
		"simply", "quite", "rather", "somewhat",
	}

	result := content
	for _, filler := range fillers {
		// Match the filler as a whole word preceded by whitespace.
		// The negative look-ahead (?!\w) prevents matching when the filler
		// is a prefix of a longer word (e.g., "justified").
		re := regexp.MustCompile(`(\s)` + regexp.QuoteMeta(filler) + `(?!\w)`)
		result = string(re.ReplaceAll([]byte(result), []byte("$1")))
	}

	// Collapse multiple spaces and newlines in a single pass
	var b strings.Builder
	b.Grow(len(result))
	prevSpace := false
	prevNewline := 0
	for _, r := range result {
		switch r {
		case '\n':
			prevSpace = false
			prevNewline++
			if prevNewline <= 2 {
				b.WriteRune(r)
			}
		case ' ', '\t':
			prevNewline = 0
			if !prevSpace {
				b.WriteRune(' ')
				prevSpace = true
			}
		default:
			prevNewline = 0
			prevSpace = false
			b.WriteRune(r)
		}
	}
	result = strings.TrimSpace(b.String())
	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
	}

	// Truncate at word boundaries (space or newline) to avoid mid-word cuts.
	first := truncateAtWordBoundary(content, firstPart, true)
	lastStart := len(content) - lastPart
	last := truncateAtWordBoundary(content, lastStart, false)
	return first + " [...] " + last
}

// truncateAtWordBoundary returns a substring adjusted to the nearest word
// boundary. If forward is true it snaps to the boundary at or before pos;
// if false it snaps to the boundary at or after pos.
func truncateAtWordBoundary(s string, pos int, forward bool) string {
	if pos <= 0 {
		return ""
	}
	if pos >= len(s) {
		return s
	}
	if forward {
		// Walk backward to find a space or newline.
		for i := pos; i > 0; i-- {
			if s[i] == ' ' || s[i] == '\n' {
				return s[:i]
			}
		}
		return s[:pos]
	}
	// Walk forward to find a space or newline.
	for i := pos; i < len(s); i++ {
		if s[i] == ' ' || s[i] == '\n' {
			return s[i+1:]
		}
	}
	return s[pos:]
}

// 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", formatNumber(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), formatNumber(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), formatNumber(compressedTokens), formatNumber(savedFromCompression)))

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

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

	pct := result.BudgetUsed * 100
	sb.WriteString(fmt.Sprintf("Total: %s/%s (%.1f%% utilized)\n",
		formatNumber(result.TotalTokens), formatNumber(budget), pct))
	sb.WriteString(fmt.Sprintf("Savings: %s tokens from compression + drops\n", formatNumber(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
}

// formatNumber is defined in format.go and shared across the package.