"""Model alias resolution and local-equivalence mappings.

Resolves shorthand model names to canonical identifiers and identifies
local models that can substitute for cloud models.
"""

from __future__ import annotations

import logging

from pydantic import BaseModel

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Alias table -- shorthand -> canonical model identifier
# ---------------------------------------------------------------------------

MODEL_ALIASES: dict[str, str] = {
    # OpenAI shorthands
    "gpt-4": "gpt-4o",
    "gpt-4-turbo": "gpt-4-turbo-2024-04-09",
    "gpt-3.5-turbo": "gpt-3.5-turbo-0125",
    "gpt-4o": "gpt-4o",
    "gpt-4o-mini": "gpt-4o-mini",
    # Anthropic shorthands
    "claude-3.5-sonnet": "claude-3-5-sonnet-20241022",
    "claude-3.5-haiku": "claude-3-5-haiku-20241022",
    "claude-3-opus": "claude-3-opus-20240229",
    "claude-sonnet": "claude-3-5-sonnet-20241022",
    "claude-haiku": "claude-3-5-haiku-20241022",
    "claude-opus": "claude-3-opus-20240229",
}

# ---------------------------------------------------------------------------
# Provider detection -- infer provider from canonical model name
# ---------------------------------------------------------------------------

PROVIDER_PREFIXES: dict[str, str] = {
    "gpt-": "openai",
    "o1-": "openai",
    "o3-": "openai",
    "claude-": "anthropic",
    "gemini-": "google",
    "command-": "cohere",
    "mistral-": "mistral",
}

# ---------------------------------------------------------------------------
# Quality equivalence -- cloud model -> local substitutes (best first)
# ---------------------------------------------------------------------------

QUALITY_EQUIVALENTS: dict[str, list[str]] = {
    # GPT-4o-mini class: these local models are close in quality for most tasks
    "gpt-4o-mini": [
        "llama3.2:8b",
        "llama3.1:8b",
        "mistral:7b",
        "qwen2.5:7b",
        "gemma2:9b",
    ],
    # GPT-3.5-turbo class
    "gpt-3.5-turbo-0125": [
        "llama3.2:8b",
        "llama3.1:8b",
        "mistral:7b",
        "phi3:mini",
        "gemma2:2b",
    ],
    # Claude 3.5 Haiku class
    "claude-3-5-haiku-20241022": [
        "llama3.2:8b",
        "llama3.1:8b",
        "mistral:7b",
        "qwen2.5:7b",
    ],
    # Larger cloud models -- no reliable local equivalents yet
    # "gpt-4o": [],
    # "claude-3-5-sonnet-20241022": [],
    # "claude-3-opus-20240229": [],
}

# ---------------------------------------------------------------------------
# Cloud API endpoint URLs per provider
# ---------------------------------------------------------------------------

PROVIDER_URLS: dict[str, str] = {
    "openai": "https://api.openai.com/v1",
    "anthropic": "https://api.anthropic.com/v1",
    "openrouter": "https://openrouter.ai/api/v1",
}


# ---------------------------------------------------------------------------
# ResolvedModel -- output of alias resolution
# ---------------------------------------------------------------------------


class ResolvedModel(BaseModel):
    """Result of resolving a model alias."""

    original: str  # what the user asked for
    canonical: str  # resolved canonical name
    provider: str  # "ollama", "openai", "anthropic", "openrouter", "unknown"
    is_alias: bool  # True if we performed an alias expansion
    local_equivalents: list[str]  # local models that could substitute
    available_locally: bool  # True if the canonical name (or a variant) is in available_models


# ---------------------------------------------------------------------------
# AliasResolver
# ---------------------------------------------------------------------------


class AliasResolver:
    """Resolves model aliases, detects providers, and finds local equivalents.

    Usage::

        resolver = AliasResolver()
        result = resolver.resolve("gpt-4", available_models=["llama3.2:8b", "mistral:7b"])
    """

    def __init__(
        self,
        aliases: dict[str, str] | None = None,
        equivalents: dict[str, list[str]] | None = None,
    ) -> None:
        self._aliases = aliases if aliases is not None else dict(MODEL_ALIASES)
        self._equivalents = equivalents if equivalents is not None else dict(QUALITY_EQUIVALENTS)

    def resolve(self, model: str, available_models: list[str] | None = None) -> ResolvedModel:
        """Resolve a model name to its canonical form and find local equivalents.

        Parameters
        ----------
        model:
            The model name the user requested (may be a shorthand).
        available_models:
            List of models actually available on local backends (Ollama).
            Used to check if the model or a local equivalent is available.

        Returns
        -------
        ResolvedModel:
            Full resolution result with canonical name, provider, and local equivalents.
        """
        available = available_models or []
        available_lower = {m.lower(): m for m in available}

        # Step 1: alias expansion
        canonical = self._aliases.get(model, model)
        is_alias = canonical != model

        # Step 2: detect provider from canonical name
        provider = self._detect_provider(canonical, available)

        # Step 3: check direct local availability
        # Ollama model names may have tags (e.g. "llama3.2:8b") -- check both with and without
        available_locally = self._is_available_locally(canonical, available_lower)

        # Step 4: find local equivalents
        local_equivalents = self._find_local_equivalents(canonical, available_lower)

        return ResolvedModel(
            original=model,
            canonical=canonical,
            provider=provider,
            is_alias=is_alias,
            local_equivalents=local_equivalents,
            available_locally=available_locally,
        )

    def add_alias(self, shorthand: str, canonical: str) -> None:
        """Register a custom alias."""
        self._aliases[shorthand] = canonical

    def add_equivalents(self, cloud_model: str, local_models: list[str]) -> None:
        """Register local equivalents for a cloud model."""
        self._equivalents[cloud_model] = local_models

    def _detect_provider(self, model: str, available_models: list[str]) -> str:
        """Infer provider from model name or availability."""
        model_lower = model.lower()

        # Check if it's available locally (Ollama)
        available_lower = {m.lower() for m in available_models}
        if model_lower in available_lower or any(m.startswith(model_lower.split(":")[0]) for m in available_lower):
            return "ollama"

        # Check known provider prefixes
        for prefix, provider in PROVIDER_PREFIXES.items():
            if model_lower.startswith(prefix):
                return provider

        return "unknown"

    def _is_available_locally(self, model: str, available_lower: dict[str, str]) -> bool:
        """Check if a model (or a close variant) exists locally."""
        model_l = model.lower()

        # Exact match
        if model_l in available_lower:
            return True

        # Match without tag (e.g. "llama3.2" matches "llama3.2:8b")
        base_name = model_l.split(":")[0]
        return any(avail.split(":")[0] == base_name for avail in available_lower)

    def _find_local_equivalents(self, canonical: str, available_lower: dict[str, str]) -> list[str]:
        """Find local models that can substitute for the given cloud model."""
        equivalents = self._equivalents.get(canonical, [])
        result: list[str] = []

        for eq in equivalents:
            eq_lower = eq.lower()
            # Exact match in available
            if eq_lower in available_lower:
                result.append(available_lower[eq_lower])
                continue
            # Partial match (base name without tag)
            eq_base = eq_lower.split(":")[0]
            for avail_lower, avail_original in available_lower.items():
                if avail_lower.split(":")[0] == eq_base:
                    result.append(avail_original)
                    break

        return result
