"""
LiteLLM universal model provider.

Supports 100+ LLM providers through a single unified API:
  - Cloud: Groq, Mistral, Together, Fireworks, Deepseek, Cohere, AI21, etc.
  - Enterprise: Azure OpenAI, AWS Bedrock, Google Vertex AI, Databricks
  - China: 智谱 (Zhipu/GLM), 通义千问 (Qwen/DashScope), 百川 (Baichuan), 月之暗面 (Moonshot)
  - Self-hosted: vLLM, TGI, LiteLLM Proxy
  - All via one interface, OpenAI-compatible tool calling included

Model naming convention (LiteLLM format):
  - "groq/llama-3.1-70b-versatile"
  - "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
  - "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"
  - "vertex_ai/gemini-1.5-pro"
  - "deepseek/deepseek-chat"
  - "zhipu/glm-4"
  - See full list: https://docs.litellm.ai/docs/providers
"""

from __future__ import annotations

import logging
from collections.abc import AsyncIterator
from typing import Any

from nexus.core.message import Message
from nexus.models.base import ModelProvider
from nexus.types import TokenUsage

logger = logging.getLogger("nexus.models.litellm")

# Known LiteLLM provider prefixes — used for auto-detection in api.py
LITELLM_PROVIDER_PREFIXES: set[str] = {
    # Major cloud providers
    "groq",
    "together_ai",
    "together",
    "fireworks_ai",
    "fireworks",
    "mistral",
    "deepseek",
    "cohere",
    "cohere_chat",
    "ai21",
    "perplexity",
    "anyscale",
    "replicate",
    # Enterprise / hyperscaler
    "azure",
    "azure_ai",
    "bedrock",
    "vertex_ai",
    "vertex_ai_beta",
    "sagemaker",
    "databricks",
    # China providers
    "zhipu",
    "dashscope",
    "qwen",
    "baichuan",
    "moonshot",
    "yi",
    "minimax",
    "spark",  # 讯飞星火
    "hunyuan",  # 腾讯混元
    "volcengine",  # 字节火山引擎（豆包）
    # Self-hosted / proxy
    "vllm",
    "text-completion-openai",
    "openrouter",
    "huggingface",
    "petals",
    # LiteLLM proxy
    "litellm_proxy",

    # openai
    "openai",
}


