"""
Context Budget — Token budget control for context injection.

PRD §10 — Token Budget Control:
    max_context_tokens = 800 (configurable)
    Priority: user memory > summary > timeline > raw
    Strategy: overflow → truncate timeline, keep summary

Uses simple word-count approximation (1 token ≈ 0.75 words).
"""

import logging
from typing import List, Optional, Dict

logger = logging.getLogger("openmemo.context_memory")

# 1 token ≈ 0.75 words  →  words ≈ tokens * 0.75
_WORDS_PER_TOKEN = 0.75


def estimate_tokens(text: str) -> int:
    """Approximate token count from text (word-based heuristic)."""
    if not text:
        return 0
    word_count = len(text.split())
    return max(1, int(word_count / _WORDS_PER_TOKEN))


class ContextBudget:
    """
    Manages token budget for context injection.

    PRD §10 priority order:
        user_memory > summary > timeline > raw

    Usage:
        budget = ContextBudget(max_tokens=800)
        result = budget.fit(
            user_memory_block="USER PROFILE...",
            summary_block="CONTEXT SUMMARY...",
            timeline_block="Timeline:...",
            raw_block="...",
        )
    """

    def __init__(self, max_tokens: int = 800):
        self.max_tokens = max_tokens

    def fit(
        self,
        user_memory_block: str = "",
        summary_block: str = "",
        timeline_block: str = "",
        raw_block: str = "",
    ) -> Dict[str, str]:
        """
        Fit context blocks within token budget.

        Returns dict of blocks that fit within budget:
            { "user_memory": ..., "summary": ..., "timeline": ..., "raw": ... }

        PRD §10 strategy: drop lower-priority blocks when over budget.
        """
        budget = self.max_tokens
        result = {}

        # Priority 1: user memory
        um_tokens = estimate_tokens(user_memory_block)
        if user_memory_block and um_tokens <= budget:
            result["user_memory"] = user_memory_block
            budget -= um_tokens
        elif user_memory_block:
            # Truncate to fit
            result["user_memory"] = self._truncate_to_tokens(
                user_memory_block, budget
            )
            budget = 0

        # Priority 2: summary (always try to keep)
        sum_tokens = estimate_tokens(summary_block)
        if summary_block and sum_tokens <= budget:
            result["summary"] = summary_block
            budget -= sum_tokens
        elif summary_block and budget > 20:
            result["summary"] = self._truncate_to_tokens(summary_block, budget)
            budget = 0

        # Priority 3: timeline
        tl_tokens = estimate_tokens(timeline_block)
        if timeline_block and tl_tokens <= budget:
            result["timeline"] = timeline_block
            budget -= tl_tokens
        elif timeline_block and budget > 20:
            result["timeline"] = self._truncate_to_tokens(timeline_block, budget)
            budget = 0

        # Priority 4: raw (lowest priority)
        if raw_block and budget > 20:
            result["raw"] = self._truncate_to_tokens(raw_block, budget)

        total_used = self.max_tokens - budget
        logger.debug("context_budget: used=%d/%d tokens", total_used, self.max_tokens)
        return result

    def combine(
        self,
        user_memory_block: str = "",
        summary_block: str = "",
        timeline_block: str = "",
        raw_block: str = "",
    ) -> str:
        """
        Fit blocks and combine into a single context string.

        Injection order per PRD §11.1:
            Context Memory (summary) → User Memory → Task Memory → Experience
        """
        fitted = self.fit(
            user_memory_block=user_memory_block,
            summary_block=summary_block,
            timeline_block=timeline_block,
            raw_block=raw_block,
        )
        parts = []
        # PRD §11.1 injection order: context summary FIRST
        if fitted.get("summary"):
            parts.append(fitted["summary"])
        if fitted.get("user_memory"):
            parts.append(fitted["user_memory"])
        if fitted.get("timeline"):
            parts.append(fitted["timeline"])
        if fitted.get("raw"):
            parts.append(fitted["raw"])
        return "\n\n".join(parts)

    def remaining_tokens(self, *blocks: str) -> int:
        """How many tokens remain after accounting for the given blocks."""
        used = sum(estimate_tokens(b) for b in blocks)
        return max(0, self.max_tokens - used)

    def _truncate_to_tokens(self, text: str, max_tokens: int) -> str:
        """Truncate text to fit within max_tokens (word-based)."""
        max_words = int(max_tokens * _WORDS_PER_TOKEN)
        words = text.split()
        if len(words) <= max_words:
            return text
        return " ".join(words[:max_words]) + " ..."
