"""LiteLLM integration — unified API for multiple LLM providers."""

from __future__ import annotations

import logging
from collections.abc import AsyncIterator, Iterator
from dataclasses import dataclass
from typing import Any
from cyberbrain26.model_layer.provider_interface import ChatResult, ProviderABC, ProviderProtocol

logger = logging.getLogger(__name__)


@dataclass
class LiteLLMConfig:
    """Configuration for LiteLLM client."""

    api_key: str | None = None
    base_url: str | None = None
    max_retries: int = 3
    timeout: float = 60.0
    default_model: str | None = None
    extra_headers: dict[str, str] | None = None


class LiteLLMClient(ProviderABC):
    """LiteLLM-based model client providing unified access to multiple providers."""

    def __init__(self, config: LiteLLMConfig | None = None) -> None:
        self.config = config or LiteLLMConfig()
        self._client: Any = None
        self._initialized: bool = False

    def _ensure_initialized(self) -> None:
        if self._initialized:
            return
        try:
            import litellm

            self._client = litellm
            self._initialized = True
            logger.info("LiteLLM client initialized")
        except ImportError as exc:
            raise ImportError(
                "litellm is required for LiteLLMClient. "
                "Install with: pip install litellm"
            ) from exc

    @property
    def available_providers(self) -> list[str]:
        self._ensure_initialized()
        try:
            import litellm

            return list(litellm.LitellmConfig().get("providers", []))[:50]
        except Exception:  # noqa: BLE001
            return [
                "openai",
                "azure",
                "anthropic",
                "google",
                "gemini",
                "mistral",
                "groq",
                "ollama",
                "cohere",
                "huggingface",
            ]

    def chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        **kwargs: Any,
    ) -> ChatResult:
        self._ensure_initialized()
        import litellm

        temperature = kwargs.get("temperature")
        max_tokens = kwargs.get("max_tokens")
        tools = kwargs.get("tools")
        tool_choice = kwargs.get("tool_choice")
        stream = kwargs.get("stream", False)
        extra = kwargs.get("extra")

        params: dict[str, Any] = {
            "model": model,
            "messages": messages,
        }
        if temperature is not None:
            params["temperature"] = temperature
        if max_tokens is not None:
            params["max_tokens"] = max_tokens
        if tools:
            params["tools"] = tools
        if tool_choice is not None:
            params["tool_choice"] = tool_choice
        if self.config.base_url:
            params["api_base"] = self.config.base_url
        if self.config.timeout:
            params["timeout"] = self.config.timeout
        if extra:
            params.update(extra)

        logger.info("LiteLLM chat: model=%s, messages=%d", model, len(messages))

        if stream:
            return self._stream_chat(litellm, params)

        response = litellm.completion(**params)
        return response.choices[0].message.content or ""

    def _stream_chat(
        self, litellm_module: Any, params: dict[str, Any]
    ) -> Iterator[dict[str, Any]]:
        for chunk in litellm_module.streaming.completion(**params):
            yield {
                "id": getattr(chunk, "id", ""),
                "model": getattr(chunk, "model", params["model"]),
                "choices": [
                    {
                        "delta": {
                            "content": getattr(
                                getattr(c, "delta", None), "content", ""
                            ),
                            "role": getattr(
                                getattr(c, "delta", None), "role", "assistant"
                            ),
                        },
                        "finish_reason": getattr(c, "finish_reason", None),
                    }
                    for c in getattr(chunk, "choices", [])
                ],
            }

    def embed(
        self,
        model: str,
        texts: list[str],
        **kwargs: Any,
    ) -> list[list[float]]:
        self._ensure_initialized()
        import litellm

        extra = kwargs.get("extra")

        params: dict[str, Any] = {
            "model": model,
            "input": texts,
        }
        if self.config.base_url:
            params["api_base"] = self.config.base_url
        if extra:
            params.update(extra)

        logger.info("LiteLLM embed: model=%s, texts=%d", model, len(texts))

        response = litellm.embedding(**params)
        return [
            list(embedding)
            for embedding in getattr(response, "data", [])
            if hasattr(embedding, "embedding")
        ]

    def list_models(self) -> list[dict[str, Any]]:
        self._ensure_initialized()
        try:
            import litellm

            return [{"model": name} for name in list(litellm.model_list)[:200]]
        except Exception as exc:  # noqa: BLE001
            logger.warning("Failed to list models from LiteLLM: %s", exc)
        return []

    async def achat(
        self,
        model: str,
        messages: list[dict[str, str]],
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | None = None,
        stream: bool = False,
        extra: dict[str, Any] | None = None,
    ) -> str | AsyncIterator[dict[str, Any]]:
        if stream:
            return self._astream_chat(
                model, messages, temperature, max_tokens,
                tools, tool_choice, extra,
            )
        return await self._achat_nonstream(
            model, messages, temperature, max_tokens,
            tools, tool_choice, extra,
        )

    async def _astream_chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        temperature: float | None,
        max_tokens: int | None,
        tools: list[dict[str, Any]] | None,
        tool_choice: str | None,
        extra: dict[str, Any] | None,
    ) -> AsyncIterator[dict[str, Any]]:
        self._ensure_initialized()
        import litellm

        params: dict[str, Any] = {
            "model": model,
            "messages": messages,
        }
        if temperature is not None:
            params["temperature"] = temperature
        if max_tokens is not None:
            params["max_tokens"] = max_tokens
        if tools:
            params["tools"] = tools
        if tool_choice is not None:
            params["tool_choice"] = tool_choice
        if self.config.base_url:
            params["api_base"] = self.config.base_url
        if self.config.timeout:
            params["timeout"] = self.config.timeout
        if extra:
            params.update(extra)

        logger.info("LiteLLM async chat: model=%s", model)

        async for chunk in litellm.acompletion(**params):
            yield {
                "id": getattr(chunk, "id", ""),
                "model": getattr(chunk, "model", params["model"]),
                "choices": [
                    {
                        "delta": {
                            "content": getattr(
                                getattr(c, "delta", None), "content", ""
                            ),
                            "role": getattr(
                                getattr(c, "delta", None), "role", "assistant"
                            ),
                        },
                        "finish_reason": getattr(c, "finish_reason", None),
                    }
                    for c in getattr(chunk, "choices", [])
                ],
            }

    async def _achat_nonstream(
        self,
        model: str,
        messages: list[dict[str, str]],
        temperature: float | None,
        max_tokens: int | None,
        tools: list[dict[str, Any]] | None,
        tool_choice: str | None,
        extra: dict[str, Any] | None,
    ) -> str:
        self._ensure_initialized()
        import litellm

        params: dict[str, Any] = {
            "model": model,
            "messages": messages,
        }
        if temperature is not None:
            params["temperature"] = temperature
        if max_tokens is not None:
            params["max_tokens"] = max_tokens
        if tools:
            params["tools"] = tools
        if tool_choice is not None:
            params["tool_choice"] = tool_choice
        if self.config.base_url:
            params["api_base"] = self.config.base_url
        if self.config.timeout:
            params["timeout"] = self.config.timeout
        if extra:
            params.update(extra)

        logger.info("LiteLLM async chat: model=%s", model)

        response = await litellm.acompletion(**params)
        return response.choices[0].message.content or ""

    async def aembed(
        self,
        model: str,
        texts: list[str],
        *,
        extra: dict[str, Any] | None = None,
    ) -> list[list[float]]:
        self._ensure_initialized()
        import litellm

        params: dict[str, Any] = {
            "model": model,
            "input": texts,
        }
        if self.config.base_url:
            params["api_base"] = self.config.base_url
        if extra:
            params.update(extra)

        logger.info("LiteLLM async embed: model=%s", model)

        response = await litellm.aembedding(**params)
        return [
            list(embedding)
            for embedding in getattr(response, "data", [])
            if hasattr(embedding, "embedding")
        ]

    def is_available(self) -> bool:
        try:
            self._ensure_initialized()
            return True
        except Exception:  # noqa: BLE001
            return False

    def health_check(self) -> dict[str, Any]:
        try:
            self._ensure_initialized()
            return {
                "healthy": True,
                "provider": "litellm",
                "models_available": len(self.list_models()),
            }
        except Exception as exc:  # noqa: BLE001
            return {"healthy": False, "provider": "litellm", "error": str(exc)}