"""
LLM cost estimation table for plyra-trace.

Provides cost-per-token lookup and fuzzy model name matching.
When llm.cost_usd is not set but model + token counts are available,
this module estimates the cost.

All costs are in USD per 1M tokens: (input_per_1m, output_per_1m).
"""

from __future__ import annotations

import re
from typing import NamedTuple


class ModelCost(NamedTuple):
    input_per_1m: float  # USD per 1M input tokens
    output_per_1m: float  # USD per 1M output tokens


# ── Cost table {normalized_model_key: ModelCost} ──────────────────────────────
# Sources: official pricing pages (as of 2026-02)
COST_TABLE: dict[str, ModelCost] = {
    # ── OpenAI ────────────────────────────────────────────────────────────────
    "gpt-4o": ModelCost(2.50, 10.00),
    "gpt-4o-mini": ModelCost(0.15, 0.60),
    "gpt-4-turbo": ModelCost(10.00, 30.00),
    "gpt-4": ModelCost(30.00, 60.00),
    "gpt-3.5-turbo": ModelCost(0.50, 1.50),
    "o1": ModelCost(15.00, 60.00),
    "o1-mini": ModelCost(3.00, 12.00),
    "o1-preview": ModelCost(15.00, 60.00),
    "o3": ModelCost(10.00, 40.00),
    "o3-mini": ModelCost(1.10, 4.40),
    "o4-mini": ModelCost(1.10, 4.40),
    # ── Anthropic ─────────────────────────────────────────────────────────────
    "claude-opus-4": ModelCost(15.00, 75.00),
    "claude-sonnet-4": ModelCost(3.00, 15.00),
    "claude-3-7-sonnet": ModelCost(3.00, 15.00),
    "claude-3-5-sonnet": ModelCost(3.00, 15.00),
    "claude-3-5-haiku": ModelCost(0.80, 4.00),
    "claude-haiku": ModelCost(0.80, 4.00),
    "claude-3-opus": ModelCost(15.00, 75.00),
    "claude-3-sonnet": ModelCost(3.00, 15.00),
    "claude-3-haiku": ModelCost(0.25, 1.25),
    "claude-2": ModelCost(8.00, 24.00),
    # ── Google ────────────────────────────────────────────────────────────────
    "gemini-2.0-flash": ModelCost(0.10, 0.40),
    "gemini-2.0-pro": ModelCost(1.25, 5.00),
    "gemini-1.5-pro": ModelCost(1.25, 5.00),
    "gemini-1.5-flash": ModelCost(0.075, 0.30),
    "gemini-1.0-pro": ModelCost(0.50, 1.50),
    # ── Groq ──────────────────────────────────────────────────────────────────
    "llama-3.3-70b": ModelCost(0.59, 0.79),
    "llama-3.1-70b": ModelCost(0.59, 0.79),
    "llama-3.1-8b": ModelCost(0.05, 0.08),
    "llama-3-70b": ModelCost(0.59, 0.79),
    "llama-3-8b": ModelCost(0.05, 0.08),
    "mixtral-8x7b": ModelCost(0.24, 0.24),
    "mixtral-8x22b": ModelCost(0.90, 0.90),
    "gemma-7b": ModelCost(0.07, 0.07),
    "gemma2-9b": ModelCost(0.20, 0.20),
    # ── Mistral ───────────────────────────────────────────────────────────────
    "mistral-large": ModelCost(3.00, 9.00),
    "mistral-medium": ModelCost(2.70, 8.10),
    "mistral-small": ModelCost(1.00, 3.00),
    "mistral-7b": ModelCost(0.25, 0.25),
    "mistral-nemo": ModelCost(0.30, 0.30),
    "codestral": ModelCost(1.00, 3.00),
    # ── Cohere ────────────────────────────────────────────────────────────────
    "command-r-plus": ModelCost(2.50, 10.00),
    "command-r": ModelCost(0.15, 0.60),
    "command": ModelCost(1.00, 2.00),
    # ── Together AI ───────────────────────────────────────────────────────────
    "meta-llama/llama-3-70b-chat-hf": ModelCost(0.90, 0.90),
    "meta-llama/llama-3-8b-chat-hf": ModelCost(0.20, 0.20),
    "meta-llama/llama-3.1-70b-instruct": ModelCost(0.88, 0.88),
    # ── Perplexity ────────────────────────────────────────────────────────────
    "sonar-large": ModelCost(1.00, 1.00),
    "sonar-medium": ModelCost(0.60, 0.60),
    # ── Fireworks ─────────────────────────────────────────────────────────────
    "fireworks/llama-v3p1-70b-instruct": ModelCost(0.90, 0.90),
    # ── DeepSeek ──────────────────────────────────────────────────────────────
    "deepseek-chat": ModelCost(0.27, 1.10),
    "deepseek-coder": ModelCost(0.27, 1.10),
    "deepseek-r1": ModelCost(0.55, 2.19),
    # ── OpenAI Embeddings ─────────────────────────────────────────────────────
    "text-embedding-3-small": ModelCost(0.02, 0.0),
    "text-embedding-3-large": ModelCost(0.13, 0.0),
    "text-embedding-ada-002": ModelCost(0.10, 0.0),
}


