"""
LLM pricing database — per 1K tokens, USD.

Updated: March 2026 from official provider documentation.
Sources:
  - OpenAI: https://platform.openai.com/docs/models
  - Anthropic: https://docs.anthropic.com/en/docs/about-claude/models
  - Google: https://ai.google.dev/pricing
  - Mistral: https://docs.mistral.ai/getting-started/models/
  - Others: respective provider pricing pages

The SDK uses these for local cost estimation (no network call needed).
"""

from typing import Dict, Tuple

# Format: "model_id": {"input": price_per_1k, "output": price_per_1k, "provider": "...", "tier": 1-3}
# tier: 1=premium, 2=standard, 3=budget

PRICING: Dict[str, dict] = {
    # ── OpenAI (March 2026 — from platform.openai.com/docs/models) ──
    # GPT-5.4 family (latest flagship)
    "gpt-5.4": {"input": 0.0025, "output": 0.015, "provider": "openai", "tier": 1},
    "gpt-5.4-mini": {"input": 0.00075, "output": 0.0045, "provider": "openai", "tier": 2},
    "gpt-5.4-nano": {"input": 0.0002, "output": 0.00125, "provider": "openai", "tier": 3},
    # GPT-4.1 family
    "gpt-4.1": {"input": 0.002, "output": 0.008, "provider": "openai", "tier": 1},
    "gpt-4.1-mini": {"input": 0.0004, "output": 0.0016, "provider": "openai", "tier": 3},
    "gpt-4.1-nano": {"input": 0.0001, "output": 0.0004, "provider": "openai", "tier": 3},
    # GPT-4o family (still widely used)
    "gpt-4o": {"input": 0.005, "output": 0.015, "provider": "openai", "tier": 1},
    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006, "provider": "openai", "tier": 3},
    # GPT-4 legacy
    "gpt-4": {"input": 0.03, "output": 0.06, "provider": "openai", "tier": 1},
    "gpt-4-turbo": {"input": 0.01, "output": 0.03, "provider": "openai", "tier": 1},
    "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015, "provider": "openai", "tier": 3},
    # o-series reasoning models
    "o1": {"input": 0.015, "output": 0.06, "provider": "openai", "tier": 1},
    "o1-mini": {"input": 0.003, "output": 0.012, "provider": "openai", "tier": 2},
    "o1-pro": {"input": 0.15, "output": 0.60, "provider": "openai", "tier": 1},
    "o3": {"input": 0.01, "output": 0.04, "provider": "openai", "tier": 1},
    "o3-mini": {"input": 0.0011, "output": 0.0044, "provider": "openai", "tier": 2},
    "o3-pro": {"input": 0.02, "output": 0.08, "provider": "openai", "tier": 1},
    "o4-mini": {"input": 0.0011, "output": 0.0044, "provider": "openai", "tier": 2},

    # ── Anthropic (March 2026 — from docs.anthropic.com/en/docs/about-claude/models) ──
    # Claude 4.6 family (latest)
    "claude-opus-4-6": {"input": 0.005, "output": 0.025, "provider": "anthropic", "tier": 1},
    "claude-sonnet-4-6": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-haiku-4-5": {"input": 0.001, "output": 0.005, "provider": "anthropic", "tier": 3},
    "claude-haiku-4-5-20251001": {"input": 0.001, "output": 0.005, "provider": "anthropic", "tier": 3},
    # Claude 4.5 legacy
    "claude-sonnet-4-5": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-sonnet-4-5-20250929": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-opus-4-5": {"input": 0.005, "output": 0.025, "provider": "anthropic", "tier": 1},
    "claude-opus-4-5-20251101": {"input": 0.005, "output": 0.025, "provider": "anthropic", "tier": 1},
    # Claude 4.1 / 4.0 legacy
    "claude-opus-4-1": {"input": 0.015, "output": 0.075, "provider": "anthropic", "tier": 1},
    "claude-opus-4-1-20250805": {"input": 0.015, "output": 0.075, "provider": "anthropic", "tier": 1},
    "claude-sonnet-4": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-sonnet-4-0": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-sonnet-4-20250514": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-opus-4": {"input": 0.015, "output": 0.075, "provider": "anthropic", "tier": 1},
    "claude-opus-4-0": {"input": 0.015, "output": 0.075, "provider": "anthropic", "tier": 1},
    "claude-opus-4-20250514": {"input": 0.015, "output": 0.075, "provider": "anthropic", "tier": 1},
    # Claude 3.x legacy (still used)
    "claude-3.5-sonnet-20241022": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-3.5-haiku-20241022": {"input": 0.0008, "output": 0.004, "provider": "anthropic", "tier": 3},
    "claude-3-opus-20240229": {"input": 0.015, "output": 0.075, "provider": "anthropic", "tier": 1},
    "claude-3-haiku-20240307": {"input": 0.00025, "output": 0.00125, "provider": "anthropic", "tier": 3},
    # Short aliases
    "claude-3.5-sonnet": {"input": 0.003, "output": 0.015, "provider": "anthropic", "tier": 2},
    "claude-3.5-haiku": {"input": 0.0008, "output": 0.004, "provider": "anthropic", "tier": 3},
    "claude-3-opus": {"input": 0.015, "output": 0.075, "provider": "anthropic", "tier": 1},

    # ── Google (from ai.google.dev/pricing) ──────────────────────────
    "gemini-2.5-pro": {"input": 0.00125, "output": 0.01, "provider": "google", "tier": 1},
    "gemini-2.5-flash": {"input": 0.00015, "output": 0.00035, "provider": "google", "tier": 3},
    "gemini-2.0-flash": {"input": 0.0001, "output": 0.0004, "provider": "google", "tier": 3},
    "gemini-2.0-flash-lite": {"input": 0.000075, "output": 0.0003, "provider": "google", "tier": 3},
    "gemini-1.5-pro": {"input": 0.00125, "output": 0.005, "provider": "google", "tier": 2},
    "gemini-1.5-flash": {"input": 0.000075, "output": 0.0003, "provider": "google", "tier": 3},
    "gemini-1.0-pro": {"input": 0.0005, "output": 0.0015, "provider": "google", "tier": 2},

    # ── Mistral (from docs.mistral.ai) ───────────────────────────────
    "mistral-large-latest": {"input": 0.002, "output": 0.006, "provider": "mistral", "tier": 1},
    "mistral-medium-latest": {"input": 0.0027, "output": 0.0081, "provider": "mistral", "tier": 2},
    "mistral-small-latest": {"input": 0.001, "output": 0.003, "provider": "mistral", "tier": 2},
    "open-mistral-nemo": {"input": 0.00015, "output": 0.00015, "provider": "mistral", "tier": 3},
    "codestral-latest": {"input": 0.001, "output": 0.003, "provider": "mistral", "tier": 2},
    "pixtral-large-latest": {"input": 0.002, "output": 0.006, "provider": "mistral", "tier": 1},
    "ministral-8b-latest": {"input": 0.0001, "output": 0.0001, "provider": "mistral", "tier": 3},

    # ── Cohere ──────────────────────────────────────────────────────
    "command-r-plus": {"input": 0.003, "output": 0.015, "provider": "cohere", "tier": 1},
    "command-r": {"input": 0.0005, "output": 0.0015, "provider": "cohere", "tier": 2},
    "command-light": {"input": 0.0003, "output": 0.0006, "provider": "cohere", "tier": 3},
    "command-a": {"input": 0.0025, "output": 0.01, "provider": "cohere", "tier": 1},

    # ── DeepSeek ────────────────────────────────────────────────────
    "deepseek-chat": {"input": 0.00014, "output": 0.00028, "provider": "deepseek", "tier": 3},
    "deepseek-reasoner": {"input": 0.00055, "output": 0.00219, "provider": "deepseek", "tier": 2},

    # ── Meta / Llama (via Together, Groq, Fireworks, etc.) ──────────
    "llama-3.3-70b": {"input": 0.0006, "output": 0.0006, "provider": "meta", "tier": 2},
    "llama-3.1-405b": {"input": 0.003, "output": 0.003, "provider": "meta", "tier": 1},
    "llama-3.1-70b": {"input": 0.0006, "output": 0.0006, "provider": "meta", "tier": 2},
    "llama-3.1-8b": {"input": 0.0001, "output": 0.0001, "provider": "meta", "tier": 3},
    "llama-4-scout": {"input": 0.00015, "output": 0.0006, "provider": "meta", "tier": 3},
    "llama-4-maverick": {"input": 0.0003, "output": 0.0006, "provider": "meta", "tier": 2},

    # ── xAI (Grok) ──────────────────────────────────────────────────
    "grok-3": {"input": 0.003, "output": 0.015, "provider": "xai", "tier": 1},
    "grok-3-mini": {"input": 0.0003, "output": 0.0005, "provider": "xai", "tier": 3},
    "grok-2": {"input": 0.002, "output": 0.01, "provider": "xai", "tier": 2},

    # ── Amazon Bedrock (Nova) ───────────────────────────────────────
    "amazon.nova-pro-v1:0": {"input": 0.0008, "output": 0.0032, "provider": "aws", "tier": 2},
    "amazon.nova-lite-v1:0": {"input": 0.00006, "output": 0.00024, "provider": "aws", "tier": 3},
    "amazon.nova-micro-v1:0": {"input": 0.000035, "output": 0.00014, "provider": "aws", "tier": 3},

    # ── AI21 ────────────────────────────────────────────────────────
    "jamba-1.5-large": {"input": 0.002, "output": 0.008, "provider": "ai21", "tier": 1},
    "jamba-1.5-mini": {"input": 0.0002, "output": 0.0004, "provider": "ai21", "tier": 3},
}

