"""LLM provider protocol, registry, capability routing, and initialization.

This module provides:
- LLMProvider protocol (the contract for all providers)
- Provider registry (register_provider, get_provider)
- Capability-based routing (ModelCapability, get_model_for_capability)
- Provider initialization (init_providers)
"""

from __future__ import annotations

import contextlib
import logging
from enum import StrEnum
from typing import TYPE_CHECKING, Protocol, runtime_checkable

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from core.common import JsonObject
    from core.types.llm import (
        LLMEmbedResult,
        LLMGenerateResult,
        LLMStreamChunk,
    )

logger = logging.getLogger(__name__)


# =============================================================================
# Error
# =============================================================================


class LLMProviderError(RuntimeError):
    """Normalized provider error."""

    def __init__(
        self,
        message: str,
        provider: str,
        model: str | None = None,
        *,
        is_transient: bool = False,
    ):
        self.message = message
        self.provider = provider
        self.model = model
        self.is_transient = is_transient
        super().__init__(f"{provider}: {message}")


# =============================================================================
# Protocol
# =============================================================================


@runtime_checkable
class LLMProvider(Protocol):
    """Protocol for LLM provider adapters.

    The provider.generate() signature accepts dicts for messages/tools
    because it sits at the wire-format boundary. The invoke.py layer converts
    Pydantic models to dicts before calling providers.
    """

    async def generate(
        self,
        *,
        prompt: str,
        model: str,
        system: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
        response_format: str = "text",
        tools: list[JsonObject] | None = None,
        tool_choice: str | JsonObject | None = None,
        messages: list[JsonObject] | None = None,
        images: list[str] | None = None,
    ) -> LLMGenerateResult:
        """Generate text using the LLM."""
        ...

    async def generate_stream(
        self,
        *,
        prompt: str,
        model: str,
        system: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
        response_format: str = "text",
        tools: list[JsonObject] | None = None,
        tool_choice: str | JsonObject | None = None,
        messages: list[JsonObject] | None = None,
        images: list[str] | None = None,
    ) -> AsyncIterator[LLMStreamChunk]:
        """Stream text generation, yielding chunks as they arrive."""
        # Minimal async-generator stub — concrete providers override.
        yield  # type: ignore[misc]  # pragma: no cover

    async def embed(
        self,
        texts: list[str],
        *,
        model: str,
    ) -> LLMEmbedResult:
        """Generate embeddings for texts."""
        ...

    async def aclose(self) -> None:
        """Close provider connections."""
        ...


# =============================================================================
# Capability Routing
# =============================================================================


class ModelCapability(StrEnum):
    """Model capabilities for routing decisions."""

    FAST = "fast"
    REASONING = "reasoning"
    CODE = "code"
    VISION = "vision"
    EMBEDDING = "embedding"
    LOCAL = "local"


_DEFAULT_MODELS: dict[str, str] = {
    "fast": "openai/gpt-4o-mini",
    "reasoning": "openai/gpt-4o",
    "code": "anthropic/claude-3.5-sonnet",
    "vision": "openai/gpt-4o",
    "embedding": "openai/text-embedding-3-small",
    "local": "llama3.1:8b",
}

_runtime_routing: dict[str, str] = {}


def set_runtime_model_config(routing: dict[str, str] | None = None) -> None:
    """Set runtime model routing overrides."""
    if routing:
        _runtime_routing.update(routing)
        logger.info("Runtime model routing updated: %s", routing)


def get_model_for_capability(
    capability: str | ModelCapability,
    *,
    fallback: str | None = None,
) -> str:
    """Get the model for a capability.

    Resolution order:
    1. Runtime routing overrides (set_runtime_model_config)
    2. ``ALBUS_MODEL_{CAPABILITY}`` env var (via Settings)
    3. Hardcoded ``_DEFAULT_MODELS``
    4. Explicit fallback parameter
    5. Ultimate fallback: first available provider's fast model

    Parameters
    ----------
    capability
        The capability needed (e.g., ModelCapability.CODE or "code").
    fallback
        Fallback model if capability not found.

    Returns
    -------
    str
        Model name in provider format (e.g., "openai/gpt-4o").
    """
    from core.config import get_settings

    cap_str = capability.value if isinstance(capability, ModelCapability) else capability.lower()

    if cap_str in _runtime_routing:
        return _runtime_routing[cap_str]

    # Check env-var overrides (ALBUS_MODEL_FAST, ALBUS_MODEL_REASONING, etc.)
    env_override = get_settings().get_model_override(cap_str)
    if env_override:
        return env_override

    model = _DEFAULT_MODELS.get(cap_str)
    if model:
        return model

    if fallback:
        return fallback

    # Smart fallback: prefer a provider that's actually configured.
    if "ollama" in _providers and "openrouter" not in _providers:
        return _DEFAULT_MODELS.get("local", "llama3.1:8b")

    return _runtime_routing.get("fast", _DEFAULT_MODELS.get("fast", "openai/gpt-4o-mini"))


