/headroom/memory/budget.py

"""Memory file budget manager — token-optimized memory file maintenance.

Manages the token budget for agent memory files:
- Ranks memories by importance × recency × access frequency
- Prunes memories that fall below budget threshold
- Detects stale memories referencing deleted/renamed entities
- Merges similar memories to save tokens

This is Headroom's superpower applied to itself — the proxy that optimizes
LLM context also optimizes its own memory files for minimum token consumption.
"""

from __future__ import annotations

import logging
import re
import subprocess
import time
from dataclasses import dataclass, field
from pathlib import Path

from headroom._subprocess import run
from headroom.memory.writers.base import MemoryEntry, _estimate_tokens

logger = logging.getLogger(__name__)


@dataclass
class BudgetConfig:
    """Configuration for memory budget management."""

    # Token budgets per agent type
    agent_budgets: dict[str, int] = field(
        default_factory=lambda: {
            "claude": 2000,  # Claude Code: 200-line MEMORY.md limit
            "cursor": 3000,  # Cursor: .mdc rules loaded on match
            "codex": 3000,  # Codex: AGENTS.md merged
            "aider": 2000,  # Aider: read files
            "gemini": 3000,  # Gemini: GEMINI.md
            "generic": 3000,  # Default
        }
    )

    # Staleness detection
    staleness_decay_rate: float = 0.1  # Importance decays ~10% per day
    staleness_min_importance: float = 0.15  # Below this = prunable
    staleness_check_git: bool = True  # Check git for deleted files

    # Deduplication
    similarity_merge_threshold: float = 0.85  # Cosine sim above = merge


@dataclass
class BudgetReport:
    """Report from a budget optimization pass."""

    total_memories: int = 0
    kept: int = 0
    pruned_staleness: int = 0
    pruned_budget: int = 0
    merged: int = 0
    tokens_before: int = 0
    tokens_after: int = 0

    @property
    def tokens_saved(self) -> int:
        return self.tokens_before - self.tokens_after


