"""Static price table for estimating AI provider spend.

Provider APIs don't return per-site cost, so we apply our own rates and label
the figures "estimated" in the UI. Keep this table current as provider prices
change — it is the single source of truth for the /usage cost math.

Token rates are USD per 1,000,000 tokens (input and output priced separately).
Tavily is priced per search credit (basic search = 1 credit, advanced = 2).

Sources at time of writing:
- Claude prices from the claude-api skill model table (2026): sonnet-4-6 $3/$15,
  haiku-4-5 $1/$5, opus-4-7 $5/$25 per 1M.
- OpenAI + Tavily rates below are best-effort estimates — verify against the
  provider pricing pages and edit here when they change.
"""

from typing import Optional, Tuple

# USD per 1M tokens: {"input": <rate>, "output": <rate>}
OPENAI_PRICING = {
    "gpt-5.5": {"input": 1.25, "output": 10.00},
    "gpt-5.4": {"input": 1.25, "output": 10.00},
    "gpt-5": {"input": 1.25, "output": 10.00},
    "gpt-4.1": {"input": 2.00, "output": 8.00},
    "gpt-4.1-mini": {"input": 0.40, "output": 1.60},
    "gpt-4.1-nano": {"input": 0.10, "output": 0.40},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
}

CLAUDE_PRICING = {
    "claude-sonnet-4-6": {"input": 3.00, "output": 15.00},
    "claude-haiku-4-5-20251001": {"input": 1.00, "output": 5.00},
    "claude-haiku-4-5": {"input": 1.00, "output": 5.00},
    "claude-opus-4-7": {"input": 5.00, "output": 25.00},
    # Retired models still referenced by older content / RETIRED_MODELS remaps.
    "claude-sonnet-4-5-20250929": {"input": 3.00, "output": 15.00},
    "claude-opus-4-6": {"input": 5.00, "output": 25.00},
}

# USD per Tavily search credit.
TAVILY_CREDIT_USD = 0.008

# USD per generated image, keyed by (model, size, quality). Image generation is
# billed per image (by size/quality/count), not by tokens — see the "Tracking
# OpenAI API Usage" research. Estimates below track OpenAI's gpt-image tiers;
# edit as the real rates change. gpt-image-2 reuses the gpt-image-1.5 tiers
# until a distinct table is published.
IMAGE_PRICING = {
    ("gpt-image-1.5", "1024x1024", "low"): 0.011,
    ("gpt-image-1.5", "1024x1024", "medium"): 0.042,
    ("gpt-image-1.5", "1024x1024", "high"): 0.167,
    ("gpt-image-1.5", "1536x1024", "low"): 0.016,
    ("gpt-image-1.5", "1536x1024", "medium"): 0.063,
    ("gpt-image-1.5", "1536x1024", "high"): 0.25,
    ("gpt-image-1.5", "1024x1536", "low"): 0.016,
    ("gpt-image-1.5", "1024x1536", "medium"): 0.063,
    ("gpt-image-1.5", "1024x1536", "high"): 0.25,
    ("gpt-image-2", "1024x1024", "low"): 0.011,
    ("gpt-image-2", "1024x1024", "medium"): 0.042,
    ("gpt-image-2", "1024x1024", "high"): 0.167,
    ("gpt-image-2", "1536x1024", "low"): 0.016,
    ("gpt-image-2", "1536x1024", "medium"): 0.063,
    ("gpt-image-2", "1536x1024", "high"): 0.25,
    ("gpt-image-2", "1024x1536", "low"): 0.016,
    ("gpt-image-2", "1024x1536", "medium"): 0.063,
    ("gpt-image-2", "1024x1536", "high"): 0.25,
}


def image_rate(model: str, size: Optional[str], quality: Optional[str]) -> Optional[float]:
    """USD per image for a (model, size, quality), or None when not in the table."""
    return IMAGE_PRICING.get((model, size or "", (quality or "").lower()))


def token_rates(provider: str, model: str) -> Tuple[Optional[float], Optional[float]]:
    """Return (input_per_1m, output_per_1m) for a provider/model, or (None, None)
    when the model isn't in the table (unknown price → cost treated as 0)."""
    table = OPENAI_PRICING if provider == "OpenAI" else CLAUDE_PRICING if provider == "Claude" else {}
    rate = table.get(model)
    if not rate:
        return None, None
    return rate["input"], rate["output"]


def estimate_cost(
    provider: str,
    model: str,
    *,
    input_tokens: int = 0,
    output_tokens: int = 0,
    tavily_credits: int = 0,
    images: int = 0,
    image_size: Optional[str] = None,
    image_quality: Optional[str] = None,
) -> Tuple[float, bool]:
    """Estimate USD cost for one aggregated usage row.

    Returns (cost, price_known). price_known is False when we had usage to price
    but no matching rate in the table, so the UI can flag "price not set".
    """
    if provider == "Tavily":
        return round(tavily_credits * TAVILY_CREDIT_USD, 6), True

    if images:
        rate = image_rate(model, image_size, image_quality)
        if rate is None:
            return 0.0, False
        return round(images * rate, 6), True

    in_rate, out_rate = token_rates(provider, model)
    if in_rate is None or out_rate is None:
        # No rate on file — cost unknown. Flag it only if tokens were actually used.
        return 0.0, not (input_tokens or output_tokens)

    cost = (input_tokens / 1_000_000) * in_rate + (output_tokens / 1_000_000) * out_rate
    return round(cost, 6), True