# ── Model downgrade map ──────────────────────────────────────────────
# For each expensive model, what's the cheaper alternative?
# Used by the optimizer when task complexity is "simple".
DOWNGRADE_MAP: Dict[str, str] = {
    # OpenAI — GPT-5.4 family
    "gpt-5.4": "gpt-5.4-mini",
    "gpt-5.4-mini": "gpt-5.4-nano",
    # OpenAI — GPT-4.1 family
    "gpt-4.1": "gpt-4.1-mini",
    "gpt-4.1-mini": "gpt-4.1-nano",
    # OpenAI — GPT-4o family
    "gpt-4o": "gpt-4o-mini",
    "gpt-4": "gpt-4o-mini",
    "gpt-4-turbo": "gpt-4o-mini",
    # OpenAI — o-series
    "o1": "o1-mini",
    "o1-mini": "o3-mini",
    "o1-pro": "o1",
    "o3": "o3-mini",
    "o3-pro": "o3",
    "o4-mini": "o3-mini",
    # Anthropic — Claude 4.6 family
    "claude-opus-4-6": "claude-sonnet-4-6",
    "claude-sonnet-4-6": "claude-haiku-4-5",
    # Anthropic — Claude 4.5 family
    "claude-opus-4-5": "claude-sonnet-4-5",
    "claude-opus-4-5-20251101": "claude-sonnet-4-5-20250929",
    "claude-sonnet-4-5": "claude-haiku-4-5",
    "claude-sonnet-4-5-20250929": "claude-haiku-4-5-20251001",
    # Anthropic — Claude 4.1 / 4.0 legacy
    "claude-opus-4-1": "claude-sonnet-4",
    "claude-opus-4-1-20250805": "claude-sonnet-4-20250514",
    "claude-opus-4": "claude-sonnet-4",
    "claude-opus-4-0": "claude-sonnet-4-0",
    "claude-opus-4-20250514": "claude-sonnet-4-20250514",
    "claude-sonnet-4": "claude-3.5-haiku",
    "claude-sonnet-4-0": "claude-3.5-haiku",
    "claude-sonnet-4-20250514": "claude-3.5-haiku-20241022",
    # Anthropic — Claude 3.x legacy
    "claude-3-opus-20240229": "claude-3.5-haiku-20241022",
    "claude-3-opus": "claude-3.5-haiku",
    "claude-3.5-sonnet-20241022": "claude-3.5-haiku-20241022",
    "claude-3.5-sonnet": "claude-3.5-haiku",
    # Google
    "gemini-2.5-pro": "gemini-2.5-flash",
    "gemini-1.5-pro": "gemini-1.5-flash",
    "gemini-1.0-pro": "gemini-1.5-flash",
    # Mistral
    "mistral-large-latest": "mistral-small-latest",
    "mistral-medium-latest": "open-mistral-nemo",
    "pixtral-large-latest": "mistral-small-latest",
    # Cohere
    "command-r-plus": "command-r",
    "command-r": "command-light",
    "command-a": "command-r",
    # DeepSeek
    "deepseek-reasoner": "deepseek-chat",
    # Meta / Llama
    "llama-3.1-405b": "llama-3.3-70b",
    "llama-3.3-70b": "llama-3.1-8b",
    "llama-3.1-70b": "llama-3.1-8b",
    "llama-4-maverick": "llama-4-scout",
    # xAI
    "grok-3": "grok-3-mini",
    "grok-2": "grok-3-mini",
    # AWS Nova
    "amazon.nova-pro-v1:0": "amazon.nova-lite-v1:0",
    "amazon.nova-lite-v1:0": "amazon.nova-micro-v1:0",
    # AI21
    "jamba-1.5-large": "jamba-1.5-mini",
}

