"""Model pricing for cost estimation.

Provides per-token pricing for supported LLM providers and a utility
function to calculate the dollar cost of a request.
"""

from __future__ import annotations


# Pricing per 1M tokens (USD) — verified 2026-02-12 from official sources
# Format: "provider/model-pattern": {"input": $, "output": $}
MODEL_PRICING: dict[str, dict[str, float]] = {
    # ── OpenAI (platform.openai.com/docs/pricing) ──
    # GPT-5 series (Aug–Dec 2025)
    "openai/gpt-5.2": {"input": 1.75, "output": 14.00},
    "openai/gpt-5.2-pro": {"input": 21.00, "output": 168.00},
    "openai/gpt-5.2-codex": {"input": 1.75, "output": 14.00},
    "openai/gpt-5.1": {"input": 1.25, "output": 10.00},
    "openai/gpt-5.1-codex": {"input": 1.25, "output": 10.00},
    "openai/gpt-5.1-codex-mini": {"input": 0.25, "output": 2.00},
    "openai/gpt-5": {"input": 1.25, "output": 10.00},
    "openai/gpt-5-pro": {"input": 15.00, "output": 120.00},
    "openai/gpt-5-mini": {"input": 0.25, "output": 2.00},
    "openai/gpt-5-nano": {"input": 0.05, "output": 0.40},
    "openai/gpt-5-codex": {"input": 1.25, "output": 10.00},
    # GPT-4 series (retiring from ChatGPT Feb 13, 2026; API still available)
    "openai/gpt-4o": {"input": 2.50, "output": 10.00},
    "openai/gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "openai/gpt-4.1": {"input": 2.00, "output": 8.00},
    "openai/gpt-4.1-mini": {"input": 0.40, "output": 1.60},
    "openai/gpt-4.1-nano": {"input": 0.10, "output": 0.40},
    # o-series reasoning models
    "openai/o1": {"input": 15.00, "output": 60.00},
    "openai/o1-pro": {"input": 150.00, "output": 600.00},
    "openai/o1-mini": {"input": 3.00, "output": 12.00},
    "openai/o3": {"input": 2.00, "output": 8.00},
    "openai/o3-pro": {"input": 20.00, "output": 80.00},
    "openai/o3-mini": {"input": 1.10, "output": 4.40},
    "openai/o3-deep-research": {"input": 10.00, "output": 40.00},
    "openai/o4-mini": {"input": 1.10, "output": 4.40},
    "openai/o4-mini-deep-research": {"input": 2.00, "output": 8.00},
    # Specialized
    "openai/computer-use-preview": {"input": 3.00, "output": 12.00},
    # ── Anthropic (claude.com/pricing, verified Feb 2026) ──
    "anthropic/claude-opus-4-6": {"input": 5.00, "output": 25.00},
    "anthropic/claude-sonnet-4-5": {"input": 3.00, "output": 15.00},
    "anthropic/claude-haiku-4-5": {"input": 1.00, "output": 5.00},
    "anthropic/claude-opus-4-5": {"input": 5.00, "output": 25.00},
    "anthropic/claude-opus-4-1": {"input": 15.00, "output": 75.00},
    "anthropic/claude-opus-4-0": {"input": 15.00, "output": 75.00},
    "anthropic/claude-sonnet-4-0": {"input": 3.00, "output": 15.00},
    "anthropic/claude-sonnet-3-7": {"input": 3.00, "output": 15.00},
    "anthropic/claude-haiku-3-5": {"input": 0.80, "output": 4.00},
    "anthropic/claude-haiku-3": {"input": 0.25, "output": 1.25},
    # ── Google (ai.google.dev/pricing) ──
    "google/gemini-3-pro-preview": {"input": 2.00, "output": 12.00},
    "google/gemini-3-flash-preview": {"input": 0.50, "output": 3.00},
    "google/gemini-2.5-pro": {"input": 1.25, "output": 10.00},
    "google/gemini-2.5-flash": {"input": 0.30, "output": 2.50},
    "google/gemini-2.5-flash-lite": {"input": 0.10, "output": 0.40},
    "google/gemini-2.0-flash": {"input": 0.10, "output": 0.40},
    "google/gemini-2.0-flash-lite": {"input": 0.075, "output": 0.30},
    # ── Groq (groq.com/pricing) ──
    "groq/kimi-k2-0905": {"input": 1.00, "output": 3.00},
    "groq/llama-3.3-70b": {"input": 0.59, "output": 0.79},
    "groq/qwen3-32b": {"input": 0.29, "output": 0.59},
    "groq/llama-4-maverick": {"input": 0.20, "output": 0.60},
    "groq/gpt-oss-120b": {"input": 0.15, "output": 0.60},
    "groq/llama-4-scout": {"input": 0.11, "output": 0.34},
    "groq/gpt-oss-20b": {"input": 0.075, "output": 0.30},
    "groq/llama-3.1-8b": {"input": 0.05, "output": 0.08},
    # ── DeepSeek V3.2 unified pricing (api-docs.deepseek.com) ──
    "deepseek/deepseek-chat": {"input": 0.28, "output": 0.42},
    "deepseek/deepseek-reasoner": {"input": 0.28, "output": 0.42},
    # ── Mistral (mistral.ai/pricing) ──
    "mistral/mistral-large": {"input": 0.50, "output": 1.50},
    "mistral/mistral-medium-3.1": {"input": 0.40, "output": 2.00},
    "mistral/mistral-small-3.2": {"input": 0.06, "output": 0.18},
    "mistral/mistral-small-3.1": {"input": 0.03, "output": 0.11},
    "mistral/mistral-small": {"input": 0.06, "output": 0.18},
    "mistral/codestral": {"input": 0.30, "output": 0.90},
    "mistral/devstral-2": {"input": 0.05, "output": 0.22},
    "mistral/pixtral-12b": {"input": 0.10, "output": 0.10},
    "mistral/mistral-nemo": {"input": 0.02, "output": 0.04},
    # ── Cohere (cohere.com/pricing) ──
    "cohere/command-a": {"input": 2.50, "output": 10.00},
    "cohere/command-r-plus": {"input": 2.50, "output": 10.00},
    "cohere/command-r": {"input": 0.15, "output": 0.60},
    "cohere/command-r7b": {"input": 0.037, "output": 0.15},
    "cohere/aya-expanse": {"input": 0.50, "output": 1.50},
    # ── OpenRouter (pass-through, use underlying model prices) ──
    "openrouter/anthropic/claude-opus-4-6": {"input": 5.00, "output": 25.00},
    "openrouter/anthropic/claude-sonnet-4-5": {"input": 3.00, "output": 15.00},
    "openrouter/openai/gpt-5.2": {"input": 1.75, "output": 14.00},
    "openrouter/openai/gpt-5": {"input": 1.25, "output": 10.00},
    "openrouter/openai/gpt-4o": {"input": 2.50, "output": 10.00},
    "openrouter/openai/o3": {"input": 2.00, "output": 8.00},
    "openrouter/google/gemini-2.5-pro": {"input": 1.25, "output": 10.00},
    "openrouter/google/gemini-3-pro-preview": {"input": 2.00, "output": 12.00},
    # ── Local models (free) ──
    "ollama/*": {"input": 0.00, "output": 0.00},
    "lmstudio/*": {"input": 0.00, "output": 0.00},
    "vllm/*": {"input": 0.00, "output": 0.00},
}

