"""LiteLLM backend -- unified access to 100+ cloud LLM providers via BYOK keys.

Wraps LiteLLM's async completion interface to provide:

* Automatic retry with exponential backoff
* Token counting across all providers
* Consistent error handling mapped to :class:`ProxyError` hierarchy
* Model aliasing built-in
* Rate limiting awareness
* Streaming via async generators
"""

from __future__ import annotations

import logging
import os
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, ConfigDict

from llmhosts.proxy.errors import (
    AuthenticationError,
    BackendTimeoutError,
    BackendUnavailableError,
    InvalidRequestError,
    ProxyError,
)

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from llmhosts.keys.manager import KeyManager

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Graceful import -- litellm is a core dep but we stay robust
# ---------------------------------------------------------------------------

try:
    import litellm
    from litellm import ModelResponse
    from litellm.exceptions import (
        AuthenticationError as LiteLLMAuthError,
    )
    from litellm.exceptions import (
        BadRequestError as LiteLLMBadRequestError,
    )
    from litellm.exceptions import (
        NotFoundError as LiteLLMNotFoundError,
    )
    from litellm.exceptions import (
        RateLimitError as LiteLLMRateLimitError,
    )
    from litellm.exceptions import (
        ServiceUnavailableError as LiteLLMServiceUnavailableError,
    )
    from litellm.exceptions import (
        Timeout as LiteLLMTimeout,
    )

    LITELLM_AVAILABLE = True
except ImportError:
    LITELLM_AVAILABLE = False
    litellm = None  # type: ignore[assignment]
    ModelResponse = None  # type: ignore[assignment,misc]

    # Placeholder exception classes so type checking doesn't break
    class _PlaceholderError(Exception):
        pass

    LiteLLMAuthError = _PlaceholderError  # type: ignore[assignment,misc]
    LiteLLMBadRequestError = _PlaceholderError  # type: ignore[assignment,misc]
    LiteLLMNotFoundError = _PlaceholderError  # type: ignore[assignment,misc]
    LiteLLMRateLimitError = _PlaceholderError  # type: ignore[assignment,misc]
    LiteLLMServiceUnavailableError = _PlaceholderError  # type: ignore[assignment,misc]
    LiteLLMTimeout = _PlaceholderError  # type: ignore[assignment,misc]


# ---------------------------------------------------------------------------
# Provider → LiteLLM env-var mapping
# ---------------------------------------------------------------------------

_PROVIDER_ENV_MAP: dict[str, str] = {
    "openai": "OPENAI_API_KEY",
    "anthropic": "ANTHROPIC_API_KEY",
    "google": "GEMINI_API_KEY",
    "mistral": "MISTRAL_API_KEY",
    "groq": "GROQ_API_KEY",
    "together": "TOGETHERAI_API_KEY",
    "deepseek": "DEEPSEEK_API_KEY",
    "openrouter": "OPENROUTER_API_KEY",
}

# Provider → LiteLLM model prefix mapping (for provider detection)
_PROVIDER_PREFIX_MAP: dict[str, str] = {
    "openai": "",  # default provider -- no prefix needed
    "anthropic": "anthropic/",
    "google": "gemini/",
    "mistral": "mistral/",
    "groq": "groq/",
    "together": "together_ai/",
    "deepseek": "deepseek/",
    "openrouter": "openrouter/",
}


# ---------------------------------------------------------------------------
# CloudModelInfo -- Pydantic v2 model for cloud model metadata
# ---------------------------------------------------------------------------


class CloudModelInfo(BaseModel):
    """Metadata for a cloud model accessible via LiteLLM."""

    model_config = ConfigDict(frozen=True)

    id: str
    provider: str
    max_tokens: int = 0
    input_cost_per_token: float = 0.0
    output_cost_per_token: float = 0.0
    supports_streaming: bool = True
    supports_tools: bool = False
    supports_vision: bool = False


# ---------------------------------------------------------------------------
# LiteLLMBackend
# ---------------------------------------------------------------------------


