"""Model registry — catalog of all LLM models across providers."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from pi_agents.ai.types import Cost, Model, ModelCost, Usage


def _load_models_data() -> dict[str, dict[str, dict[str, Any]]]:
    """Load model catalog from JSON data file."""
    data_path = Path(__file__).parent / "models_data.json"
    with open(data_path) as f:
        return json.load(f)


def _build_model(data: dict[str, Any]) -> Model:
    """Build a Model dataclass from raw dict data."""
    cost_data = data.get("cost", {})
    return Model(
        id=data["id"],
        name=data["name"],
        api=data["api"],
        provider=data["provider"],
        base_url=data.get("baseUrl", ""),
        reasoning=data.get("reasoning", False),
        input=data.get("input", ["text", "image"]),
        cost=ModelCost(
            input=cost_data.get("input", 0),
            output=cost_data.get("output", 0),
            cache_read=cost_data.get("cacheRead", 0),
            cache_write=cost_data.get("cacheWrite", 0),
        ),
        context_window=data.get("contextWindow", 0),
        max_tokens=data.get("maxTokens", 0),
        headers=data.get("headers"),
        compat=None,  # compat parsed separately if needed
    )


# Initialize registry on module load
_MODELS_RAW = _load_models_data()
_MODEL_REGISTRY: dict[str, dict[str, Model]] = {}

for _provider, _models in _MODELS_RAW.items():
    _provider_models: dict[str, Model] = {}
    for _model_id, _model_data in _models.items():
        _provider_models[_model_id] = _build_model(_model_data)
    _MODEL_REGISTRY[_provider] = _provider_models


def get_providers() -> list[str]:
    """Get a list of all available LLM providers."""
    return list(_MODEL_REGISTRY.keys())


def get_models(provider: str) -> list[Model]:
    """Get all models available for a specific provider."""
    models = _MODEL_REGISTRY.get(provider)
    return list(models.values()) if models else []


def get_model(provider: str, model_id: str) -> Model | None:
    """Get a specific model definition by provider name and model ID."""
    provider_models = _MODEL_REGISTRY.get(provider)
    if provider_models is None:
        return None
    return provider_models.get(model_id)


def calculate_cost(model: Model, usage: Usage) -> Cost:
    """Calculate token costs for a model based on usage data.

    Formula: cost = (model_cost_per_million / 1_000_000) * token_count
    """
    usage.cost.input = (model.cost.input / 1_000_000) * usage.input
    usage.cost.output = (model.cost.output / 1_000_000) * usage.output
    usage.cost.cache_read = (model.cost.cache_read / 1_000_000) * usage.cache_read
    usage.cost.cache_write = (model.cost.cache_write / 1_000_000) * usage.cache_write
    usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cache_read + usage.cost.cache_write
    return usage.cost


def supports_xhigh(model: Model) -> bool:
    """Check if a model supports the 'xhigh' thinking level.

    Supported by: GPT-5.2+, GPT-5.3, GPT-5.4, Opus 4.6
    """
    if "gpt-5.2" in model.id or "gpt-5.3" in model.id or "gpt-5.4" in model.id:
        return True
    if "opus-4-6" in model.id or "opus-4.6" in model.id:
        return True
    return False


def models_are_equal(a: Model | None, b: Model | None) -> bool:
    """Check if two models are equal by comparing their id and provider."""
    if a is None or b is None:
        return False
    return a.id == b.id and a.provider == b.provider
