from __future__ import annotations

from typing import List, Optional

# USD per 1K input tokens (approximate, as of early 2025)
_PRICES_PER_1K_INPUT: dict[str, float] = {
    # OpenAI
    "gpt-4o": 0.005,
    "gpt-4o-mini": 0.00015,
    "gpt-4-turbo": 0.01,
    "gpt-4": 0.03,
    "gpt-4-32k": 0.06,
    "gpt-3.5-turbo": 0.0005,
    "gpt-3.5-turbo-16k": 0.001,
    # Anthropic
    "claude-3-5-sonnet-20241022": 0.003,
    "claude-3-5-sonnet-20240620": 0.003,
    "claude-3-5-haiku-20241022": 0.0008,
    "claude-3-opus-20240229": 0.015,
    "claude-3-sonnet-20240229": 0.003,
    "claude-3-haiku-20240307": 0.00025,
    # Aliases
    "claude-3-5-sonnet": 0.003,
    "claude-3-5-haiku": 0.0008,
    "claude-3-opus": 0.015,
    "claude-3-sonnet": 0.003,
    "claude-3-haiku": 0.00025,
}
_DEFAULT_PRICE_PER_1K = 0.005


class BudgetExceededError(Exception):
    """Raised when a request would exceed the configured cost budget."""

    def __init__(self, message: str, estimated_tokens: int, estimated_usd: float, limit_usd: float) -> None:
        super().__init__(message)
        self.estimated_tokens = estimated_tokens
        self.estimated_usd = estimated_usd
        self.limit_usd = limit_usd


def _try_tiktoken(messages: list[dict], model: str) -> Optional[int]:
    try:
        import tiktoken  # type: ignore
        try:
            enc = tiktoken.encoding_for_model(model)
        except KeyError:
            enc = tiktoken.get_encoding("cl100k_base")
        total = 0
        for m in messages:
            content = m.get("content", "")
            if isinstance(content, str):
                total += len(enc.encode(content)) + 4  # role overhead
        return total
    except ImportError:
        return None


def estimate_tokens(messages: List[dict], model: str = "") -> int:
    """Estimate token count. Uses tiktoken if available, falls back to char/4."""
    if model:
        result = _try_tiktoken(messages, model)
        if result is not None:
            return result
    # Fallback: sum character lengths / 4
    total_chars = sum(
        len(m.get("content", "")) if isinstance(m.get("content"), str) else 0
        for m in messages
    )
    return max(1, total_chars // 4)


def estimate_cost(messages: List[dict], model: str = "") -> float:
    """Estimate input cost in USD for the given messages and model."""
    tokens = estimate_tokens(messages, model)
    price = _PRICES_PER_1K_INPUT.get(model, _DEFAULT_PRICE_PER_1K)
    return round((tokens / 1000) * price, 8)