def _normalize_model_name(model: str) -> str:
    """
    Normalize a model name for matching:
    - Lowercase
    - Strip provider prefixes (azure/, together/, groq/, openai/)
    - Strip version suffixes (-preview, -YYYY-MM-DD)
    - Strip minor version numbers from model strings when possible
    """
    m = model.lower().strip()

    # Strip provider prefixes like "azure/", "together/", "groq/", "openai/"
    # but keep "meta-llama/" as it's part of the model name on Together
    for prefix in (
        "azure/",
        "openai/",
        "groq/",
        "anthropic/",
        "google/",
        "fireworks/accounts/fireworks/models/",
        "perplexity/",
    ):
        if m.startswith(prefix):
            m = m[len(prefix) :]
            break

    # Strip date version suffixes like -2024-05-13
    m = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", m)

    # Strip -preview, -latest, -turbo-preview suffixes
    m = re.sub(r"-(preview|latest|instruct|chat|it)$", "", m)

    return m


def estimate_cost(model_name: str, prompt_tokens: int, completion_tokens: int) -> float | None:
    """
    Estimate LLM cost in USD.

    Args:
        model_name: Any model name string (normalized internally)
        prompt_tokens: Number of input tokens
        completion_tokens: Number of output tokens

    Returns:
        Estimated cost in USD, or None if model not found in table.

    Examples:
        >>> estimate_cost("gpt-4o-mini", 1000, 500)
        0.00045
        >>> estimate_cost("azure/gpt-4o", 1000, 500)   # strips "azure/"
        0.007500000000000001
    """
    normalized = _normalize_model_name(model_name)

    # Exact match first
    if normalized in COST_TABLE:
        cost_entry = COST_TABLE[normalized]
        return (prompt_tokens * cost_entry.input_per_1m + completion_tokens * cost_entry.output_per_1m) / 1_000_000

    # Prefix match: find all table keys that are prefixes of the normalized name
    # e.g. "gpt-4o-mini-2024-07-18" → matches "gpt-4o-mini"
    best_match: str | None = None
    best_len = 0
    for key in COST_TABLE:
        if normalized.startswith(key) and len(key) > best_len:
            best_match = key
            best_len = len(key)

    if best_match:
        cost_entry = COST_TABLE[best_match]
        return (prompt_tokens * cost_entry.input_per_1m + completion_tokens * cost_entry.output_per_1m) / 1_000_000

    # Suffix/substring match: any key that appears in the normalized name
    for key in COST_TABLE:
        if key in normalized:
            cost_entry = COST_TABLE[key]
            return (prompt_tokens * cost_entry.input_per_1m + completion_tokens * cost_entry.output_per_1m) / 1_000_000

    return None
