internal/memory/context_window.go

// Package memory provides context window management with smart truncation.
package memory

import (
	"fmt"
	"strings"
	"time"
)

const (
	// MaxContextTokens is the maximum context window size
	MaxContextTokens = 8192

	// SummaryThreshold is when to start summarizing older entries
	SummaryThreshold = 4000

	// DefaultMaxTokens for context building
	DefaultMaxTokens = 4096

	// TokensPerChar is an approximation (1 token ≈ 4 chars for English)
	TokensPerChar = 0.25
)

// ContextWindow manages the construction of context for LLM prompts
type ContextWindow struct {
	// Entries included in the context
	Entries []*MemoryEntry

	// FormattedContext is the final context string
	FormattedContext string

	// TotalTokens estimated in the context
	TotalTokens int

	// MaxTokens allowed
	MaxTokens int

	// IncludedTypes tracks which memory types are represented
	IncludedTypes map[MemoryType]int

	// TruncatedCount is how many entries were truncated
	TruncatedCount int

	// HasSummary indicates if older entries were summarized
	HasSummary bool

	// SummaryText is the summary of older entries
	SummaryText string
}

// ContextWindowBuilder builds optimized context windows
type ContextWindowBuilder struct {
	maxTokens        int
	includeMetadata  bool
	includeSources   bool
	includeTimestamp bool
	formatStyle      ContextFormatStyle
	priorityOrder    []MemoryType
}

// ContextFormatStyle defines how to format entries
type ContextFormatStyle int

const (
	// FormatPlain uses simple text formatting
	FormatPlain ContextFormatStyle = iota
	// FormatMarkdown uses markdown formatting
	FormatMarkdown
	// FormatXML uses XML-like tags
	FormatXML
	// FormatJSON uses JSON formatting
	FormatJSON
)

// NewContextWindowBuilder creates a new builder with defaults
func NewContextWindowBuilder() *ContextWindowBuilder {
	return &ContextWindowBuilder{
		maxTokens:        DefaultMaxTokens,
		includeMetadata:  false,
		includeSources:   true,
		includeTimestamp: true,
		formatStyle:      FormatMarkdown,
		priorityOrder: []MemoryType{
			EpisodicMemory,
			SemanticMemory,
			ProceduralMemory,
			TemporalMemory,
		},
	}
}

// WithMaxTokens sets the maximum token limit
func (b *ContextWindowBuilder) WithMaxTokens(max int) *ContextWindowBuilder {
	b.maxTokens = max
	return b
}

// WithMetadata includes entry metadata
func (b *ContextWindowBuilder) WithMetadata(include bool) *ContextWindowBuilder {
	b.includeMetadata = include
	return b
}

// WithSources includes source information
func (b *ContextWindowBuilder) WithSources(include bool) *ContextWindowBuilder {
	b.includeSources = include
	return b
}

// WithTimestamps includes timestamps
func (b *ContextWindowBuilder) WithTimestamps(include bool) *ContextWindowBuilder {
	b.includeTimestamp = include
	return b
}

// WithFormatStyle sets the formatting style
func (b *ContextWindowBuilder) WithFormatStyle(style ContextFormatStyle) *ContextWindowBuilder {
	b.formatStyle = style
	return b
}

// WithPriorityOrder sets the memory type priority for inclusion
func (b *ContextWindowBuilder) WithPriorityOrder(order []MemoryType) *ContextWindowBuilder {
	b.priorityOrder = order
	return b
}

