"""Context budget management for Oclawma.

This module provides token budget tracking with strict enforcement
and warnings for the 8K context window limitation.
"""

from dataclasses import dataclass
from enum import Enum
from typing import Callable, Optional


class BudgetThreshold(Enum):
    """Budget threshold levels for warnings."""

    NORMAL = "normal"
    WARNING = "warning"  # 75% - 6K tokens
    CRITICAL = "critical"  # 87% - 7K tokens
    EXHAUSTED = "exhausted"  # 100% - 8K tokens


@dataclass
class BudgetStatus:
    """Current status of the context budget."""

    used_tokens: int
    total_tokens: int
    remaining_tokens: int
    usage_percent: float
    threshold: BudgetThreshold
    warning_message: Optional[str] = None


class BudgetExceededError(Exception):
    """Raised when the context budget is exceeded."""

    def __init__(self, used: int, limit: int, message: str = ""):
        self.used = used
        self.limit = limit
        self.message = message or (
            f"🛑 CONTEXT BUDGET EXHAUSTED: {used:,} / {limit:,} tokens used (100%)\n\n"
            "To continue:\n"
            "  1. Start a new conversation session\n"
            "  2. Use /compact to summarize context\n"
            "  3. Clear unused attachments with /clear\n"
            "  4. Focus on specific tasks to reduce context buildup"
        )
        super().__init__(self.message)


