"""Context budget management for Oclawma.

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

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from enum import Enum


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: str | None = 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: Callable[[str], None] | None = None,
        on_critical: Callable[[str], None] | 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,
        }
