"""U15 — LiteLLM 统一 Provider 适配层。

用 LiteLLM 的 ``acompletion()`` 统一接口替换 6 个直接 API provider 适配器
（OpenAI/Anthropic/Gemini/Doubao/Wenxin/Yuanbao）。LiteLLM 内部处理各家 API
的差异（消息格式、tool_calls 格式、streaming SSE 协议），本模块只负责：

1. 把 ``LLMRequest`` 翻译成 LiteLLM ``acompletion`` 的 kwargs；
2. 把 LiteLLM 响应（OpenAI ChatCompletion 格式）翻译回 ``LLMResponse`` /
   ``StreamChunk``；
3. 把 LiteLLM 异常包装成 ``LLMProviderError``，保留 fallback 链可用性。

设计取舍（ponytail）：
- 自建的 fallback 链 / usage tracking / 部门配额 全部保留在 ``LLMGateway``，
  本 provider 不重复实现。LiteLLM 自带的 fallback / retry 在此处禁用
  （``num_retries=0``），避免与 gateway 的 fallback 重复叠加导致超时放大。
- ``wenxin`` 在 LiteLLM 中无原生支持，回退到 ``openai/`` 前缀 + 自定义
  ``api_base``（千帆 OpenAI 兼容端点）。升级路径：LiteLLM 上游支持后切换到
  ``wenxin/`` 前缀。
- 旧的 6 个 provider 类（``OpenAICompatibleProvider`` 等）保留为死代码一个
  release，便于回滚；新代码通过 ``create_litellm_provider`` 工厂构造。
"""

from __future__ import annotations

import inspect
import json
import logging
import time
from collections.abc import AsyncGenerator, Iterable

from agentkit.core.exceptions import LLMProviderError
from agentkit.llm.protocol import (
    LLMProvider,
    LLMRequest,
    LLMResponse,
    StreamChunk,
    TokenUsage,
    ToolCall,
)

logger = logging.getLogger(__name__)


# provider_type → LiteLLM model 前缀映射。
# LiteLLM 通过 model 字符串前缀路由到对应 SDK：``openai/gpt-4o``、
# ``anthropic/claude-...``、``gemini/gemini-...``、``volcengine/...``、
# ``hunyuan/...``。未知类型回退到 ``openai/``（最大兼容性）。
_PROVIDER_TYPE_TO_PREFIX: dict[str, str] = {
    "openai": "openai/",
    "anthropic": "anthropic/",
    "gemini": "gemini/",
    # 豆包由火山引擎提供，LiteLLM 前缀为 volcengine/
    "doubao": "volcengine/",
    # ponytail: LiteLLM 暂无 wenxin 原生支持，回退到 openai/ + 自定义 api_base
    # （千帆 OpenAI 兼容端点）。ceiling: wenxin 专属参数（如 AK/SK 鉴权）不支持；
    # 升级路径：上游支持后改 "wenxin/"。
    "wenxin": "openai/",
    # 腾讯混元
    "yuanbao": "hunyuan/",
}


def _model_prefix_for(provider_type: str) -> str:
    """根据 provider_type 返回 LiteLLM model 前缀；未知类型回退到 openai/。"""
    return _PROVIDER_TYPE_TO_PREFIX.get(provider_type, "openai/")