// Build constructs the context window from entries
func (b *ContextWindowBuilder) Build(entries []*MemoryEntry) *ContextWindow {
	window := &ContextWindow{
		MaxTokens:     b.maxTokens,
		IncludedTypes: make(map[MemoryType]int),
	}

	if len(entries) == 0 {
		window.FormattedContext = ""
		return window
	}

	// Sort entries by priority and recency
	sortedEntries := b.sortByPriority(entries)

	// Build context with token budget
	var contextParts []string
	usedTokens := 0

	for _, entry := range sortedEntries {
		formatted := b.formatEntry(entry)
		entryTokens := estimateTokens(formatted)

		if usedTokens+entryTokens > b.maxTokens {
			window.TruncatedCount++
			continue
		}

		contextParts = append(contextParts, formatted)
		window.Entries = append(window.Entries, entry)
		window.IncludedTypes[entry.Type]++
		usedTokens += entryTokens
	}

	window.FormattedContext = strings.Join(contextParts, "\n\n")
	window.TotalTokens = usedTokens

	return window
}

// BuildWithSummary builds context with summarization of older entries
func (b *ContextWindowBuilder) BuildWithSummary(
	entries []*MemoryEntry,
	summarizer func(entries []*MemoryEntry) string,
) *ContextWindow {
	window := &ContextWindow{
		MaxTokens:     b.maxTokens,
		IncludedTypes: make(map[MemoryType]int),
	}

	if len(entries) == 0 {
		window.FormattedContext = ""
		return window
	}

	// Reserve tokens for summary
	summaryBudget := b.maxTokens / 4
	contentBudget := b.maxTokens - summaryBudget

	// Sort entries by timestamp (newest first)
	sortedEntries := make([]*MemoryEntry, len(entries))
	copy(sortedEntries, entries)

	// Sort by timestamp descending
	for i := 0; i < len(sortedEntries)-1; i++ {
		for j := i + 1; j < len(sortedEntries); j++ {
			if sortedEntries[j].Timestamp.After(sortedEntries[i].Timestamp) {
				sortedEntries[i], sortedEntries[j] = sortedEntries[j], sortedEntries[i]
			}
		}
	}

	// Add recent entries until budget exhausted
	var recentParts []string
	var olderEntries []*MemoryEntry
	usedTokens := 0

	for i, entry := range sortedEntries {
		formatted := b.formatEntry(entry)
		entryTokens := estimateTokens(formatted)

		if usedTokens+entryTokens > contentBudget {
			// Remaining entries go to summary
			olderEntries = sortedEntries[i:]
			break
		}

		recentParts = append(recentParts, formatted)
		window.Entries = append(window.Entries, entry)
		window.IncludedTypes[entry.Type]++
		usedTokens += entryTokens
	}

	// Generate summary for older entries
	if len(olderEntries) > 0 && summarizer != nil {
		summary := summarizer(olderEntries)
		window.HasSummary = true
		window.SummaryText = summary
		window.TruncatedCount = len(olderEntries)

		// Format summary
		summaryFormatted := b.formatSummary(summary)
		usedTokens += estimateTokens(summaryFormatted)

		// Prepend summary to context
		window.FormattedContext = summaryFormatted + "\n\n---\n\n" + strings.Join(recentParts, "\n\n")
	} else {
		window.FormattedContext = strings.Join(recentParts, "\n\n")
	}

	window.TotalTokens = usedTokens
	return window
}

// sortByPriority sorts entries by memory type priority and recency
func (b *ContextWindowBuilder) sortByPriority(entries []*MemoryEntry) []*MemoryEntry {
	// Create priority map
	priorityMap := make(map[MemoryType]int)
	for i, mt := range b.priorityOrder {
		priorityMap[mt] = i
	}

	sorted := make([]*MemoryEntry, len(entries))
	copy(sorted, entries)

	// Sort by priority, then by timestamp (descending)
	for i := 0; i < len(sorted)-1; i++ {
		for j := i + 1; j < len(sorted); j++ {
			iPriority := priorityMap[sorted[i].Type]
			jPriority := priorityMap[sorted[j].Type]

			swap := false
			if iPriority > jPriority {
				swap = true
			} else if iPriority == jPriority {
				if sorted[j].Timestamp.After(sorted[i].Timestamp) {
					swap = true
				}
			}

			if swap {
				sorted[i], sorted[j] = sorted[j], sorted[i]
			}
		}
	}

	return sorted
}

