"""
Issue #121: Token budget enforcement for context assembly.

Provides accurate tiktoken-based token counting and priority-ordered
truncation for get_context() and get_full_context().
"""

from __future__ import annotations

import logging
from typing import Optional

logger = logging.getLogger(__name__)

# Lazy-loaded tiktoken encoder
_encoder_cache: dict[str, "tiktoken.Encoding"] = {}


def _get_encoder(model: str = "gpt-4o"):
    """Get (cached) tiktoken encoder for a model."""
    if model not in _encoder_cache:
        try:
            import tiktoken
            _encoder_cache[model] = tiktoken.encoding_for_model(model)
        except Exception:
            # Fallback to cl100k_base (GPT-4 / GPT-3.5 family)
            import tiktoken
            _encoder_cache[model] = tiktoken.get_encoding("cl100k_base")
            logger.warning(
                f"tiktoken has no encoding for model '{model}', "
                "falling back to cl100k_base"
            )
    return _encoder_cache[model]


def count_tokens(text: str, model: str = "gpt-4o") -> int:
    """Count tokens in text using tiktoken."""
    if not text:
        return 0
    enc = _get_encoder(model)
    return len(enc.encode(text))


def truncate_to_tokens(text: str, max_tokens: int, model: str = "gpt-4o") -> str:
    """Truncate text to fit within max_tokens, splitting on line boundaries.

    Tries to keep complete lines. If even the first line exceeds the budget,
    truncates mid-line with a trailing ``...`` marker.
    """
    if not text or max_tokens <= 0:
        return ""
    enc = _get_encoder(model)
    tokens = enc.encode(text)
    if len(tokens) <= max_tokens:
        return text

    # Truncate at token level, then trim to last complete line
    truncated = enc.decode(tokens[:max_tokens])
    # Find last newline to avoid cutting mid-line
    last_nl = truncated.rfind("\n")
    if last_nl > 0:
        truncated = truncated[:last_nl]
    return truncated + "\n[... truncated to fit token budget]"


class TokenBudget:
    """Tracks remaining token budget and assembles sections in priority order.

    Usage::

        budget = TokenBudget(max_tokens=4000, model="gpt-4o")
        budget.add_section(goals_text)    # added if fits
        budget.add_section(decisions_text) # added if fits, truncated if partial
        result = budget.build()
    """

    def __init__(self, max_tokens: int, model: str = "gpt-4o"):
        self.max_tokens = max_tokens
        self.model = model
        self._sections: list[str] = []
        self._tokens_used: int = 0
        # Account for separator tokens between sections ("\n\n" ≈ 1-2 tokens)
        self._separator_overhead = 2

    @property
    def tokens_used(self) -> int:
        return self._tokens_used

    @property
    def tokens_remaining(self) -> int:
        return max(0, self.max_tokens - self._tokens_used)

    @property
    def is_exhausted(self) -> bool:
        return self._tokens_used >= self.max_tokens

    def add_section(self, text: str) -> bool:
        """Add a section if budget permits. Truncates if partial fit.

        Returns True if the full section was added, False if truncated or skipped.
        """
        if not text or self.is_exhausted:
            return False

        section_tokens = count_tokens(text, self.model)
        separator_cost = self._separator_overhead if self._sections else 0

        if section_tokens + separator_cost <= self.tokens_remaining:
            # Full section fits
            self._sections.append(text)
            self._tokens_used += section_tokens + separator_cost
            return True

        # Partial fit — truncate to remaining budget
        available = self.tokens_remaining - separator_cost
        if available <= 10:
            # Not enough space for meaningful content
            return False

        truncated = truncate_to_tokens(text, available, self.model)
        if truncated and count_tokens(truncated, self.model) > 0:
            self._sections.append(truncated)
            self._tokens_used += count_tokens(truncated, self.model) + separator_cost
            return False  # indicates truncation happened

        return False

    def build(self) -> str:
        """Return assembled context string."""
        return "\n\n".join(self._sections)