class LiteLLMBackend:
    """Wraps LiteLLM for unified access to 100+ cloud providers.

    Benefits over raw httpx:
    - Automatic retry with exponential backoff
    - Token counting across all providers
    - Consistent error handling
    - Model aliasing built-in
    - Rate limiting awareness

    Usage::

        backend = LiteLLMBackend(keys=key_manager)
        providers = await backend.initialize()
        response = await backend.completion("gpt-4o", messages=[...])
    """

    def __init__(self, keys: KeyManager) -> None:
        self._keys = keys
        self._configured_providers: set[str] = set()
        self._initialized: bool = False

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

    async def initialize(self) -> list[str]:
        """Configure LiteLLM with available BYOK keys.

        Reads all stored keys from :class:`KeyManager`, injects them as
        environment variables for LiteLLM, and returns the list of providers
        that are configured and ready to use.

        Returns:
            List of provider names with valid keys configured.

        Raises:
            RuntimeError: If ``litellm`` is not installed.
        """
        if not LITELLM_AVAILABLE:
            raise RuntimeError(
                "litellm is not installed. Install with: pip install 'llmhosts[smart]' or pip install litellm"
            )

        env_vars = self._setup_litellm_keys(self._keys)
        for var_name, var_value in env_vars.items():
            os.environ[var_name] = var_value

        self._configured_providers = {
            provider for provider, env_var in _PROVIDER_ENV_MAP.items() if env_var in env_vars
        }

        # Configure LiteLLM defaults
        litellm.drop_params = True  # silently drop unsupported params
        litellm.set_verbose = False  # keep logs clean

        self._initialized = True
        logger.info(
            "LiteLLM initialized with %d provider(s): %s",
            len(self._configured_providers),
            ", ".join(sorted(self._configured_providers)) or "none",
        )
        return sorted(self._configured_providers)

    # -- completion ----------------------------------------------------------

    async def completion(
        self,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool = False,
        **kwargs: Any,
    ) -> dict[str, Any] | AsyncIterator[dict[str, Any]]:
        """Send a completion request via LiteLLM.

        Handles:
        - Provider detection from model name
        - API key injection (already in env from :meth:`initialize`)
        - Streaming vs non-streaming
        - Token counting
        - Error translation to our :class:`ProxyError` types

        Parameters
        ----------
        model:
            Model identifier.  Can be a bare name like ``"gpt-4o"`` (resolved
            to the default provider) or prefixed like ``"anthropic/claude-3.5-sonnet"``.
        messages:
            List of ``{"role": ..., "content": ...}`` message dicts.
        stream:
            If *True*, return an :class:`AsyncIterator` yielding chunk dicts.
        **kwargs:
            Additional params forwarded to ``litellm.acompletion`` (e.g.
            ``temperature``, ``max_tokens``, ``top_p``, ``stop``).

        Returns:
            A response dict (non-streaming) or an async iterator of chunk dicts
            (streaming).

        Raises:
            RuntimeError: If :meth:`initialize` has not been called.
            ProxyError: Translated from LiteLLM exceptions.
        """
        self._ensure_initialized()

        # Clean up kwargs -- remove None values
        clean_kwargs: dict[str, Any] = {k: v for k, v in kwargs.items() if v is not None}

        if stream:
            return self._stream_completion(model, messages, **clean_kwargs)

        return await self._non_stream_completion(model, messages, **clean_kwargs)

    async def _non_stream_completion(
        self,
        model: str,
        messages: list[dict[str, Any]],
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Execute a non-streaming completion and return a normalized dict."""
        try:
            response: ModelResponse = await litellm.acompletion(  # type: ignore[union-attr]
                model=model,
                messages=messages,
                stream=False,
                **kwargs,
            )
            return self._normalize_response(response)
        except Exception as exc:
            raise self._translate_exception(exc, model=model) from exc

    async def _stream_completion(
        self,
        model: str,
        messages: list[dict[str, Any]],
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        """Execute a streaming completion, yielding normalized chunk dicts."""
        try:
            response = await litellm.acompletion(  # type: ignore[union-attr]
                model=model,
                messages=messages,
                stream=True,
                **kwargs,
            )
            async for chunk in response:
                yield self._normalize_stream_chunk(chunk)
        except Exception as exc:
            raise self._translate_exception(exc, model=model) from exc

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

    async def list_models(self, provider: str | None = None) -> list[CloudModelInfo]:
        """List available models for configured providers.

        Uses LiteLLM's model cost map as the canonical source of known models,
        filtered to only the providers we have keys for.

        Parameters
        ----------
        provider:
            Optional provider name to filter by.  If *None*, returns models
            for all configured providers.

        Returns:
            List of :class:`CloudModelInfo` for available models.
        """
        self._ensure_initialized()

        target_providers = {provider} if provider else self._configured_providers
        models: list[CloudModelInfo] = []

        model_cost_map: dict[str, Any] = getattr(litellm, "model_cost", {})  # type: ignore[union-attr]
        if not model_cost_map:
            logger.debug("LiteLLM model_cost map is empty; cannot enumerate models")
            return models

        for model_name, cost_info in model_cost_map.items():
            if not isinstance(cost_info, dict):
                continue

            model_provider = cost_info.get("litellm_provider", "")
            # Normalize provider name for matching
            normalized = self._normalize_provider_name(model_provider)
            if normalized not in target_providers:
                continue

            models.append(
                CloudModelInfo(
                    id=model_name,
                    provider=normalized,
                    max_tokens=cost_info.get("max_tokens", 0) or 0,
                    input_cost_per_token=cost_info.get("input_cost_per_token", 0.0) or 0.0,
                    output_cost_per_token=cost_info.get("output_cost_per_token", 0.0) or 0.0,
                    supports_streaming=True,
                    supports_tools=cost_info.get("supports_function_calling", False) or False,
                    supports_vision=cost_info.get("supports_vision", False) or False,
                )
            )

        logger.debug("Enumerated %d cloud models for providers: %s", len(models), target_providers)
        return models

    # -- cost calculation ----------------------------------------------------

    def get_model_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Calculate cost using LiteLLM's pricing data.

        Parameters
        ----------
        model:
            The model name (e.g. ``"gpt-4o"``).
        input_tokens:
            Number of input (prompt) tokens.
        output_tokens:
            Number of output (completion) tokens.

        Returns:
            Estimated cost in USD.  Returns ``0.0`` if pricing data is
            unavailable or litellm is not installed.
        """
        if not LITELLM_AVAILABLE:
            return 0.0

        try:
            cost = litellm.completion_cost(  # type: ignore[union-attr]
                model=model,
                prompt_tokens=input_tokens,
                completion_tokens=output_tokens,
            )
            return float(cost)
        except Exception:
            logger.debug("Could not compute cost for model %s", model, exc_info=True)
            return 0.0

    # -- provider queries ----------------------------------------------------

    def is_configured(self, provider: str) -> bool:
        """Check if *provider* has a valid key configured."""
        return provider.strip().lower() in self._configured_providers

    @property
    def configured_providers(self) -> set[str]:
        """The set of provider names with keys configured."""
        return frozenset(self._configured_providers)  # type: ignore[return-value]

    # -- static helpers ------------------------------------------------------

    @staticmethod
    def _setup_litellm_keys(keys: KeyManager) -> dict[str, str]:
        """Map our key storage to LiteLLM environment variables.

        Reads decrypted keys from :class:`KeyManager` and returns a dict
        of environment variable name → key value for all available providers.
        Does **not** mutate ``os.environ``; the caller handles that.
        """
        env_vars: dict[str, str] = {}
        for provider, env_var in _PROVIDER_ENV_MAP.items():
            api_key = keys.get_key(provider)
            if api_key:
                env_vars[env_var] = api_key
                logger.debug("Mapped key for provider '%s' → %s", provider, env_var)
        return env_vars

    # -- internal helpers ----------------------------------------------------

    def _ensure_initialized(self) -> None:
        """Raise if :meth:`initialize` has not been called."""
        if not LITELLM_AVAILABLE:
            raise RuntimeError(
                "litellm is not installed. Install with: pip install 'llmhosts[smart]' or pip install litellm"
            )
        if not self._initialized:
            raise RuntimeError("LiteLLMBackend.initialize() must be called before use")

    @staticmethod
    def _normalize_response(response: Any) -> dict[str, Any]:
        """Convert a LiteLLM ``ModelResponse`` to a plain dict.

        Extracts the fields we care about into a normalized structure
        compatible with our :class:`UnifiedResponse`.
        """
        choices = getattr(response, "choices", [])
        message = choices[0].message if choices else None
        content = getattr(message, "content", "") or "" if message else ""
        finish_reason = getattr(choices[0], "finish_reason", "stop") if choices else "stop"

        usage = getattr(response, "usage", None)
        prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0 if usage else 0
        completion_tokens = getattr(usage, "completion_tokens", 0) or 0 if usage else 0

        # Extract tool calls if present
        tool_calls = None
        if message and getattr(message, "tool_calls", None):
            tool_calls = [
                {
                    "id": tc.id,
                    "type": tc.type,
                    "function": {"name": tc.function.name, "arguments": tc.function.arguments},
                }
                for tc in message.tool_calls
            ]

        return {
            "content": content,
            "finish_reason": finish_reason or "stop",
            "model": getattr(response, "model", ""),
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
            "tool_calls": tool_calls,
        }

    @staticmethod
    def _normalize_stream_chunk(chunk: Any) -> dict[str, Any]:
        """Convert a LiteLLM streaming chunk to a plain dict.

        Returns a dict with ``content`` (the text delta) and ``finish_reason``.
        """
        choices = getattr(chunk, "choices", [])
        delta = choices[0].delta if choices else None
        content = getattr(delta, "content", "") or "" if delta else ""
        finish_reason = getattr(choices[0], "finish_reason", None) if choices else None

        return {
            "content": content,
            "finish_reason": finish_reason,
            "model": getattr(chunk, "model", ""),
        }

    @staticmethod
    def _normalize_provider_name(litellm_provider: str) -> str:
        """Normalize a LiteLLM provider string to our provider names.

        LiteLLM uses names like ``"openai"``, ``"anthropic"``,
        ``"vertex_ai"``, ``"together_ai"``, etc.
        """
        mapping: dict[str, str] = {
            "openai": "openai",
            "anthropic": "anthropic",
            "gemini": "google",
            "vertex_ai": "google",
            "vertex_ai_beta": "google",
            "mistral": "mistral",
            "groq": "groq",
            "together_ai": "together",
            "deepseek": "deepseek",
            "openrouter": "openrouter",
        }
        return mapping.get(litellm_provider.lower(), litellm_provider.lower())

    @staticmethod
    def _translate_exception(exc: Exception, *, model: str = "") -> ProxyError:
        """Map a LiteLLM exception to our :class:`ProxyError` hierarchy.

        Falls through to a generic :class:`ProxyError` for unknown exception
        types.
        """
        # Authentication failures
        if isinstance(exc, LiteLLMAuthError):
            return AuthenticationError(f"Authentication failed for model '{model}': {exc}")

        # Bad request / invalid model
        if isinstance(exc, (LiteLLMBadRequestError, LiteLLMNotFoundError)):
            return InvalidRequestError(f"Invalid request for model '{model}': {exc}")

        # Rate limiting
        if isinstance(exc, LiteLLMRateLimitError):
            return ProxyError(
                f"Rate limit exceeded for model '{model}': {exc}",
                status_code=429,
                error_type="rate_limit_error",
            )

        # Timeout
        if isinstance(exc, LiteLLMTimeout):
            return BackendTimeoutError(f"Request to '{model}' timed out: {exc}")

        # Service unavailable
        if isinstance(exc, LiteLLMServiceUnavailableError):
            return BackendUnavailableError(f"Provider for '{model}' is unavailable: {exc}")

        # Fallback: wrap any other exception
        return ProxyError(
            f"LiteLLM error for model '{model}': {exc}",
            status_code=502,
            error_type="backend_error",
        )