// formatEntry formats a single entry based on style
func (b *ContextWindowBuilder) formatEntry(entry *MemoryEntry) string {
	switch b.formatStyle {
	case FormatMarkdown:
		return b.formatEntryMarkdown(entry)
	case FormatXML:
		return b.formatEntryXML(entry)
	case FormatJSON:
		return b.formatEntryJSON(entry)
	default:
		return b.formatEntryPlain(entry)
	}
}

func (b *ContextWindowBuilder) formatEntryPlain(entry *MemoryEntry) string {
	var parts []string

	if b.includeTimestamp {
		parts = append(parts, fmt.Sprintf("[%s]", entry.Timestamp.Format(time.RFC3339)))
	}

	if b.includeSources {
		parts = append(parts, fmt.Sprintf("(%s/%s)", entry.Type.String(), entry.Source))
	}

	parts = append(parts, entry.Content)

	return strings.Join(parts, " ")
}

func (b *ContextWindowBuilder) formatEntryMarkdown(entry *MemoryEntry) string {
	var sb strings.Builder

	// Header with type and source
	sb.WriteString(fmt.Sprintf("### %s Memory", strings.Title(entry.Type.String())))

	if b.includeSources && entry.Source != "" {
		sb.WriteString(fmt.Sprintf(" (%s)", entry.Source))
	}
	sb.WriteString("\n")

	if b.includeTimestamp {
		sb.WriteString(fmt.Sprintf("*%s*\n\n", entry.Timestamp.Format("Jan 2, 2006 3:04 PM")))
	}

	sb.WriteString(entry.Content)

	if b.includeMetadata && len(entry.Tags) > 0 {
		sb.WriteString(fmt.Sprintf("\n\n**Tags**: %s", strings.Join(entry.Tags, ", ")))
	}

	return sb.String()
}

func (b *ContextWindowBuilder) formatEntryXML(entry *MemoryEntry) string {
	var sb strings.Builder

	sb.WriteString(fmt.Sprintf("<memory type=\"%s\"", entry.Type.String()))

	if b.includeSources {
		sb.WriteString(fmt.Sprintf(" source=\"%s\"", entry.Source))
	}

	if b.includeTimestamp {
		sb.WriteString(fmt.Sprintf(" timestamp=\"%s\"", entry.Timestamp.Format(time.RFC3339)))
	}

	sb.WriteString(">\n")
	sb.WriteString(entry.Content)
	sb.WriteString("\n</memory>")

	return sb.String()
}

func (b *ContextWindowBuilder) formatEntryJSON(entry *MemoryEntry) string {
	// Simple JSON-like format without full JSON marshaling
	var parts []string

	parts = append(parts, fmt.Sprintf(`"type": "%s"`, entry.Type.String()))

	if b.includeSources {
		parts = append(parts, fmt.Sprintf(`"source": "%s"`, entry.Source))
	}

	if b.includeTimestamp {
		parts = append(parts, fmt.Sprintf(`"timestamp": "%s"`, entry.Timestamp.Format(time.RFC3339)))
	}

	parts = append(parts, fmt.Sprintf(`"content": "%s"`, escapeJSON(entry.Content)))

	return "{\n  " + strings.Join(parts, ",\n  ") + "\n}"
}

func (b *ContextWindowBuilder) formatSummary(summary string) string {
	switch b.formatStyle {
	case FormatMarkdown:
		return fmt.Sprintf("## Earlier Context Summary\n\n%s", summary)
	case FormatXML:
		return fmt.Sprintf("<summary>\n%s\n</summary>", summary)
	default:
		return fmt.Sprintf("[Summary of earlier context: %s]", summary)
	}
}

// estimateTokens estimates the token count for a string
func estimateTokens(s string) int {
	// Rough approximation: 1 token ≈ 4 characters for English text
	// This is a simplified heuristic; production should use tiktoken
	return int(float64(len(s)) * TokensPerChar)
}

