"""
Token cost estimation for LLM models.

Pricing data is hardcoded and may become stale. Last updated: 2026-01.

Sources:
- https://www.helicone.ai/llm-cost/provider/openai/model/gpt-5.1-codex
- https://pricepertoken.com/pricing-page/model/openai-codex-mini
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass
class ModelPricing:
    """Pricing for a model in $ per million tokens."""
    input_per_million: float
    output_per_million: float
    cached_input_per_million: float | None = None  # Some models have cached input discount

    def estimate_cost(
        self,
        input_tokens: int,
        output_tokens: int,
        cached_tokens: int = 0,
    ) -> float:
        """
        Estimate cost in dollars.

        Args:
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens
            cached_tokens: Number of cached input tokens (if applicable)

        Returns:
            Estimated cost in USD
        """
        input_cost = (input_tokens / 1_000_000) * self.input_per_million
        output_cost = (output_tokens / 1_000_000) * self.output_per_million

        cached_cost = 0.0
        if cached_tokens and self.cached_input_per_million:
            cached_cost = (cached_tokens / 1_000_000) * self.cached_input_per_million

        return input_cost + output_cost + cached_cost


# Model pricing table ($ per million tokens)
# Last updated: 2026-01
MODEL_PRICING: dict[str, ModelPricing] = {
    # OpenAI Codex models
    "gpt-5.1-codex": ModelPricing(
        input_per_million=1.25,
        output_per_million=10.00,
        cached_input_per_million=0.125,  # 90% discount for cached
    ),
    "gpt-5.1-codex-mini": ModelPricing(
        input_per_million=0.25,
        output_per_million=2.00,
        cached_input_per_million=0.025,
    ),
    "gpt-5.1-codex-max": ModelPricing(
        input_per_million=1.25,
        output_per_million=10.00,
        cached_input_per_million=0.125,
    ),
    # GPT-5 base models (for reference)
    "gpt-5": ModelPricing(
        input_per_million=1.25,
        output_per_million=10.00,
    ),
    "gpt-5-mini": ModelPricing(
        input_per_million=0.25,
        output_per_million=2.00,
    ),
    # Claude models (Anthropic)
    "claude-sonnet-4-20250514": ModelPricing(
        input_per_million=3.00,
        output_per_million=15.00,
    ),
    "claude-opus-4-20250514": ModelPricing(
        input_per_million=15.00,
        output_per_million=75.00,
    ),
    "claude-3-5-sonnet-20241022": ModelPricing(
        input_per_million=3.00,
        output_per_million=15.00,
    ),
}

# Aliases for common model names
MODEL_ALIASES: dict[str, str] = {
    "codex": "gpt-5.1-codex",
    "codex-mini": "gpt-5.1-codex-mini",
    "codex-max": "gpt-5.1-codex-max",
    "gpt5": "gpt-5",
    "gpt5-mini": "gpt-5-mini",
    "sonnet": "claude-sonnet-4-20250514",
    "opus": "claude-opus-4-20250514",
}


def get_pricing(model: str) -> ModelPricing | None:
    """
    Get pricing for a model.

    Args:
        model: Model name or alias

    Returns:
        ModelPricing or None if unknown
    """
    # Check aliases first
    resolved = MODEL_ALIASES.get(model.lower(), model)

    # Exact match
    if resolved in MODEL_PRICING:
        return MODEL_PRICING[resolved]

    # Try lowercase
    if resolved.lower() in MODEL_PRICING:
        return MODEL_PRICING[resolved.lower()]

    # Try prefix matching (e.g., "gpt-5.1-codex-mini-2026-01" -> "gpt-5.1-codex-mini")
    for known_model in MODEL_PRICING:
        if resolved.lower().startswith(known_model.lower()):
            return MODEL_PRICING[known_model]

    return None


def estimate_cost(
    model: str,
    input_tokens: int,
    output_tokens: int,
    cached_tokens: int = 0,
) -> float | None:
    """
    Estimate cost for a model run.

    Args:
        model: Model name
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        cached_tokens: Number of cached input tokens

    Returns:
        Cost in USD, or None if model pricing unknown
    """
    pricing = get_pricing(model)
    if pricing is None:
        return None

    return pricing.estimate_cost(input_tokens, output_tokens, cached_tokens)


def format_cost(cost: float | None) -> str:
    """Format cost as a human-readable string."""
    if cost is None:
        return "?"
    if cost < 0.01:
        return f"${cost:.4f}"
    elif cost < 1.00:
        return f"${cost:.3f}"
    else:
        return f"${cost:.2f}"


def estimate_session_cost(
    model: str,
    token_usage: dict[str, Any],
) -> dict[str, Any]:
    """
    Estimate cost for a session given its token usage.

    Args:
        model: Model used
        token_usage: Dict with input_tokens, output_tokens, etc.

    Returns:
        Dict with cost info: {cost, cost_formatted, pricing_known}
    """
    input_tokens = token_usage.get("input_tokens", 0)
    output_tokens = token_usage.get("output_tokens", 0)
    cached_tokens = token_usage.get("cached_tokens", 0)

    cost = estimate_cost(model, input_tokens, output_tokens, cached_tokens)

    return {
        "cost": cost,
        "cost_formatted": format_cost(cost),
        "pricing_known": cost is not None,
        "model": model,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
    }
