"""LLM 客户端:协议 + litellm 实现(DeepSeek / MiniMax)+ 测试用脚本桩。

面向接口编程——上层只依赖 LLM 协议,换模型/换厂商只改注入(litellm 已抽象多厂商)。
"""

from __future__ import annotations

import re
from typing import Callable, Protocol

from genesis import config

_THINK_BLOCK = re.compile(r"<think>.*?</think>\s*", re.S)   # 推理模型(MiniMax-M3)的思考块,解析前剥离


class LLM(Protocol):
    def complete(self, messages: list[dict], **kwargs: object) -> str:
        """传 [{'role','content'}...],返回模型文本输出。"""
        ...


class LiteLLMClient:
    """litellm 后端,按 config.LLM_PROVIDER 路由到 MiniMax(默认)或 DeepSeek。"""

    def __init__(
        self,
        model: str | None = None,
        *,
        api_key: str | None = None,
        api_base: str | None = None,
        temperature: float = 0.8,
    ) -> None:
        import litellm  # 延迟导入:离线/单测环境不强依赖

        self._litellm = litellm
        if config.LLM_PROVIDER == "minimax":
            # MiniMax 走 OpenAI 兼容端点("openai/<model>" + 自定义 api_base)
            self.model = model or f"openai/{config.MINIMAX_MODEL}"
            self.api_key = api_key or config.MINIMAX_API_KEY
            self.api_base = api_base or config.MINIMAX_BASE_URL
        else:
            # litellm 用 "deepseek/<model>" 路由到 DeepSeek 兼容接口
            self.model = model or f"deepseek/{config.DEEPSEEK_MODEL}"
            self.api_key = api_key or config.DEEPSEEK_API_KEY
            self.api_base = api_base or config.DEEPSEEK_BASE_URL
        self.temperature = temperature

    def complete(self, messages: list[dict], **kwargs: object) -> str:
        # 必须带超时:否则单次卡死会拖垮整轮并发结算(实测偶发挂起)
        # 思考模式:DeepSeek 与 MiniMax-M3 共用 {"thinking":{"type":...}} 参数,默认 enabled(慢且贵)——
        # 本项目默认关闭,仅在调用方显式传 _thinking=True 的"深思熟虑时刻"(如对质投票)开启。
        thinking = bool(kwargs.pop("_thinking", False))
        extra_body = dict(kwargs.pop("extra_body", {}) or {})
        # 开思考的值两家不同:DeepSeek=enabled,MiniMax=adaptive;关思考都用 disabled
        on_type = "adaptive" if config.LLM_PROVIDER == "minimax" else "enabled"
        extra_body.setdefault("thinking", {"type": on_type if thinking else "disabled"})
        resp = self._litellm.completion(
            model=self.model,
            messages=messages,
            api_key=self.api_key,
            api_base=self.api_base,
            temperature=kwargs.pop("temperature", self.temperature),
            timeout=kwargs.pop("timeout", 60),
            num_retries=kwargs.pop("num_retries", 1),
            extra_body=extra_body,
            **kwargs,
        )
        content = resp.choices[0].message.content or ""
        # 双保险:即便思考没关干净,也剥离内联 <think> 块,保证下游 JSON 解析不被污染
        return _THINK_BLOCK.sub("", content).strip()


class ScriptedLLM:
    """确定性测试桩:按队列依次返回,或用 responder(messages)->str 动态应答。"""

    def __init__(
        self,
        responses: list[str] | None = None,
        *,
        responder: Callable[[list[dict]], str] | None = None,
    ) -> None:
        self._responses = list(responses or [])
        self._responder = responder
        self.calls: list[list[dict]] = []  # 记录每次调用的 messages,便于断言

    def complete(self, messages: list[dict], **kwargs: object) -> str:
        self.calls.append(messages)
        if self._responder is not None:
            return self._responder(messages)
        if self._responses:
            return self._responses.pop(0)
        raise AssertionError("ScriptedLLM 响应已耗尽")