"""Shared pricing utility using LiteLLM's model_prices_and_context_window.json.

Source: https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json

Usage:
    from utils.pricing import calc_cost, get_model_pricing, get_pricing_summary

    cost = calc_cost("claude-opus-4-5-20251101", input_tokens=1000000, output_tokens=500000, cached_tokens=200000)
"""

import json
import urllib.request
from pathlib import Path

LITELLM_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
CACHE_PATH = Path("/home/sagar/trace/analysis-14022026/output/litellm_pricing.json")

# Models not in LiteLLM — rates per 1M tokens
# Only models that LiteLLM doesn't have; others are resolved automatically
FALLBACK_PRICING = {
    "gpt-5-codex":          {"input": 2.00, "output": 16.0, "cached": 1.00},
    "gpt-5.3-codex":        {"input": 1.75, "output": 14.0, "cached": 0.175},
}

# Default fallback when model is completely unknown
DEFAULT_PRICING = {"input": 3.0, "output": 15.0, "cached": 0.30}

_litellm_data: dict | None = None
_pricing_cache: dict[str, dict] = {}
_resolution_log: dict[str, str] = {}  # model -> "litellm" | "fallback" | "default"


def _load_litellm() -> dict:
    """Fetch LiteLLM pricing JSON. Cache to disk for offline use."""
    global _litellm_data
    if _litellm_data is not None:
        return _litellm_data

    # Try fetching fresh
    try:
        req = urllib.request.Request(LITELLM_URL, headers={"User-Agent": "qc-trace-analytics"})
        with urllib.request.urlopen(req, timeout=10) as resp:
            _litellm_data = json.loads(resp.read())
        CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
        CACHE_PATH.write_text(json.dumps(_litellm_data))
        print(f"  [pricing] Fetched LiteLLM pricing ({len(_litellm_data)} models)")
        return _litellm_data
    except Exception as e:
        print(f"  [pricing] LiteLLM fetch failed ({e}), trying cache...")

    # Fall back to cached file
    if CACHE_PATH.exists():
        _litellm_data = json.loads(CACHE_PATH.read_text())
        print(f"  [pricing] Loaded cached LiteLLM pricing ({len(_litellm_data)} models)")
        return _litellm_data

    print("  [pricing] No LiteLLM data available, using fallbacks only")
    _litellm_data = {}
    return _litellm_data


def _extract_pricing(entry: dict) -> dict:
    """Extract pricing from a LiteLLM entry (rates per token -> per 1M tokens)."""
    return {
        "input": (entry.get("input_cost_per_token") or 0) * 1e6,
        "output": (entry.get("output_cost_per_token") or 0) * 1e6,
        "cached": (entry.get("cache_read_input_token_cost") or entry.get("input_cost_per_token") or 0) * 1e6,
    }


def get_model_pricing(model: str) -> dict:
    """Get pricing for a model. Returns {"input", "output", "cached"} rates per 1M tokens.

    Resolution order:
    1. Exact match in LiteLLM
    2. Provider-prefixed match (anthropic., openai/, gemini/, vertex_ai/)
    3. Substring match in LiteLLM keys
    4. FALLBACK_PRICING dict
    5. DEFAULT_PRICING
    """
    if model in _pricing_cache:
        return _pricing_cache[model]

    # Skip synthetic/unknown/non-string models
    if not model or not isinstance(model, str) or model == "<synthetic>":
        _pricing_cache[model] = DEFAULT_PRICING
        _resolution_log[model] = "default"
        return DEFAULT_PRICING

    data = _load_litellm()

    # 1. Exact match
    if model in data:
        result = _extract_pricing(data[model])
        _pricing_cache[model] = result
        _resolution_log[model] = "litellm"
        return result

    # 2. Provider-prefixed match
    prefixes = [f"anthropic/{model}", f"openai/{model}", f"gemini/{model}",
                f"vertex_ai/{model}", f"anthropic.{model}"]
    for key in prefixes:
        if key in data:
            result = _extract_pricing(data[key])
            _pricing_cache[model] = result
            _resolution_log[model] = "litellm"
            return result

    # 3. Substring match — find keys containing the model name
    candidates = [k for k in data if model in k]
    if candidates:
        # Prefer shortest key (most specific match)
        best = min(candidates, key=len)
        result = _extract_pricing(data[best])
        _pricing_cache[model] = result
        _resolution_log[model] = "litellm"
        return result

    # 4. Fallback dict
    if model in FALLBACK_PRICING:
        _pricing_cache[model] = FALLBACK_PRICING[model]
        _resolution_log[model] = "fallback"
        return FALLBACK_PRICING[model]

    # 5. Default
    _pricing_cache[model] = DEFAULT_PRICING
    _resolution_log[model] = "default"
    return DEFAULT_PRICING


def calc_cost(model: str, input_tokens: int | float, output_tokens: int | float,
              cached_tokens: int | float = 0) -> float:
    """Calculate cost in USD for a model usage.

    Args:
        model: Model name (e.g. "claude-opus-4-5-20251101")
        input_tokens: Total input tokens billed
        output_tokens: Output tokens (including thinking tokens if applicable)
        cached_tokens: Tokens served from cache (subtracted from input for billing)
    """
    p = get_model_pricing(str(model) if model else "")
    inp = (input_tokens or 0) / 1e6
    cached = (cached_tokens or 0) / 1e6
    uncached = max(inp - cached, 0)
    out = (output_tokens or 0) / 1e6
    return uncached * p["input"] + cached * p["cached"] + out * p["output"]


def get_pricing_summary() -> dict:
    """Return a summary of pricing resolution for cost_methodology reporting."""
    return {
        "models_resolved": dict(_resolution_log),
        "litellm_count": sum(1 for v in _resolution_log.values() if v == "litellm"),
        "fallback_count": sum(1 for v in _resolution_log.values() if v == "fallback"),
        "default_count": sum(1 for v in _resolution_log.values() if v == "default"),
        "model_rates": {model: get_model_pricing(model) for model in _resolution_log},
    }


def get_pricing_table() -> list[dict]:
    """Return a table of all resolved model prices for display."""
    return [
        {"model": model, "source": _resolution_log.get(model, "unknown"), **get_model_pricing(model)}
        for model in sorted(_resolution_log.keys())
        if model and model != "<synthetic>"
    ]
