"""BackendManager -- single point of truth for all LLMHosts inference backends.

Manages Ollama (local) and LiteLLM (cloud) backends, providing unified
model listing, health tracking, and request routing.
"""

from __future__ import annotations

import logging
from datetime import datetime, timezone
from enum import Enum
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, ConfigDict, Field

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from llmhosts.backends.litellm_backend import LiteLLMBackend
    from llmhosts.discovery.ollama import OllamaDiscovery

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Backend type enumeration
# ---------------------------------------------------------------------------


class BackendType(str, Enum):
    """Known backend types."""

    OLLAMA = "ollama"
    LITELLM = "litellm"


# ---------------------------------------------------------------------------
# BackendInfo -- Pydantic v2 model
# ---------------------------------------------------------------------------


class BackendInfo(BaseModel):
    """Status and metadata for a registered backend."""

    model_config = ConfigDict(frozen=False)

    name: str
    backend_type: BackendType
    healthy: bool = True
    models: list[str] = Field(default_factory=list)
    provider_count: int = 0
    last_checked: datetime | None = None


# ---------------------------------------------------------------------------
# BackendManager
# ---------------------------------------------------------------------------


class BackendManager:
    """Manages all available backends: Ollama (local) + LiteLLM (cloud).

    Single point of truth for what backends exist, their health, and their
    capabilities.  Used by :class:`BackendDispatcher` to route requests.

    Usage::

        manager = BackendManager()
        await manager.initialize(ollama=ollama_discovery, litellm=litellm_backend)
        result = await manager.completion("litellm", "gpt-4o", messages=[...])
    """

    def __init__(self) -> None:
        self._ollama: OllamaDiscovery | None = None
        self._litellm: LiteLLMBackend | None = None
        self._backends: dict[str, BackendInfo] = {}

    # -- lifecycle -----------------------------------------------------------

    async def initialize(
        self,
        ollama: OllamaDiscovery | None = None,
        litellm: LiteLLMBackend | None = None,
    ) -> None:
        """Register all available backends and probe their health.

        Parameters
        ----------
        ollama:
            An initialised :class:`OllamaDiscovery` instance, or *None* to
            skip local backend registration.
        litellm:
            An initialised :class:`LiteLLMBackend` instance, or *None* to
            skip cloud backend registration.
        """
        self._ollama = ollama
        self._litellm = litellm

        if ollama is not None:
            await self._register_ollama(ollama)

        if litellm is not None:
            self._register_litellm(litellm)

        logger.info(
            "BackendManager ready -- %d backend(s): %s",
            len(self._backends),
            ", ".join(self._backends.keys()) or "none",
        )

    # -- completion routing --------------------------------------------------

    async def completion(
        self,
        backend_type: str,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool = False,
        **kwargs: Any,
    ) -> dict[str, Any] | AsyncIterator[dict[str, Any]]:
        """Route a completion request to the correct backend.

        Parameters
        ----------
        backend_type:
            ``"ollama"`` or ``"litellm"``.
        model:
            Model identifier (e.g. ``"llama3.2:8b"`` or ``"gpt-4o"``).
        messages:
            Conversation messages.
        stream:
            Whether to stream the response.
        **kwargs:
            Additional parameters forwarded to the backend.

        Returns:
            A response dict (non-streaming) or :class:`AsyncIterator` of
            chunk dicts (streaming).

        Raises:
            ValueError: If *backend_type* is not registered or the backend
                is unhealthy.
        """
        bt = backend_type.lower()

        if bt == BackendType.OLLAMA:
            return await self._completion_ollama(model, messages, stream, **kwargs)

        if bt == BackendType.LITELLM:
            return await self._completion_litellm(model, messages, stream, **kwargs)

        raise ValueError(f"Unknown backend type: {backend_type!r}")

    # -- model listing -------------------------------------------------------

    async def list_all_models(self) -> dict[str, list[str]]:
        """List all available models grouped by backend type.

        Returns:
            A dict mapping backend type name → list of model ID strings.
        """
        result: dict[str, list[str]] = {}

        if self._ollama is not None and BackendType.OLLAMA in self._backends:
            info = self._backends[BackendType.OLLAMA]
            result["ollama"] = list(info.models)

        if self._litellm is not None and BackendType.LITELLM in self._backends:
            info = self._backends[BackendType.LITELLM]
            result["litellm"] = list(info.models)

        return result

    # -- backend queries -----------------------------------------------------

    def get_backend_info(self, backend_type: str) -> BackendInfo | None:
        """Get info for a specific backend.

        Parameters
        ----------
        backend_type:
            ``"ollama"`` or ``"litellm"``.

        Returns:
            :class:`BackendInfo` or *None* if not registered.
        """
        return self._backends.get(backend_type.lower())

    def get_all_backends(self) -> dict[str, BackendInfo]:
        """Get all registered backends.

        Returns:
            A dict mapping backend type name → :class:`BackendInfo`.
        """
        return dict(self._backends)

    def has_backend(self, backend_type: str) -> bool:
        """Return *True* if a backend of this type is registered and healthy."""
        info = self._backends.get(backend_type.lower())
        return info is not None and info.healthy

    @property
    def has_cloud(self) -> bool:
        """Return *True* if a cloud (LiteLLM) backend is available."""
        return self.has_backend(BackendType.LITELLM)

    @property
    def has_local(self) -> bool:
        """Return *True* if a local (Ollama) backend is available."""
        return self.has_backend(BackendType.OLLAMA)

    # -- internal: registration ----------------------------------------------

    async def _register_ollama(self, ollama: OllamaDiscovery) -> None:
        """Probe Ollama and register as a backend."""
        available = await ollama.is_available()
        models: list[str] = []
        if available:
            ollama_models = await ollama.list_models()
            models = [m.name for m in ollama_models]

        self._backends[BackendType.OLLAMA] = BackendInfo(
            name="ollama-local",
            backend_type=BackendType.OLLAMA,
            healthy=available and len(models) > 0,
            models=models,
            last_checked=datetime.now(tz=timezone.utc),
        )
        logger.info(
            "Registered Ollama backend: available=%s, models=%d",
            available,
            len(models),
        )

    def _register_litellm(self, litellm_backend: LiteLLMBackend) -> None:
        """Register the LiteLLM cloud backend."""
        providers = litellm_backend.configured_providers
        # We don't eagerly enumerate all cloud models (there are thousands);
        # we just record that the backend is available with its providers.
        self._backends[BackendType.LITELLM] = BackendInfo(
            name="litellm-cloud",
            backend_type=BackendType.LITELLM,
            healthy=len(providers) > 0,
            models=[],  # populated lazily via list_models if needed
            provider_count=len(providers),
            last_checked=datetime.now(tz=timezone.utc),
        )
        logger.info(
            "Registered LiteLLM backend: %d provider(s): %s",
            len(providers),
            ", ".join(sorted(providers)) or "none",
        )

    # -- internal: completion dispatch ---------------------------------------

    async def _completion_ollama(
        self,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool,
        **kwargs: Any,
    ) -> dict[str, Any] | AsyncIterator[dict[str, Any]]:
        """Dispatch to Ollama via :class:`OllamaDiscovery.chat`."""
        if self._ollama is None:
            raise ValueError("Ollama backend is not registered")

        info = self._backends.get(BackendType.OLLAMA)
        if info and not info.healthy:
            raise ValueError("Ollama backend is unhealthy")

        # Build options from kwargs
        options: dict[str, Any] = {}
        for key in ("temperature", "top_p", "num_predict"):
            if key in kwargs and kwargs[key] is not None:
                options[key] = kwargs[key]
        # Map max_tokens → num_predict for Ollama
        if "max_tokens" in kwargs and kwargs["max_tokens"] is not None:
            options["num_predict"] = kwargs.pop("max_tokens")

        chat_kwargs: dict[str, Any] = {}
        if "stop" in kwargs and kwargs["stop"] is not None:
            options["stop"] = kwargs["stop"]
        if options:
            chat_kwargs["options"] = options

        result = await self._ollama.chat(model, messages, stream=stream, **chat_kwargs)

        if stream:
            return self._wrap_ollama_stream(result)  # type: ignore[arg-type]

        # Non-streaming: normalize Ollama response to our dict format
        data: dict[str, Any] = result  # type: ignore[assignment]
        msg = data.get("message", {})
        return {
            "content": msg.get("content", ""),
            "finish_reason": data.get("done_reason", "stop") or "stop",
            "model": model,
            "prompt_tokens": data.get("prompt_eval_count", 0) or 0,
            "completion_tokens": data.get("eval_count", 0) or 0,
            "total_tokens": (data.get("prompt_eval_count", 0) or 0) + (data.get("eval_count", 0) or 0),
            "tool_calls": None,
        }

    async def _wrap_ollama_stream(self, stream: AsyncIterator[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]:
        """Normalize Ollama streaming chunks to our standard format."""
        async for chunk in stream:
            msg = chunk.get("message", {})
            yield {
                "content": msg.get("content", ""),
                "finish_reason": "stop" if chunk.get("done", False) else None,
                "model": chunk.get("model", ""),
            }

    async def _completion_litellm(
        self,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool,
        **kwargs: Any,
    ) -> dict[str, Any] | AsyncIterator[dict[str, Any]]:
        """Dispatch to LiteLLM cloud backend."""
        if self._litellm is None:
            raise ValueError("LiteLLM backend is not registered")

        info = self._backends.get(BackendType.LITELLM)
        if info and not info.healthy:
            raise ValueError("LiteLLM backend is unhealthy (no providers configured)")

        return await self._litellm.completion(model, messages, stream=stream, **kwargs)