# ── Provider detection from model name ────────────────────────────
def detect_provider(model: str) -> str:
    """Detect LLM provider from model name."""
    if model in PRICING:
        return PRICING[model]["provider"]
    model_lower = model.lower()
    if "gpt" in model_lower or model_lower.startswith("o1") or model_lower.startswith("o3") or model_lower.startswith("o4"):
        return "openai"
    if "claude" in model_lower:
        return "anthropic"
    if "gemini" in model_lower:
        return "google"
    if "mistral" in model_lower or "codestral" in model_lower or "pixtral" in model_lower or "ministral" in model_lower:
        return "mistral"
    if "command" in model_lower or "jamba" in model_lower:
        return "cohere"
    if "deepseek" in model_lower:
        return "deepseek"
    if "llama" in model_lower:
        return "meta"
    if "grok" in model_lower:
        return "xai"
    if "amazon.nova" in model_lower:
        return "aws"
    return "unknown"


def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """Estimate cost in USD for a given model and token counts."""
    pricing = PRICING.get(model)
    if not pricing:
        return 0.0
    return (input_tokens / 1000) * pricing["input"] + (output_tokens / 1000) * pricing["output"]


def get_cheapest_alternative(model: str, provider: str = None) -> str:
    """Get the cheapest model from the same provider."""
    if provider is None:
        provider = detect_provider(model)
    candidates = [
        (name, info)
        for name, info in PRICING.items()
        if info["provider"] == provider
    ]
    if not candidates:
        return model
    cheapest = min(candidates, key=lambda x: x[1]["input"] + x[1]["output"])
    return cheapest[0]
