"""Model registry with capabilities database.

Centralized database of LLM models with their capabilities, context limits,
pricing, and provider information. Supports dynamic registration of custom
models and automatic provider detection.
"""

from __future__ import annotations

from dataclasses import dataclass
from datetime import date
from typing import Any


@dataclass(frozen=True)
class ModelInfo:
    """Information about an LLM model.

    Attributes:
        name: Model identifier.
        provider: Provider name (openai, anthropic, etc.).
        context_window: Maximum context window in tokens.
        max_output_tokens: Maximum output tokens.
        supports_tools: Whether model supports tool/function calling.
        supports_vision: Whether model supports image inputs.
        supports_streaming: Whether model supports streaming responses.
        supports_json_mode: Whether model supports JSON output mode.
        tokenizer_backend: Tokenizer backend to use.
        input_cost_per_1m: Cost per 1M input tokens in USD.
        output_cost_per_1m: Cost per 1M output tokens in USD.
        cached_input_cost_per_1m: Cost per 1M cached input tokens.
        pricing_date: Date pricing was last updated.
        aliases: Alternative names for the model.
        notes: Additional notes about the model.
    """

    name: str
    provider: str
    context_window: int = 128000
    max_output_tokens: int = 4096
    supports_tools: bool = True
    supports_vision: bool = False
    supports_streaming: bool = True
    supports_json_mode: bool = True
    tokenizer_backend: str | None = None
    input_cost_per_1m: float | None = None
    output_cost_per_1m: float | None = None
    cached_input_cost_per_1m: float | None = None
    pricing_date: date | None = None
    aliases: tuple[str, ...] = ()
    notes: str = ""


# Built-in model database
# Pricing as of January 2025 - verify current rates
_MODELS: dict[str, ModelInfo] = {}