class LiteLLMProvider(ModelProvider):
    """
    Universal LLM provider powered by LiteLLM.

    This single provider unlocks 100+ LLM backends. It delegates to
    litellm.acompletion(), which handles auth, retries, and format
    normalization for each downstream provider.

    Usage:
        # Groq — ultra-fast inference
        provider = LiteLLMProvider(model="groq/llama-3.1-70b-versatile")

        # AWS Bedrock
        provider = LiteLLMProvider(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")

        # Deepseek
        provider = LiteLLMProvider(model="deepseek/deepseek-chat")

        # Google Vertex AI
        provider = LiteLLMProvider(model="vertex_ai/gemini-1.5-pro")

        # 通义千问
        provider = LiteLLMProvider(model="dashscope/qwen-max")

        # With API key override
        provider = LiteLLMProvider(
            model="groq/llama-3.1-70b-versatile",
            api_key="gsk_xxx",
        )

        # With LiteLLM proxy server
        provider = LiteLLMProvider(
            model="my-model",
            api_base="http://localhost:4000",
        )

        response = await provider.chat(messages=[{"role": "user", "content": "Hello"}])
    """

    def __init__(
            self,
            model: str = "groq/llama-3.1-70b-versatile",
            api_key: str | None = None,
            api_base: str | None = None,
            default_temperature: float = 0.7,
            default_max_tokens: int | None = None,
            # LiteLLM-specific options
            num_retries: int = 2,
            request_timeout: float = 120.0,
            fallbacks: list[str] | None = None,
            metadata: dict[str, Any] | None = None,
            **extra_kwargs: Any,
    ) -> None:
        self._model = model
        self._api_key = api_key
        self._api_base = api_base
        self._default_temperature = default_temperature
        self._default_max_tokens = default_max_tokens
        self._num_retries = num_retries
        self._request_timeout = request_timeout
        self._fallbacks = fallbacks
        self._metadata = metadata or {}
        self._extra_kwargs = extra_kwargs
        self._litellm: Any = None  # lazy import

    def _get_litellm(self) -> Any:
        """Lazy-import litellm to avoid import overhead."""
        if self._litellm is None:
            try:
                import litellm

                # Suppress litellm's verbose logging
                litellm.suppress_debug_info = True
                try:
                    litellm._logging._disable_debugging()
                except Exception:
                    pass
                import logging as _logging
                for _name in ("LiteLLM", "LiteLLM Router", "LiteLLM Proxy",
                              "litellm", "litellm.litellm", "litellm.utils",
                              "litellm.main", "litellm.cost_calculator"):
                    _logging.getLogger(_name).setLevel(_logging.WARNING)

                self._litellm = litellm
            except ImportError:
                raise ImportError(
                    "litellm package is required for LiteLLMProvider. "
                    "Install it with: pip install nexus-agents[litellm]"
                )
        return self._litellm

    async def chat(
            self,
            messages: list[dict[str, Any]],
            tools: list[dict[str, Any]] | None = None,
            temperature: float | None = None,
            max_tokens: int | None = None,
            **kwargs: Any,
    ) -> dict[str, Any]:
        litellm = self._get_litellm()

        # Build request params
        params: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "temperature": (
                temperature if temperature is not None else self._default_temperature
            ),
            "num_retries": self._num_retries,
            "timeout": self._request_timeout,
        }

        # Max tokens
        effective_max_tokens = max_tokens or self._default_max_tokens
        if effective_max_tokens:
            params["max_tokens"] = effective_max_tokens

        # API key / base
        if self._api_key:
            params["api_key"] = self._api_key
        if self._api_base:
            params["api_base"] = self._api_base

        # Tool calling — LiteLLM uses OpenAI format
        if tools:
            params["tools"] = tools
            params["tool_choice"] = kwargs.pop("tool_choice", "auto")

        # Fallback models — LiteLLM's killer feature
        if self._fallbacks:
            params["fallbacks"] = self._fallbacks

        # Metadata for cost tracking
        if self._metadata:
            params["metadata"] = self._metadata

        # Merge extra kwargs (provider-specific params)
        params.update(self._extra_kwargs)
        params.update(kwargs)

        # Call LiteLLM's unified async completion
        n_tools = len(tools) if tools else 0
        logger.debug(
            f"[LiteLLM] Request → model={self._model}, "
            f"messages={len(messages)}, tools={n_tools}"
        )

        response = await litellm.acompletion(**params)
        choice = response.choices[0]
        resp_message = choice.message

        # Extract tool calls (OpenAI-compatible format)
        tool_calls: list[dict[str, Any]] | None = None
        if hasattr(resp_message, "tool_calls") and resp_message.tool_calls:
            tool_calls = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in resp_message.tool_calls
            ]
            logger.debug(
                f"[LiteLLM] Response contains {len(tool_calls)} tool call(s): "
                f"{[tc['function']['name'] for tc in tool_calls]}"
            )

        # Build normalized Message
        content = resp_message.content or ""
        message = Message.agent(
            content=content,
            sender_id=self._model,
        )
        if tool_calls:
            message.tool_calls = tool_calls

        # Token usage — LiteLLM normalizes this across all providers
        usage_data = response.usage if hasattr(response, "usage") and response.usage else None
        usage = TokenUsage(
            prompt_tokens=getattr(usage_data, "prompt_tokens", 0) or 0,
            completion_tokens=getattr(usage_data, "completion_tokens", 0) or 0,
            total_tokens=getattr(usage_data, "total_tokens", 0) or 0,
        )

        # LiteLLM provides cost tracking — attach it
        try:
            cost = litellm.completion_cost(completion_response=response)
            usage.cost_usd = cost
        except Exception:
            pass  # Cost data not available for all providers

        logger.debug(
            f"[LiteLLM] Usage: {usage.prompt_tokens}+{usage.completion_tokens}={usage.total_tokens} tokens"
            + (f", cost=${usage.cost_usd:.6f}" if usage.cost_usd else "")
        )

        return {
            "message": message,
            "usage": usage,
            "model": getattr(response, "model", self._model),
            "finish_reason": getattr(choice, "finish_reason", "stop"),
            "tool_calls": tool_calls,
            "raw": response,
        }

    async def close(self) -> None:
        """LiteLLM manages its own connection pool — nothing to close."""
        self._litellm = None

    async def chat_stream(
            self,
            messages: list[dict[str, Any]],
            tools: list[dict[str, Any]] | None = None,
            temperature: float | None = None,
            max_tokens: int | None = None,
            **kwargs: Any,
    ) -> AsyncIterator[str]:
        """
        Stream a chat completion, yielding text chunks as they arrive.

        Uses litellm.acompletion(stream=True) for true token-by-token
        streaming across all 100+ supported providers.
        """
        litellm = self._get_litellm()

        params: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "temperature": (
                temperature if temperature is not None else self._default_temperature
            ),
            "stream": True,
            "num_retries": self._num_retries,
            "timeout": self._request_timeout,
        }

        effective_max_tokens = max_tokens or self._default_max_tokens
        if effective_max_tokens:
            params["max_tokens"] = effective_max_tokens

        if self._api_key:
            params["api_key"] = self._api_key
        if self._api_base:
            params["api_base"] = self._api_base

        # No tool calling in streaming mode for simplicity
        if self._fallbacks:
            params["fallbacks"] = self._fallbacks
        if self._metadata:
            params["metadata"] = self._metadata

        params.update(self._extra_kwargs)
        params.update(kwargs)

        response = await litellm.acompletion(**params)

        async for chunk in response:
            delta = chunk.choices[0].delta if chunk.choices else None
            if delta and getattr(delta, "content", None):
                yield delta.content

    @property
    def model_name(self) -> str:
        return self._model
