"""Model context limit registry.

Provides a lookup table mapping model names to context window sizes,
sourced from litellm's model_prices_and_context_window.json.

Lookup priority (handled by resolve_model_context_limit in model_config.py):
1. PydanticAI runtime model metadata
2. User cache (~/.anima/model_context_windows.json)
3. Embedded snapshot (src/anima/data/model_context_windows.json)
4. config.toml defaults.context_window_tokens
5. DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS (128_000)
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from deal import pre, post


LITELLM_URL = (
    "https://raw.githubusercontent.com/BerriAI/litellm/main/"
    "model_prices_and_context_window.json"
)

USER_CACHE_FILENAME = "model_context_windows.json"


@dataclass(frozen=True)
class ModelInfo:
    """Context window metadata for a model.

    >>> m = ModelInfo(model_id="gpt-4o", context_window=128000, max_output_tokens=16384)
    >>> m.context_window
    128000
    """

    model_id: str
    context_window: int
    max_output_tokens: int | None = None


def load_embedded_registry() -> dict[str, ModelInfo]:
    """Load the embedded model context window snapshot.

    Returns a dict keyed by model ID. Returns empty dict if the
    embedded file is missing or invalid.

    >>> registry = load_embedded_registry()
    >>> isinstance(registry, dict)
    True
    """
    data_dir = Path(__file__).parent.parent / "data"
    snapshot_path = data_dir / "model_context_windows.json"
    if not snapshot_path.exists():
        return {}
    try:
        raw = json.loads(snapshot_path.read_text())
        return _parse_registry_json(raw)
    except (json.JSONDecodeError, OSError):
        return {}


def load_user_cache(cache_dir: Path | None = None) -> dict[str, ModelInfo] | None:
    """Load the user-level model cache from ~/.anima/.

    Returns None if the cache file doesn't exist or is invalid.

    >>> load_user_cache(Path("/nonexistent")) is None
    True
    """
    if cache_dir is None:
        import os
        cache_dir = Path(os.environ.get("ANIMA_STATE", "")) or Path.home() / ".anima"
    cache_path = cache_dir / USER_CACHE_FILENAME
    if not cache_path.exists():
        return None
    try:
        raw = json.loads(cache_path.read_text())
        return _parse_registry_json(raw)
    except (json.JSONDecodeError, OSError):
        return None


@pre(lambda model_id, registry=None: isinstance(model_id, str) and len(model_id) > 0)
@post(lambda result: result > 0)
def get_context_limit(
    model_id: str,
    registry: dict[str, ModelInfo] | None = None,
) -> int:
    """Look up context window size for a model.

    Tries exact match first, then longest prefix match. Falls back
    to DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS if no match found.

    >>> from anima.core.model_registry import get_context_limit, ModelInfo
    >>> reg = {"gpt-4o": ModelInfo("gpt-4o", 128000), "gpt-5.2": ModelInfo("gpt-5.2", 272000)}
    >>> get_context_limit("gpt-5.2", reg)
    272000
    >>> get_context_limit("gpt-4o", reg)
    128000
    >>> get_context_limit("unknown-model", reg)  # falls back to default
    128000
    """
    from anima.shell.model_config import DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS

    if registry is None:
        return DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS

    # Exact match
    if model_id in registry:
        return registry[model_id].context_window

    # Strip provider prefix (e.g., "openai:gpt-5.2" → "gpt-5.2")
    bare_id = model_id.split(":", 1)[-1] if ":" in model_id else model_id
    if bare_id in registry:
        return registry[bare_id].context_window

    # Longest prefix match (e.g., "claude-sonnet-4-20250514" matches "claude-sonnet-4")
    best_match: ModelInfo | None = None
    best_len = 0
    for key, info in registry.items():
        if bare_id.startswith(key) and len(key) > best_len:
            best_match = info
            best_len = len(key)

    if best_match is not None:
        return best_match.context_window

    return DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS


def update_registry_from_litellm(output_path: Path | None = None) -> int:
    """Fetch model data from litellm GitHub and write to user cache.

    Downloads the litellm model_prices_and_context_window.json, extracts
    context window data, and writes to the user cache file.

    Args:
        output_path: Where to write the cache. Defaults to
            ~/.anima/model_context_windows.json.

    Returns:
        Number of models written to the cache.
    """
    import urllib.request

    if output_path is None:
        import os
        cache_dir = Path(os.environ.get("ANIMA_STATE", "")) or Path.home() / ".anima"
        cache_dir.mkdir(parents=True, exist_ok=True)
        output_path = cache_dir / USER_CACHE_FILENAME

    response = urllib.request.urlopen(LITELLM_URL, timeout=30)  # noqa: S310
    raw_data: dict[str, Any] = json.loads(response.read())

    extracted: dict[str, dict[str, int | None]] = {}
    for key, value in raw_data.items():
        if key.startswith("sample_spec"):
            continue
        if not isinstance(value, dict):
            continue
        max_input = value.get("max_input_tokens")
        if not isinstance(max_input, int) or max_input <= 0:
            continue
        max_output = value.get("max_output_tokens")
        extracted[key] = {
            "context_window": max_input,
            "max_output_tokens": max_output if isinstance(max_output, int) else None,
        }

    output_path.write_text(json.dumps(extracted, indent=2, sort_keys=True))
    return len(extracted)


def _parse_registry_json(raw: Any) -> dict[str, ModelInfo]:
    """Parse a registry JSON dict into ModelInfo objects."""
    if not isinstance(raw, dict):
        return {}
    result: dict[str, ModelInfo] = {}
    for key, value in raw.items():
        if not isinstance(value, dict):
            continue
        context_window = value.get("context_window")
        if not isinstance(context_window, int) or context_window <= 0:
            continue
        result[key] = ModelInfo(
            model_id=key,
            context_window=context_window,
            max_output_tokens=value.get("max_output_tokens"),
        )
    return result


__all__ = [
    "ModelInfo",
    "get_context_limit",
    "load_embedded_registry",
    "load_user_cache",
    "update_registry_from_litellm",
    "LITELLM_URL",
    "USER_CACHE_FILENAME",
]