def _register_builtin_models() -> None:
    """Register built-in models."""

    # ============================================================
    # OpenAI Models
    # ============================================================

    # GPT-4o family
    _MODELS["gpt-4o"] = ModelInfo(
        name="gpt-4o",
        provider="openai",
        context_window=128000,
        max_output_tokens=16384,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=2.50,
        output_cost_per_1m=10.00,
        cached_input_cost_per_1m=1.25,
        pricing_date=date(2025, 1, 6),
        aliases=("gpt-4o-2024-11-20", "gpt-4o-2024-08-06", "gpt-4o-2024-05-13"),
        notes="Latest GPT-4o with vision and tools",
    )

    _MODELS["gpt-4o-mini"] = ModelInfo(
        name="gpt-4o-mini",
        provider="openai",
        context_window=128000,
        max_output_tokens=16384,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=0.15,
        output_cost_per_1m=0.60,
        cached_input_cost_per_1m=0.075,
        pricing_date=date(2025, 1, 6),
        aliases=("gpt-4o-mini-2024-07-18",),
        notes="Cost-effective GPT-4o variant",
    )

    # o1 reasoning models
    _MODELS["o1"] = ModelInfo(
        name="o1",
        provider="openai",
        context_window=200000,
        max_output_tokens=100000,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=15.00,
        output_cost_per_1m=60.00,
        cached_input_cost_per_1m=7.50,
        pricing_date=date(2025, 1, 6),
        notes="Full reasoning model with extended thinking",
    )

    _MODELS["o1-mini"] = ModelInfo(
        name="o1-mini",
        provider="openai",
        context_window=128000,
        max_output_tokens=65536,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=1.10,
        output_cost_per_1m=4.40,
        cached_input_cost_per_1m=0.55,
        pricing_date=date(2025, 1, 6),
        notes="Fast reasoning model",
    )

    _MODELS["o3-mini"] = ModelInfo(
        name="o3-mini",
        provider="openai",
        context_window=200000,
        max_output_tokens=100000,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=1.10,
        output_cost_per_1m=4.40,
        cached_input_cost_per_1m=0.55,
        pricing_date=date(2025, 1, 6),
        notes="Latest reasoning model",
    )

    # GPT-4 Turbo
    _MODELS["gpt-4-turbo"] = ModelInfo(
        name="gpt-4-turbo",
        provider="openai",
        context_window=128000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=10.00,
        output_cost_per_1m=30.00,
        cached_input_cost_per_1m=5.00,
        pricing_date=date(2025, 1, 6),
        aliases=("gpt-4-turbo-preview", "gpt-4-turbo-2024-04-09"),
        notes="GPT-4 Turbo with vision",
    )

    # GPT-4
    _MODELS["gpt-4"] = ModelInfo(
        name="gpt-4",
        provider="openai",
        context_window=8192,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=30.00,
        output_cost_per_1m=60.00,
        pricing_date=date(2025, 1, 6),
        aliases=("gpt-4-0613",),
        notes="Original GPT-4",
    )

    _MODELS["gpt-4-32k"] = ModelInfo(
        name="gpt-4-32k",
        provider="openai",
        context_window=32768,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=60.00,
        output_cost_per_1m=120.00,
        pricing_date=date(2025, 1, 6),
        notes="Extended context GPT-4",
    )

    # GPT-3.5
    _MODELS["gpt-3.5-turbo"] = ModelInfo(
        name="gpt-3.5-turbo",
        provider="openai",
        context_window=16385,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="tiktoken",
        input_cost_per_1m=0.50,
        output_cost_per_1m=1.50,
        cached_input_cost_per_1m=0.25,
        pricing_date=date(2025, 1, 6),
        aliases=("gpt-3.5-turbo-0125", "gpt-3.5-turbo-1106"),
        notes="Fast and cost-effective",
    )

    # ============================================================
    # Anthropic Models
    # ============================================================

    _MODELS["claude-3-5-sonnet-20241022"] = ModelInfo(
        name="claude-3-5-sonnet-20241022",
        provider="anthropic",
        context_window=200000,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="anthropic",
        input_cost_per_1m=3.00,
        output_cost_per_1m=15.00,
        cached_input_cost_per_1m=0.30,
        pricing_date=date(2025, 1, 6),
        aliases=("claude-3-5-sonnet-latest", "claude-sonnet-4-20250514"),
        notes="Claude 3.5 Sonnet - Best balance of speed and capability",
    )

    _MODELS["claude-3-5-haiku-20241022"] = ModelInfo(
        name="claude-3-5-haiku-20241022",
        provider="anthropic",
        context_window=200000,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="anthropic",
        input_cost_per_1m=0.80,
        output_cost_per_1m=4.00,
        cached_input_cost_per_1m=0.08,
        pricing_date=date(2025, 1, 6),
        aliases=("claude-3-5-haiku-latest",),
        notes="Claude 3.5 Haiku - Fast and cost-effective",
    )

    _MODELS["claude-3-opus-20240229"] = ModelInfo(
        name="claude-3-opus-20240229",
        provider="anthropic",
        context_window=200000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="anthropic",
        input_cost_per_1m=15.00,
        output_cost_per_1m=75.00,
        cached_input_cost_per_1m=1.50,
        pricing_date=date(2025, 1, 6),
        aliases=("claude-3-opus-latest",),
        notes="Claude 3 Opus - Most capable",
    )

    _MODELS["claude-3-haiku-20240307"] = ModelInfo(
        name="claude-3-haiku-20240307",
        provider="anthropic",
        context_window=200000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="anthropic",
        input_cost_per_1m=0.25,
        output_cost_per_1m=1.25,
        cached_input_cost_per_1m=0.03,
        pricing_date=date(2025, 1, 6),
        notes="Claude 3 Haiku - Legacy fast model",
    )

    # ============================================================
    # Google Models
    # ============================================================

    _MODELS["gemini-2.0-flash"] = ModelInfo(
        name="gemini-2.0-flash",
        provider="google",
        context_window=1000000,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="google",
        input_cost_per_1m=0.10,
        output_cost_per_1m=0.40,
        pricing_date=date(2025, 1, 6),
        aliases=("gemini-2.0-flash-exp",),
        notes="Gemini 2.0 Flash - Fast multimodal",
    )

    _MODELS["gemini-1.5-pro"] = ModelInfo(
        name="gemini-1.5-pro",
        provider="google",
        context_window=2000000,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="google",
        input_cost_per_1m=1.25,
        output_cost_per_1m=5.00,
        pricing_date=date(2025, 1, 6),
        aliases=("gemini-1.5-pro-latest",),
        notes="Gemini 1.5 Pro - 2M context window",
    )

    _MODELS["gemini-1.5-flash"] = ModelInfo(
        name="gemini-1.5-flash",
        provider="google",
        context_window=1000000,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=True,
        supports_streaming=True,
        tokenizer_backend="google",
        input_cost_per_1m=0.075,
        output_cost_per_1m=0.30,
        pricing_date=date(2025, 1, 6),
        aliases=("gemini-1.5-flash-latest",),
        notes="Gemini 1.5 Flash - Cost-effective",
    )

    # ============================================================
    # Meta Llama Models (open source)
    # ============================================================

    _MODELS["llama-3.3-70b"] = ModelInfo(
        name="llama-3.3-70b",
        provider="meta",
        context_window=128000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("llama-3.3-70b-instruct", "meta-llama/Llama-3.3-70B-Instruct"),
        notes="Llama 3.3 70B - Open source",
    )

    _MODELS["llama-3.1-405b"] = ModelInfo(
        name="llama-3.1-405b",
        provider="meta",
        context_window=128000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("llama-3.1-405b-instruct", "meta-llama/Llama-3.1-405B-Instruct"),
        notes="Llama 3.1 405B - Largest open source",
    )

    _MODELS["llama-3.1-70b"] = ModelInfo(
        name="llama-3.1-70b",
        provider="meta",
        context_window=128000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("llama-3.1-70b-instruct", "meta-llama/Llama-3.1-70B-Instruct"),
        notes="Llama 3.1 70B",
    )

    _MODELS["llama-3.1-8b"] = ModelInfo(
        name="llama-3.1-8b",
        provider="meta",
        context_window=128000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("llama-3.1-8b-instruct", "meta-llama/Llama-3.1-8B-Instruct"),
        notes="Llama 3.1 8B - Fast and efficient",
    )

    # ============================================================
    # Mistral Models
    # ============================================================

    _MODELS["mistral-large"] = ModelInfo(
        name="mistral-large",
        provider="mistral",
        context_window=128000,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        input_cost_per_1m=2.00,
        output_cost_per_1m=6.00,
        pricing_date=date(2025, 1, 6),
        aliases=("mistral-large-latest",),
        notes="Mistral Large - Best capability",
    )

    _MODELS["mistral-small"] = ModelInfo(
        name="mistral-small",
        provider="mistral",
        context_window=32768,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        input_cost_per_1m=0.20,
        output_cost_per_1m=0.60,
        pricing_date=date(2025, 1, 6),
        aliases=("mistral-small-latest",),
        notes="Mistral Small - Cost-effective",
    )

    _MODELS["mixtral-8x7b"] = ModelInfo(
        name="mixtral-8x7b",
        provider="mistral",
        context_window=32768,
        max_output_tokens=4096,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("mixtral-8x7b-instruct",),
        notes="Mixtral 8x7B - MoE architecture",
    )

    _MODELS["mistral-7b"] = ModelInfo(
        name="mistral-7b",
        provider="mistral",
        context_window=32768,
        max_output_tokens=4096,
        supports_tools=False,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("mistral-7b-instruct",),
        notes="Mistral 7B - Open source",
    )

    # ============================================================
    # DeepSeek Models
    # ============================================================

    _MODELS["deepseek-v3"] = ModelInfo(
        name="deepseek-v3",
        provider="deepseek",
        context_window=128000,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        input_cost_per_1m=0.14,
        output_cost_per_1m=0.28,
        pricing_date=date(2025, 1, 6),
        notes="DeepSeek V3 - High performance, low cost",
    )

    _MODELS["deepseek-coder"] = ModelInfo(
        name="deepseek-coder",
        provider="deepseek",
        context_window=16384,
        max_output_tokens=4096,
        supports_tools=False,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        notes="DeepSeek Coder - Specialized for code",
    )

    # ============================================================
    # Qwen Models
    # ============================================================

    _MODELS["qwen2.5-72b"] = ModelInfo(
        name="qwen2.5-72b",
        provider="alibaba",
        context_window=131072,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("qwen2.5-72b-instruct",),
        notes="Qwen 2.5 72B - Strong multilingual",
    )

    _MODELS["qwen2.5-7b"] = ModelInfo(
        name="qwen2.5-7b",
        provider="alibaba",
        context_window=131072,
        max_output_tokens=8192,
        supports_tools=True,
        supports_vision=False,
        supports_streaming=True,
        tokenizer_backend="huggingface",
        aliases=("qwen2.5-7b-instruct",),
        notes="Qwen 2.5 7B - Efficient",
    )