# Fallback: wildcard patterns for provider-level defaults
_PROVIDER_DEFAULTS: dict[str, dict[str, float]] = {
    "ollama": {"input": 0.00, "output": 0.00},
    "lmstudio": {"input": 0.00, "output": 0.00},
    "vllm": {"input": 0.00, "output": 0.00},
}


def get_model_pricing(model_name: str) -> dict[str, float] | None:
    """Look up pricing for a model.

    Tries exact match first, then prefix matching, then provider defaults.

    Args:
        model_name: Model identifier (e.g. "openai/gpt-4o", "ollama/llama3").

    Returns:
        Dict with "input" and "output" prices per 1M tokens, or None.
    """
    if not model_name:
        return None

    normalized = model_name.lower().strip()

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

    # Prefix match (e.g. "openai/gpt-4o-2024-08-06" matches "openai/gpt-4o")
    for pattern, pricing in MODEL_PRICING.items():
        if pattern.endswith("/*"):
            prefix = pattern[:-1]  # "ollama/"
            if normalized.startswith(prefix):
                return pricing
        elif normalized.startswith(pattern):
            return pricing

    # Provider-level default
    provider = normalized.split("/")[0] if "/" in normalized else ""
    if provider in _PROVIDER_DEFAULTS:
        return _PROVIDER_DEFAULTS[provider]

    return None


def calculate_cost(
    model_name: str,
    input_tokens: int,
    output_tokens: int,
) -> float:
    """Calculate the dollar cost for a request.

    Args:
        model_name: Model identifier.
        input_tokens: Number of input/prompt tokens.
        output_tokens: Number of output/completion tokens.

    Returns:
        Cost in USD. Returns 0.0 if pricing is unavailable.
    """
    pricing = get_model_pricing(model_name)
    if pricing is None:
        return 0.0

    input_cost = (input_tokens / 1_000_000) * pricing["input"]
    output_cost = (output_tokens / 1_000_000) * pricing["output"]
    return input_cost + output_cost
