"""Context manager for tracking token usage and conversation state."""

from dataclasses import dataclass, field
from typing import Any


@dataclass
class TokenUsage:
    """Token usage tracking."""

    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def total(self) -> int:
        """Total tokens used."""
        return self.input_tokens + self.output_tokens

    def add(self, input_tokens: int = 0, output_tokens: int = 0) -> None:
        """Add token usage."""
        self.input_tokens += input_tokens
        self.output_tokens += output_tokens

    def reset(self) -> None:
        """Reset token counts."""
        self.input_tokens = 0
        self.output_tokens = 0


@dataclass
class ContextWindow:
    """Tracks context window usage.

    Different models have different context windows:
    - Claude Sonnet/Opus: 200K tokens
    - Claude Haiku: 200K tokens
    """

    max_tokens: int = 200_000
    used_tokens: int = 0
    warning_threshold: float = 0.8  # 80%

    @property
    def available(self) -> int:
        """Available tokens remaining."""
        return max(0, self.max_tokens - self.used_tokens)

    @property
    def usage_percentage(self) -> float:
        """Percentage of context used."""
        if self.max_tokens == 0:
            return 0.0
        return self.used_tokens / self.max_tokens

    @property
    def is_warning(self) -> bool:
        """Whether usage is above warning threshold."""
        return self.usage_percentage >= self.warning_threshold

    @property
    def is_full(self) -> bool:
        """Whether context is effectively full (>95%)."""
        return self.usage_percentage >= 0.95

    def add_usage(self, tokens: int) -> None:
        """Add tokens to usage."""
        self.used_tokens += tokens

    def reset(self) -> None:
        """Reset usage tracking."""
        self.used_tokens = 0


# Model context window sizes
MODEL_CONTEXT_SIZES = {
    "claude-sonnet-4-20250514": 200_000,
    "claude-opus-4-20250514": 200_000,
    "claude-haiku-4-20250514": 200_000,
    "sonnet": 200_000,
    "opus": 200_000,
    "haiku": 200_000,
}


@dataclass
class ContextManager:
    """Manages conversation context and token tracking.

    Provides:
    - Token usage tracking per conversation
    - Context window monitoring with warnings
    - Model-specific context limits

    Usage:
        ctx = ContextManager(model="sonnet")

        # Track usage from responses
        ctx.add_message_tokens(input_tokens=100, output_tokens=500)

        # Check if near limit
        if ctx.is_warning:
            print(f"Context {ctx.usage_percentage:.0%} full")

        # Get summary
        print(ctx.get_status())
    """

    model: str = "sonnet"
    token_usage: TokenUsage = field(default_factory=TokenUsage)
    context_window: ContextWindow = field(default_factory=ContextWindow)
    message_count: int = 0

    def __post_init__(self):
        # Set context window size based on model
        max_tokens = MODEL_CONTEXT_SIZES.get(self.model, 200_000)
        self.context_window = ContextWindow(max_tokens=max_tokens)

    def set_model(self, model: str) -> None:
        """Update model and adjust context window.

        Args:
            model: Model identifier.
        """
        self.model = model
        max_tokens = MODEL_CONTEXT_SIZES.get(model, 200_000)
        self.context_window.max_tokens = max_tokens

    def add_message_tokens(self, input_tokens: int = 0, output_tokens: int = 0) -> None:
        """Track tokens from a message exchange.

        Args:
            input_tokens: Tokens in the input/prompt.
            output_tokens: Tokens in the response.
        """
        self.token_usage.add(input_tokens, output_tokens)
        self.context_window.add_usage(input_tokens + output_tokens)
        self.message_count += 1

    def estimate_message_tokens(self, content: str) -> int:
        """Estimate tokens in a message.

        Rough estimate: ~4 characters per token for English text.

        Args:
            content: Message content.

        Returns:
            Estimated token count.
        """
        return len(content) // 4

    @property
    def is_warning(self) -> bool:
        """Whether context usage is above warning threshold."""
        return self.context_window.is_warning

    @property
    def is_full(self) -> bool:
        """Whether context is effectively full."""
        return self.context_window.is_full

    @property
    def usage_percentage(self) -> float:
        """Current context usage percentage."""
        return self.context_window.usage_percentage

    def get_status(self) -> dict[str, Any]:
        """Get current context status.

        Returns:
            Status dictionary with usage information.
        """
        return {
            "model": self.model,
            "message_count": self.message_count,
            "input_tokens": self.token_usage.input_tokens,
            "output_tokens": self.token_usage.output_tokens,
            "total_tokens": self.token_usage.total,
            "context_used": self.context_window.used_tokens,
            "context_max": self.context_window.max_tokens,
            "context_available": self.context_window.available,
            "usage_percentage": self.usage_percentage,
            "is_warning": self.is_warning,
            "is_full": self.is_full,
        }

    def get_status_line(self) -> str:
        """Get a single-line status summary.

        Returns:
            Status string like "42K/200K (21%) - 15 messages"
        """
        used_k = self.context_window.used_tokens // 1000
        max_k = self.context_window.max_tokens // 1000
        pct = self.usage_percentage * 100

        status = f"{used_k}K/{max_k}K ({pct:.0f}%) - {self.message_count} messages"

        if self.is_full:
            status = "⚠️ FULL - " + status
        elif self.is_warning:
            status = "⚡ " + status

        return status

    def reset(self) -> None:
        """Reset all tracking."""
        self.token_usage.reset()
        self.context_window.reset()
        self.message_count = 0
