from __future__ import annotations

import json
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, List, Optional

CATALOG_FILE = Path(__file__).with_name("pricing_catalog.json")

PROVIDER_ALIASES = {
    "openai": "openai",
    "azure": "azure_openai",
    "azure_openai": "azure_openai",
    "claude": "claude",
    "anthropic": "claude",
    "gemini": "gemini",
    "google": "gemini",
}


@lru_cache(maxsize=1)
def load_pricing_catalog() -> Dict[str, Any]:
    with CATALOG_FILE.open("r", encoding="utf-8") as f:
        data = json.load(f)
    if not isinstance(data, dict):
        return {}
    return data


def normalize_pricing_provider(provider: Optional[str]) -> str:
    key = (provider or "openai").strip().lower()
    return PROVIDER_ALIASES.get(key, key)


def get_provider_block(provider: Optional[str]) -> Dict[str, Any]:
    catalog = load_pricing_catalog()
    key = normalize_pricing_provider(provider)
    block = catalog.get(key)
    if isinstance(block, dict):
        return block
    return {}


def get_provider_tier(provider: Optional[str]) -> str:
    block = get_provider_block(provider)
    tier = block.get("tier")
    if isinstance(tier, str) and tier.strip():
        return tier.strip()
    return "standard"


def _normalize_model_key(model: Optional[str]) -> str:
    return (model or "").strip().lower()


def _model_candidates(models: List[Dict[str, Any]], model_key: str) -> Optional[Dict[str, Any]]:
    exact = None
    by_model: Dict[str, Dict[str, Any]] = {}
    for rec in models:
        if not isinstance(rec, dict):
            continue
        mk = _normalize_model_key(rec.get("model"))
        if not mk:
            continue
        by_model[mk] = rec
        if mk == model_key:
            exact = rec

    if exact is not None:
        return exact

    # Fallback for snapshot-style names (example: gpt-4.1-mini-2025-xx-xx).
    # Longest prefix wins so mini/nano variants beat base families.
    for key in sorted(by_model.keys(), key=len, reverse=True):
        if model_key.startswith(key):
            return by_model[key]
    return None


def get_model_pricing_record(provider: Optional[str], model: Optional[str]) -> Optional[Dict[str, Any]]:
    block = get_provider_block(provider)
    rows = block.get("models")
    if not isinstance(rows, list):
        return None
    key = _normalize_model_key(model)
    if not key:
        return None
    return _model_candidates(rows, key)


def _pick_numeric_rate(value: Any) -> Optional[float]:
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, dict):
        # Prefer standard short-context branch when both are present.
        for preferred in ("lte_200k", "standard", "default"):
            if isinstance(value.get(preferred), (int, float)):
                return float(value[preferred])
        for v in value.values():
            if isinstance(v, (int, float)):
                return float(v)
    return None


def get_numeric_token_rates(provider: Optional[str], model: Optional[str]) -> Optional[Dict[str, float]]:
    rec = get_model_pricing_record(provider, model)
    if not rec:
        return None

    input_rate = _pick_numeric_rate(rec.get("input"))
    output_rate = _pick_numeric_rate(rec.get("output"))
    cached_rate = _pick_numeric_rate(rec.get("cached"))

    if input_rate is None or output_rate is None:
        return None

    return {
        "input": input_rate,
        "output": output_rate,
        "cached_input": cached_rate if cached_rate is not None else 0.0,
    }


def get_model_pricing(provider: Optional[str], model: Optional[str]) -> Dict[str, Any]:
    normalized_provider = normalize_pricing_provider(provider)
    block = get_provider_block(normalized_provider)
    record = get_model_pricing_record(normalized_provider, model)
    return {
        "provider": normalized_provider,
        "model": (model or "").strip(),
        "tier": get_provider_tier(normalized_provider),
        "unit": block.get("unit", "USD_per_1M_tokens"),
        "source": block.get("source"),
        "pricing": record,
    }