class MemoryBudgetManager:
    """Manages token budgets for agent memory files.

    Responsible for deciding which memories to keep, prune, or merge
    to stay within token budgets while maximizing information density.
    """

    def __init__(
        self,
        project_path: Path | None = None,
        config: BudgetConfig | None = None,
    ) -> None:
        self._project_path = project_path or Path.cwd()
        self._config = config or BudgetConfig()
        self._git_files_cache: set[str] | None = None

    def optimize(
        self,
        memories: list[MemoryEntry],
        agent_type: str = "generic",
    ) -> tuple[list[MemoryEntry], BudgetReport]:
        """Optimize a set of memories for an agent's token budget.

        Returns the optimized list and a report of what changed.

        Args:
            memories: All candidate memories.
            agent_type: Target agent type (determines budget).

        Returns:
            Tuple of (optimized memories, report).
        """
        report = BudgetReport(total_memories=len(memories))
        report.tokens_before = sum(_estimate_tokens(m.content) for m in memories)

        budget = self._config.agent_budgets.get(agent_type, 3000)

        # Step 1: Apply temporal decay to importance scores
        decayed = self._apply_decay(memories)

        # Step 2: Detect and flag stale memories
        fresh, stale = self._detect_staleness(decayed)
        report.pruned_staleness = len(stale)

        # Step 3: Merge similar memories
        merged = self._merge_similar(fresh)
        report.merged = len(fresh) - len(merged)

        # Step 4: Rank by score and apply budget
        ranked = sorted(merged, key=lambda m: m.score, reverse=True)
        budgeted: list[MemoryEntry] = []
        tokens_used = 0
        for m in ranked:
            entry_tokens = _estimate_tokens(m.content) + 10
            if tokens_used + entry_tokens > budget:
                report.pruned_budget += 1
                continue
            tokens_used += entry_tokens
            budgeted.append(m)

        report.kept = len(budgeted)
        report.tokens_after = tokens_used

        return budgeted, report

    def _apply_decay(self, memories: list[MemoryEntry]) -> list[MemoryEntry]:
        """Apply temporal decay to importance scores."""
        now = time.time()
        rate = self._config.staleness_decay_rate
        min_importance = self._config.staleness_min_importance

        result: list[MemoryEntry] = []
        for m in memories:
            age_days = (now - m.created_at) / 86400
            # Exponential decay: importance × e^(-rate × age)
            import math

            decayed_importance = m.importance * math.exp(-rate * age_days)

            # Access count counteracts decay
            if m.access_count > 0:
                access_boost = min(0.3, m.access_count * 0.05)
                decayed_importance = min(1.0, decayed_importance + access_boost)

            if decayed_importance >= min_importance:
                m.importance = decayed_importance
                result.append(m)

        return result

    def _detect_staleness(
        self,
        memories: list[MemoryEntry],
    ) -> tuple[list[MemoryEntry], list[MemoryEntry]]:
        """Separate fresh memories from stale ones.

        A memory is stale if it references files/paths that no longer exist.
        Uses git ls-files when available, falls back to filesystem checks.
        """
        git_files = self._get_git_files() if self._config.staleness_check_git else set()

        fresh: list[MemoryEntry] = []
        stale: list[MemoryEntry] = []

        for m in memories:
            if self._is_stale(m, git_files):
                stale.append(m)
            else:
                fresh.append(m)

        return fresh, stale

    def _is_stale(self, memory: MemoryEntry, git_files: set[str]) -> bool:
        """Check if a memory references entities that no longer exist."""
        # Check entity_refs for file paths
        for ref in memory.entity_refs:
            if ref.startswith("/") or ref.startswith("./"):
                # Absolute or relative path — check if file exists
                if ref.startswith("./"):
                    ref = ref[2:]
                # Normalize to relative
                try:
                    rel = str(Path(ref).relative_to(self._project_path))
                except ValueError:
                    rel = ref

                if rel not in git_files and not Path(ref).exists():
                    return True

        # Check content for backtick-quoted file paths
        path_refs = re.findall(r"`([/\w.]+(?:/[\w.]+)+)`", memory.content)
        for path_ref in path_refs:
            # Only flag if it looks like a full path and is missing
            if "/" in path_ref:
                try:
                    rel = str(Path(path_ref).relative_to(self._project_path))
                except ValueError:
                    rel = path_ref
                if rel not in git_files and not Path(path_ref).exists():
                    return True

        return False

    def _get_git_files(self) -> set[str]:
        """Get set of tracked files in the git repo."""
        if self._git_files_cache is not None:
            return self._git_files_cache

        try:
            result = run(
                ["git", "ls-files"],
                capture_output=True,
                text=True,
                cwd=self._project_path,
                timeout=5,
            )
            if result.returncode == 0:
                self._git_files_cache = set(result.stdout.strip().split("\n"))
            else:
                self._git_files_cache = set()
        except (subprocess.TimeoutExpired, FileNotFoundError):
            self._git_files_cache = set()

        return self._git_files_cache

    def _merge_similar(self, memories: list[MemoryEntry]) -> list[MemoryEntry]:
        """Merge memories with very similar content.

        Uses simple text overlap heuristic (no embeddings required).
        """
        if len(memories) <= 1:
            return memories

        merged: list[MemoryEntry] = []
        used: set[int] = set()

        for i, m1 in enumerate(memories):
            if i in used:
                continue

            # Find similar memories
            group = [m1]
            for j, m2 in enumerate(memories[i + 1 :], start=i + 1):
                if j in used:
                    continue
                if (
                    self._text_similarity(m1.content, m2.content)
                    > self._config.similarity_merge_threshold
                ):
                    group.append(m2)
                    used.add(j)

            if len(group) == 1:
                merged.append(m1)
            else:
                # Merge: keep highest-importance content, combine entity refs
                best = max(group, key=lambda m: m.importance)
                all_entities = set()
                for m in group:
                    all_entities.update(m.entity_refs)
                best.entity_refs = list(all_entities)
                best.access_count = max(m.access_count for m in group)
                merged.append(best)

        return merged

    @staticmethod
    def _text_similarity(a: str, b: str) -> float:
        """Simple Jaccard similarity on word sets."""
        words_a = set(a.lower().split())
        words_b = set(b.lower().split())
        if not words_a or not words_b:
            return 0.0
        intersection = words_a & words_b
        union = words_a | words_b
        return len(intersection) / len(union)