func escapeJSON(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `"`, `\"`)
	s = strings.ReplaceAll(s, "\n", `\n`)
	s = strings.ReplaceAll(s, "\r", `\r`)
	s = strings.ReplaceAll(s, "\t", `\t`)
	return s
}

// DynamicContextSizer adjusts context limits based on query complexity
type DynamicContextSizer struct {
	minTokens int
	maxTokens int
}

// NewDynamicContextSizer creates a new dynamic sizer
func NewDynamicContextSizer(min, max int) *DynamicContextSizer {
	return &DynamicContextSizer{
		minTokens: min,
		maxTokens: max,
	}
}

// GetLimit returns the appropriate token limit based on query complexity
func (s *DynamicContextSizer) GetLimit(query string, complexity float64) int {
	// Complexity range: 0.0 (simple) to 1.0 (complex)

	// Simple queries get smaller context
	if complexity < 0.3 {
		return s.minTokens + int(float64(s.maxTokens-s.minTokens)*0.25)
	}

	// Complex queries get larger context
	if complexity > 0.7 {
		return s.minTokens + int(float64(s.maxTokens-s.minTokens)*0.75)
	}

	// Medium complexity
	return s.minTokens + int(float64(s.maxTokens-s.minTokens)*0.5)
}

// EstimateQueryComplexity provides a simple heuristic for query complexity
func EstimateQueryComplexity(query string) float64 {
	// Factors that increase complexity:
	// - Question length
	// - Number of clauses (commas, "and", "or")
	// - Presence of technical terms
	// - Temporal references
	// - Comparison requests

	score := 0.0

	// Length factor
	wordCount := len(strings.Fields(query))
	if wordCount > 20 {
		score += 0.3
	} else if wordCount > 10 {
		score += 0.15
	}

	// Clause complexity
	clauseIndicators := []string{",", " and ", " or ", " but ", " however "}
	for _, indicator := range clauseIndicators {
		if strings.Contains(strings.ToLower(query), indicator) {
			score += 0.1
		}
	}

	// Temporal references
	temporalTerms := []string{"yesterday", "last week", "before", "after", "when", "since", "until"}
	for _, term := range temporalTerms {
		if strings.Contains(strings.ToLower(query), term) {
			score += 0.1
			break
		}
	}

	// Comparison/analysis requests
	analysisTerms := []string{"compare", "difference", "why", "how", "explain", "analyze"}
	for _, term := range analysisTerms {
		if strings.Contains(strings.ToLower(query), term) {
			score += 0.15
			break
		}
	}

	// Cap at 1.0
	if score > 1.0 {
		score = 1.0
	}

	return score
}

// ContextWindowStats provides statistics about a context window
type ContextWindowStats struct {
	TotalEntries     int            `json:"total_entries"`
	IncludedEntries  int            `json:"included_entries"`
	TruncatedEntries int            `json:"truncated_entries"`
	TotalTokens      int            `json:"total_tokens"`
	MaxTokens        int            `json:"max_tokens"`
	UtilizationPct   float64        `json:"utilization_pct"`
	TypeDistribution map[string]int `json:"type_distribution"`
	HasSummary       bool           `json:"has_summary"`
}

// Stats returns statistics for a context window
func (w *ContextWindow) Stats() ContextWindowStats {
	typeDist := make(map[string]int)
	for mt, count := range w.IncludedTypes {
		typeDist[mt.String()] = count
	}

	return ContextWindowStats{
		TotalEntries:     len(w.Entries) + w.TruncatedCount,
		IncludedEntries:  len(w.Entries),
		TruncatedEntries: w.TruncatedCount,
		TotalTokens:      w.TotalTokens,
		MaxTokens:        w.MaxTokens,
		UtilizationPct:   float64(w.TotalTokens) / float64(w.MaxTokens) * 100,
		TypeDistribution: typeDist,
		HasSummary:       w.HasSummary,
	}
}