class ContextBudget:
    """Manages token budget tracking for 8K context models.

    Provides real-time tracking with warnings at 75% (6K) and 87% (7K)
    thresholds, and hard stop at 100% (8K).

    Attributes:
        total_budget: Total token budget (default 8192)
        warning_threshold: First warning at this many tokens (default 6144)
        critical_threshold: Critical warning at this many tokens (default 7168)
    """

    DEFAULT_BUDGET = 8192
    WARNING_THRESHOLD = 6144  # 75%
    CRITICAL_THRESHOLD = 7168  # 87%

    def __init__(
        self,
        total_budget: int = DEFAULT_BUDGET,
        warning_threshold: int = WARNING_THRESHOLD,
        critical_threshold: int = CRITICAL_THRESHOLD,
        on_warning: Optional[Callable[[str], None]] = None,
        on_critical: Optional[Callable[[str], None]] = None,
    ):
        """Initialize the context budget tracker.

        Args:
            total_budget: Total token budget
            warning_threshold: Tokens at which to show first warning (75%)
            critical_threshold: Tokens at which to show critical warning (87%)
            on_warning: Callback for warning threshold
            on_critical: Callback for critical threshold
        """
        self.total_budget = total_budget
        self.warning_threshold = warning_threshold
        self.critical_threshold = critical_threshold
        self._used_tokens = 0
        self._on_warning = on_warning
        self._on_critical = on_critical
        self._warning_shown = False
        self._critical_shown = False

    @property
    def used_tokens(self) -> int:
        """Get the number of tokens currently used."""
        return self._used_tokens

    @property
    def remaining_tokens(self) -> int:
        """Get the number of tokens remaining."""
        return self.total_budget - self._used_tokens

    @property
    def usage_percent(self) -> float:
        """Get the percentage of budget used."""
        return (self._used_tokens / self.total_budget) * 100

    def allocate(self, tokens: int, strict: bool = True) -> int:
        """Allocate tokens from the budget.

        Args:
            tokens: Number of tokens to allocate
            strict: If True, raises BudgetExceededError when budget exhausted

        Returns:
            Number of tokens allocated

        Raises:
            BudgetExceededError: If strict=True and budget would be exceeded
            ValueError: If tokens is negative
        """
        if tokens < 0:
            raise ValueError(f"Cannot allocate negative tokens: {tokens}")

        if tokens == 0:
            return 0

        new_total = self._used_tokens + tokens

        # Check if budget would be exceeded
        if new_total > self.total_budget:
            if strict:
                raise BudgetExceededError(new_total, self.total_budget)
            else:
                # Allocate what we can
                available = self.total_budget - self._used_tokens
                self._used_tokens = self.total_budget
                self._check_thresholds()
                return available

        self._used_tokens = new_total
        self._check_thresholds()
        return tokens

    def release(self, tokens: int) -> int:
        """Release tokens back to the budget.

        Args:
            tokens: Number of tokens to release

        Returns:
            Number of tokens actually released

        Raises:
            ValueError: If tokens is negative
        """
        if tokens < 0:
            raise ValueError(f"Cannot release negative tokens: {tokens}")

        released = min(tokens, self._used_tokens)
        self._used_tokens -= released

        # Reset warning flags if we're back below thresholds
        if self._used_tokens < self.warning_threshold:
            self._warning_shown = False
            self._critical_shown = False
        elif self._used_tokens < self.critical_threshold:
            self._critical_shown = False

        return released

    def reset(self) -> None:
        """Reset the budget to zero used tokens."""
        self._used_tokens = 0
        self._warning_shown = False
        self._critical_shown = False

    def get_status(self) -> BudgetStatus:
        """Get the current budget status.

        Returns:
            BudgetStatus with current usage information
        """
        threshold = self._get_current_threshold()
        warning_msg = None

        if threshold == BudgetThreshold.WARNING:
            warning_msg = (
                f"⚠️  Context budget at 75% ({self._used_tokens:,} / {self.total_budget:,} tokens).\n"
                "Consider: /compact to summarize, or start fresh."
            )
        elif threshold == BudgetThreshold.CRITICAL:
            warning_msg = (
                f"🚨 Context budget at 87% ({self._used_tokens:,} / {self.total_budget:,} tokens)!\n"
                "Action recommended: /compact immediately or start new session."
            )
        elif threshold == BudgetThreshold.EXHAUSTED:
            warning_msg = (
                f"🛑 Context budget exhausted ({self._used_tokens:,} / {self.total_budget:,} tokens)!\n"
                "Start a new session or use /compact to continue."
            )

        return BudgetStatus(
            used_tokens=self._used_tokens,
            total_tokens=self.total_budget,
            remaining_tokens=self.remaining_tokens,
            usage_percent=self.usage_percent,
            threshold=threshold,
            warning_message=warning_msg,
        )

    def _get_current_threshold(self) -> BudgetThreshold:
        """Determine which threshold level we're at."""
        if self._used_tokens >= self.total_budget:
            return BudgetThreshold.EXHAUSTED
        elif self._used_tokens >= self.critical_threshold:
            return BudgetThreshold.CRITICAL
        elif self._used_tokens >= self.warning_threshold:
            return BudgetThreshold.WARNING
        return BudgetThreshold.NORMAL

    def _check_thresholds(self) -> None:
        """Check and trigger threshold callbacks."""
        if self._used_tokens >= self.critical_threshold and not self._critical_shown:
            self._critical_shown = True
            if self._on_critical:
                self._on_critical(self.get_status().warning_message or "")
        elif self._used_tokens >= self.warning_threshold and not self._warning_shown:
            self._warning_shown = True
            if self._on_warning:
                self._on_warning(self.get_status().warning_message or "")

    def can_allocate(self, tokens: int) -> bool:
        """Check if the requested tokens can be allocated.

        Args:
            tokens: Number of tokens to check

        Returns:
            True if tokens can be allocated, False otherwise
        """
        return self._used_tokens + tokens <= self.total_budget

    def estimate_tokens(self, text: str) -> int:
        """Estimate token count for text (rough approximation).

        This is a simple approximation: ~4 characters per token.
        For accurate counts, use a proper tokenizer.

        Args:
            text: Text to estimate

        Returns:
            Estimated token count
        """
        # Rough approximation: 4 characters per token on average
        return max(1, len(text) // 4)

    def visualize(self, width: int = 40) -> str:
        """Create a visual representation of the budget.

        Args:
            width: Width of the progress bar in characters

        Returns:
            ASCII progress bar with status
        """
        status = self.get_status()

        # Determine color/character based on threshold
        if status.threshold == BudgetThreshold.EXHAUSTED:
            fill_char = "█"
            color_code = "\033[91m"  # Red
        elif status.threshold == BudgetThreshold.CRITICAL:
            fill_char = "█"
            color_code = "\033[93m"  # Yellow/Orange
        elif status.threshold == BudgetThreshold.WARNING:
            fill_char = "█"
            color_code = "\033[93m"  # Yellow
        else:
            fill_char = "█"
            color_code = "\033[92m"  # Green

        reset_code = "\033[0m"

        # Calculate filled portion
        filled = int((self._used_tokens / self.total_budget) * width)
        filled = min(filled, width)  # Cap at width
        empty = width - filled

        # Build the bar
        bar = fill_char * filled + "░" * empty

        # Status indicator
        if status.threshold == BudgetThreshold.EXHAUSTED:
            indicator = "🛑"
        elif status.threshold == BudgetThreshold.CRITICAL:
            indicator = "🚨"
        elif status.threshold == BudgetThreshold.WARNING:
            indicator = "⚠️ "
        else:
            indicator = "✅"

        return (
            f"{indicator} Context Budget: [{color_code}{bar}{reset_code}] "
            f"{self._used_tokens:,} / {self.total_budget:,} tokens "
            f"({self.usage_percent:.1f}%)"
        )

    def to_dict(self) -> dict:
        """Convert budget state to dictionary.

        Returns:
            Dictionary representation of budget state
        """
        status = self.get_status()
        return {
            "total_budget": self.total_budget,
            "used_tokens": self._used_tokens,
            "remaining_tokens": self.remaining_tokens,
            "usage_percent": round(self.usage_percent, 2),
            "threshold": status.threshold.value,
            "warning_threshold": self.warning_threshold,
            "critical_threshold": self.critical_threshold,
        }