class LitellmProvider(LLMProvider):
    """基于 LiteLLM ``acompletion`` 的统一 Provider。

    一个实例对应一个 provider 配置（api_key + base_url + provider_type）。
    ``chat`` / ``chat_stream`` 把 ``LLMRequest`` 转发给 LiteLLM 并翻译响应。

    注意：本类不持有 httpx 客户端 — LiteLLM 内部管理连接池。
    """

    def __init__(
        self,
        model_prefix: str,
        api_key: str,
        base_url: str | None = None,
        provider_type: str = "openai",
        **default_kwargs: object,
    ) -> None:
        self._model_prefix = model_prefix
        self._api_key = api_key
        self._base_url = base_url or None  # 空字符串视作未设置
        self._provider_type = provider_type
        self._default_kwargs: dict[str, object] = dict(default_kwargs)

    async def chat(self, request: LLMRequest) -> LLMResponse:
        """非流式 chat — 调用 ``litellm.acompletion`` 并翻译响应。"""
        import litellm

        kwargs = self._build_kwargs(request, stream=False)

        start = time.monotonic()
        try:
            response = await litellm.acompletion(**kwargs)
        except Exception as e:
            # LiteLLM 抛出各种异常（openai.APIError / anthropic.APIError / 自定义），
            # 统一包装成 LLMProviderError 以便 gateway fallback 链识别。
            raise LLMProviderError(self._provider_type, str(e)) from e
        latency_ms = (time.monotonic() - start) * 1000

        return self._parse_response(response, request.model, latency_ms)

    async def chat_stream(self, request: LLMRequest) -> AsyncGenerator[StreamChunk, None]:
        """流式 chat — 调用 ``litellm.acompletion(stream=True)`` 并翻译 chunks。

        异步生成器安全：本项目规则禁止在第一个 ``yield`` 前使用 ``return``，
        本函数无早退分支，无需 ``return; yield`` 守卫。
        """
        import litellm

        kwargs = self._build_kwargs(request, stream=True)

        accumulated_tool_calls: dict[int, dict[str, object]] = {}
        final_usage: TokenUsage | None = None
        final_model: str = request.model

        try:
            # litellm.acompletion(stream=True) 的返回类型取决于版本 / 调用方式：
            # - 真实 litellm：返回 coroutine，await 后得到 async generator；
            # - 测试 mock（async def + yield）：调用即返回 async generator，不可 await。
            # 用 isawaitable 兼容两种路径。
            raw = litellm.acompletion(**kwargs)
            stream = await raw if inspect.isawaitable(raw) else raw
            async for chunk in stream:
                parsed = self._parse_stream_chunk(
                    chunk,
                    request.model,
                    accumulated_tool_calls,
                )
                # 更新累计状态
                if parsed.usage is not None:
                    final_usage = parsed.usage
                if parsed.model:
                    final_model = parsed.model
                # 内容块（含空内容）yield；usage-only 块也 yield（部分 provider
                # 在最后一个 chunk 才给 usage，内容为空）。
                yield parsed

            # 流结束：yield 终止 chunk（聚合 tool_calls + usage）
            tool_calls_list = self._finalize_tool_calls(accumulated_tool_calls)
            yield StreamChunk(
                content="",
                model=final_model,
                tool_calls=tool_calls_list,
                usage=final_usage,
                is_final=True,
            )
        except Exception as e:
            raise LLMProviderError(self._provider_type, str(e)) from e

    # ------------------------------------------------------------------
    # 内部辅助
    # ------------------------------------------------------------------

    def _build_kwargs(self, request: LLMRequest, *, stream: bool) -> dict[str, object]:
        """从 LLMRequest 构造 litellm.acompletion kwargs。"""
        kwargs: dict[str, object] = {
            "model": f"{self._model_prefix}{request.model}",
            "messages": request.messages,
            "temperature": request.temperature,
            "max_tokens": request.max_tokens,
            "api_key": self._api_key,
            "stream": stream,
            # 禁用 LiteLLM 自带 retry — 由 gateway 的 fallback 链 / RetryPolicy 负责
            "num_retries": 0,
        }
        if self._base_url:
            kwargs["api_base"] = self._base_url
        if request.tools:
            kwargs["tools"] = request.tools
            kwargs["tool_choice"] = request.tool_choice
        if request.timeout is not None:
            kwargs["timeout"] = request.timeout
        # U17 — 透传 LiteLLM cache 参数（cache_key 或 no-cache）到 litellm.acompletion
        cache_params = request._cache
        if cache_params is not None:
            kwargs["cache"] = cache_params
        # 合并构造时传入的默认 kwargs（如 max_connections 等provider特定参数）
        kwargs.update(self._default_kwargs)
        return kwargs

    def _parse_response(
        self,
        response: object,
        request_model: str,
        latency_ms: float,
    ) -> LLMResponse:
        """把 litellm 响应（OpenAI ChatCompletion 格式）翻译成 LLMResponse。"""
        # LiteLLM 响应统一为 OpenAI ChatCompletion 格式：
        # response.choices[0].message.{content, tool_calls}
        # response.usage.{prompt_tokens, completion_tokens}
        # response.model
        choices = getattr(response, "choices", None) or []
        content = ""
        tool_calls: list[ToolCall] = []
        if choices:
            message = getattr(choices[0], "message", None)
            if message is not None:
                content = getattr(message, "content", None) or ""
                raw_tool_calls = getattr(message, "tool_calls", None)
                if raw_tool_calls:
                    tool_calls = _parse_tool_calls(raw_tool_calls)

        usage_obj = getattr(response, "usage", None)
        usage = _parse_usage(usage_obj) if usage_obj is not None else TokenUsage()

        model_name = getattr(response, "model", None) or request_model

        # U17 — 检测 LiteLLM 缓存命中（_hidden_params 含 cache_key 或 cache_hit）
        cache_hit = False
        hidden = getattr(response, "_hidden_params", None)
        if isinstance(hidden, dict) and ("cache_key" in hidden or hidden.get("cache_hit")):
            cache_hit = True

        return LLMResponse(
            content=content,
            model=model_name,
            usage=usage,
            tool_calls=tool_calls,
            latency_ms=latency_ms,
            cache_hit=cache_hit,
        )

    def _parse_stream_chunk(
        self,
        chunk: object,
        request_model: str,
        accumulated_tool_calls: dict[int, dict[str, object]],
    ) -> StreamChunk:
        """解析单个流式 chunk（非 final）。累计 tool_calls 到传入字典。"""
        choices = getattr(chunk, "choices", None) or []
        content = ""
        model_name = getattr(chunk, "model", None) or request_model
        usage: TokenUsage | None = None

        if choices:
            delta = getattr(choices[0], "delta", None)
            if delta is not None:
                content = getattr(delta, "content", None) or ""
                raw_tool_calls = getattr(delta, "tool_calls", None)
                if raw_tool_calls:
                    _accumulate_stream_tool_calls(raw_tool_calls, accumulated_tool_calls)

        # 部分 provider 在最后一个 chunk 附带 usage
        usage_obj = getattr(chunk, "usage", None)
        if usage_obj is not None:
            usage = _parse_usage(usage_obj)

        return StreamChunk(
            content=content,
            model=model_name,
            tool_calls=[],  # 流式中间块不带 tool_calls，由 final chunk 聚合
            usage=usage,
            is_final=False,
        )

    def _finalize_tool_calls(
        self,
        accumulated: dict[int, dict[str, object]],
    ) -> list[ToolCall]:
        """把累计的流式 tool_calls 字典转成 ToolCall 列表。"""
        tool_calls: list[ToolCall] = []
        for idx in sorted(accumulated.keys()):
            tc_data = accumulated[idx]
            args_str = tc_data.get("arguments_str", "")
            try:
                arguments = json.loads(args_str) if args_str else {}
            except json.JSONDecodeError:
                arguments = {"raw": args_str}
            tool_calls.append(
                ToolCall(
                    id=tc_data.get("id", ""),
                    name=tc_data.get("name", ""),
                    arguments=arguments,
                )
            )
        return tool_calls