# Initialize built-in models
_register_builtin_models()

# Build alias lookup
_ALIASES: dict[str, str] = {}
for model_name, info in _MODELS.items():
    for alias in info.aliases:
        _ALIASES[alias.lower()] = model_name


class ModelRegistry:
    """Registry of LLM models and their capabilities.

    Singleton registry providing access to model information.
    Supports built-in models and custom registration.

    Example:
        # Get model info
        info = ModelRegistry.get("gpt-4o")
        print(f"Context: {info.context_window}")

        # Register custom model
        ModelRegistry.register(
            "my-model",
            provider="custom",
            context_window=32000,
        )

        # List models by provider
        openai_models = ModelRegistry.list_models(provider="openai")
    """

    @classmethod
    def get(cls, model: str) -> ModelInfo | None:
        """Get model information.

        Args:
            model: Model name or alias.

        Returns:
            ModelInfo if found, None otherwise.
        """
        model_lower = model.lower()

        # Direct lookup
        if model_lower in _MODELS:
            return _MODELS[model_lower]

        # Alias lookup
        if model_lower in _ALIASES:
            return _MODELS[_ALIASES[model_lower]]

        # Prefix matching
        for name, info in _MODELS.items():
            if model_lower.startswith(name):
                return info

        return None

    @classmethod
    def register(
        cls,
        model: str,
        provider: str,
        context_window: int = 128000,
        **kwargs: Any,
    ) -> ModelInfo:
        """Register a custom model.

        Args:
            model: Model name.
            provider: Provider name.
            context_window: Maximum context window.
            **kwargs: Additional ModelInfo fields.

        Returns:
            Registered ModelInfo.
        """
        info = ModelInfo(
            name=model,
            provider=provider,
            context_window=context_window,
            **kwargs,
        )
        _MODELS[model.lower()] = info

        # Register aliases
        for alias in info.aliases:
            _ALIASES[alias.lower()] = model.lower()

        return info

    @classmethod
    def list_models(
        cls,
        provider: str | None = None,
        supports_tools: bool | None = None,
        supports_vision: bool | None = None,
        min_context: int | None = None,
    ) -> list[ModelInfo]:
        """List models matching criteria.

        Args:
            provider: Filter by provider.
            supports_tools: Filter by tool support.
            supports_vision: Filter by vision support.
            min_context: Minimum context window.

        Returns:
            List of matching ModelInfo.
        """
        results = []
        for info in _MODELS.values():
            if provider and info.provider != provider:
                continue
            if supports_tools is not None and info.supports_tools != supports_tools:
                continue
            if supports_vision is not None and info.supports_vision != supports_vision:
                continue
            if min_context and info.context_window < min_context:
                continue
            results.append(info)
        return results

    @classmethod
    def list_providers(cls) -> list[str]:
        """List all known providers.

        Returns:
            List of provider names.
        """
        return list({info.provider for info in _MODELS.values()})

    @classmethod
    def get_context_limit(cls, model: str, default: int = 128000) -> int:
        """Get context limit for a model.

        Args:
            model: Model name.
            default: Default if model not found.

        Returns:
            Context window size.
        """
        info = cls.get(model)
        return info.context_window if info else default

    @classmethod
    def estimate_cost(
        cls,
        model: str,
        input_tokens: int,
        output_tokens: int,
        cached_tokens: int = 0,
    ) -> float | None:
        """Estimate API cost for a model.

        Args:
            model: Model name.
            input_tokens: Number of input tokens.
            output_tokens: Number of output tokens.
            cached_tokens: Number of cached input tokens.

        Returns:
            Estimated cost in USD, or None if pricing unknown.
        """
        info = cls.get(model)
        if not info or info.input_cost_per_1m is None:
            return None

        input_cost = (input_tokens / 1_000_000) * info.input_cost_per_1m
        output_cost = (output_tokens / 1_000_000) * (info.output_cost_per_1m or 0)

        if cached_tokens and info.cached_input_cost_per_1m:
            # Adjust for cached tokens
            regular_input = input_tokens - cached_tokens
            cached_cost = (cached_tokens / 1_000_000) * info.cached_input_cost_per_1m
            input_cost = (regular_input / 1_000_000) * info.input_cost_per_1m + cached_cost

        return input_cost + output_cost


# Convenience functions
def get_model_info(model: str) -> ModelInfo | None:
    """Get information about a model.

    Args:
        model: Model name or alias.

    Returns:
        ModelInfo if found, None otherwise.
    """
    return ModelRegistry.get(model)


def list_models(
    provider: str | None = None,
    **kwargs: Any,
) -> list[ModelInfo]:
    """List models matching criteria.

    Args:
        provider: Filter by provider.
        **kwargs: Additional filter criteria.

    Returns:
        List of matching ModelInfo.
    """
    return ModelRegistry.list_models(provider=provider, **kwargs)


def register_model(
    model: str,
    provider: str,
    context_window: int = 128000,
    **kwargs: Any,
) -> ModelInfo:
    """Register a custom model.

    Args:
        model: Model name.
        provider: Provider name.
        context_window: Maximum context window.
        **kwargs: Additional ModelInfo fields.

    Returns:
        Registered ModelInfo.
    """
    return ModelRegistry.register(model, provider, context_window, **kwargs)