def get_default_model() -> str:
    """Get the default model (fast capability)."""
    return get_model_for_capability(ModelCapability.FAST)


def list_capabilities() -> list[str]:
    """List all configured capabilities."""
    return [c.value for c in ModelCapability]


def get_runtime_model_config() -> dict[str, object]:
    """Get current runtime model configuration.

    Returns
    -------
    dict[str, object]
        Dict with defaults and any runtime overrides.
    """
    return {
        "defaults": dict(_DEFAULT_MODELS),
        "runtime_overrides": dict(_runtime_routing),
        "effective": {
            **_DEFAULT_MODELS,
            **_runtime_routing,
        },
    }


# =============================================================================
# Provider Registry
# =============================================================================

_providers: dict[str, LLMProvider] = {}


def register_provider(name: str, provider: LLMProvider) -> None:
    """Register an LLM provider by name.

    Called during ``init_providers()`` or by user code.
    """
    _providers[name] = provider
    logger.debug("Registered LLM provider: %s", name)


def get_provider(model: str | None = None) -> LLMProvider:
    """Get the appropriate provider for a model.

    Routing priority:
    1. Local models (no slash) -> Ollama
    2. ``openai/*`` models -> Direct OpenAI (if available), else OpenRouter
    3. Everything else -> OpenRouter -> Ollama fallback
    """
    if model:
        model_lower = model.lower()

        is_local = "/" not in model and not model_lower.startswith(
            ("gpt-", "claude", "gemini", "openai/", "anthropic/", "google/"),
        )
        if is_local and "ollama" in _providers:
            return _providers["ollama"]

        if model_lower.startswith("openai/") and "openai_direct" in _providers:
            return _providers["openai_direct"]

    if "openrouter" in _providers:
        return _providers["openrouter"]

    if "ollama" in _providers:
        return _providers["ollama"]

    raise LLMProviderError(
        message="No LLM providers configured. Set OPENAI_API_KEY, OPENROUTER_API_KEY, or run Ollama.",
        provider="none",
        model=model,
        is_transient=False,
    )


def list_available_providers() -> list[str]:
    """List configured provider names."""
    return list(_providers.keys())


async def close_all_providers() -> None:
    """Close all provider connections."""
    for provider in _providers.values():
        with contextlib.suppress(Exception):
            await provider.aclose()


# =============================================================================
# Provider Initialization
# =============================================================================

_INITIALIZED = False


def init_providers() -> None:
    """Initialize providers from environment and register them.

    Priority for ``openai/*`` models:
    1. Direct OpenAI API (if OPENAI_API_KEY set) — avoids OpenRouter billing
    2. OpenRouter (if OPENROUTER_API_KEY set)
    3. Ollama (local fallback)
    """
    from core.config import get_settings

    global _INITIALIZED

    if _INITIALIZED:
        return

    settings = get_settings()
    openrouter_key = settings.openrouter_api_key
    openai_key = settings.openai_api_key
    ollama_host = settings.ollama_host

    # Direct OpenAI (bypasses OpenRouter when available)
    if openai_key:
        from core.llm.providers.openrouter import OpenRouterAdapter

        class _DirectOpenAI(OpenRouterAdapter):
            """Thin wrapper that routes directly to OpenAI API."""

            BASE_URL = "https://api.openai.com/v1"

            def _normalize_model(self, model: str) -> str:
                if model.startswith("openai/"):
                    return model[len("openai/") :]
                return model

        register_provider("openai_direct", _DirectOpenAI(api_key=openai_key))
        logger.debug("Initialized direct OpenAI adapter")

    # OpenRouter (200+ models, fallback for non-direct providers)
    if openrouter_key:
        from core.llm.providers.openrouter import OpenRouterAdapter

        register_provider("openrouter", OpenRouterAdapter(api_key=openrouter_key))
        logger.debug("Initialized OpenRouter adapter")

    # Ollama (local/offline fallback — requires aiohttp)
    try:
        from core.llm.providers.ollama import OllamaAdapter

        register_provider("ollama", OllamaAdapter(host=ollama_host))
        logger.debug("Initialized Ollama adapter (host=%s)", ollama_host)
    except ImportError:
        logger.debug("Ollama adapter not available (aiohttp not installed)")

    _INITIALIZED = True


__all__ = [
    "LLMProvider",
    "LLMProviderError",
    "ModelCapability",
    "close_all_providers",
    "get_default_model",
    "get_model_for_capability",
    "get_provider",
    "get_runtime_model_config",
    "init_providers",
    "list_available_providers",
    "list_capabilities",
    "register_provider",
    "set_runtime_model_config",
]