# ----------------------------------------------------------------------
# 响应解析辅助（模块级，便于测试 mock）
# ----------------------------------------------------------------------


def _parse_tool_calls(raw_tool_calls: Iterable[object]) -> list[ToolCall]:
    """解析非流式响应的 tool_calls（OpenAI 格式 list[ChoiceMessageToolCall]）。"""
    result: list[ToolCall] = []
    for tc in raw_tool_calls:
        # LiteLLM 返回的对象有 .id / .function.{name,arguments}
        tc_id = getattr(tc, "id", "") or ""
        func = getattr(tc, "function", None)
        if func is None:
            continue
        name = getattr(func, "name", "") or ""
        args = getattr(func, "arguments", "{}")
        if isinstance(args, str):
            try:
                args_dict = json.loads(args) if args else {}
            except json.JSONDecodeError:
                args_dict = {"raw": args}
        elif isinstance(args, dict):
            args_dict = args
        else:
            args_dict = {"raw": str(args)}
        result.append(ToolCall(id=tc_id, name=name, arguments=args_dict))
    return result


def _parse_usage(usage_obj: object) -> TokenUsage:
    """解析 usage 对象（OpenAI CompletionUsage 或 dict）。"""
    prompt = getattr(usage_obj, "prompt_tokens", None)
    completion = getattr(usage_obj, "completion_tokens", None)
    if prompt is None and isinstance(usage_obj, dict):
        prompt = usage_obj.get("prompt_tokens", 0)
        completion = usage_obj.get("completion_tokens", 0)
    return TokenUsage(
        prompt_tokens=int(prompt or 0),
        completion_tokens=int(completion or 0),
    )


def _accumulate_stream_tool_calls(
    raw_tool_calls: Iterable[object],
    accumulated: dict[int, dict[str, object]],
) -> None:
    """累计流式 chunk 里的 tool_calls 片段（OpenAI delta.tool_calls 格式）。

    每个 delta tool_call 有 index/id/function.{name,arguments}，arguments
    是分片字符串，需跨 chunk 拼接。
    """
    for tc in raw_tool_calls:
        idx = getattr(tc, "index", 0) or 0
        if idx not in accumulated:
            accumulated[idx] = {
                "id": "",
                "name": "",
                "arguments_str": "",
            }
        tc_id = getattr(tc, "id", None)
        if tc_id:
            accumulated[idx]["id"] = tc_id
        func = getattr(tc, "function", None)
        if func is not None:
            name = getattr(func, "name", None)
            if name:
                accumulated[idx]["name"] = name
            args_fragment = getattr(func, "arguments", None)
            if args_fragment:
                accumulated[idx]["arguments_str"] += args_fragment


# ----------------------------------------------------------------------
# 工厂函数
# ----------------------------------------------------------------------


def create_litellm_provider(
    provider_type: str,
    api_key: str,
    base_url: str | None = None,
    **kwargs: object,
) -> LitellmProvider:
    """根据 provider_type 创建 LitellmProvider 实例。

    Args:
        provider_type: 配置中的 provider 类型字符串
            （"openai"|"anthropic"|"gemini"|"doubao"|"wenxin"|"yuanbao"）。
            未知类型回退到 "openai/" 前缀。
        api_key: provider API key。
        base_url: 可选自定义 API base URL。
        **kwargs: 透传给 LitellmProvider 的额外默认参数。

    Returns:
        配置好 model_prefix 的 LitellmProvider。
    """
    prefix = _model_prefix_for(provider_type)
    return LitellmProvider(
        model_prefix=prefix,
        api_key=api_key,
        base_url=base_url,
        provider_type=provider_type,
        **kwargs,
    )
