"""
模型适配器 - 统一模型调用接口

提供与 KlynxAgent 兼容的 .invoke(messages) 接口。
使用 LiteLLM 作为统一标准接口，完美支持全球各大模型厂商 (OpenAI, Anthropic, Google, DeepSeek 等)。
通过显式传递 httpx.Client 避开局域网代理限制。
"""

import logging
import os
import random
import time
from typing import Any, List, Optional

import httpx
import litellm
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage

# 禁用从远端请求价格表导致的 SSL 请求超时警告
litellm.suppress_debug_info = True
logger = logging.getLogger(__name__)


class _InferenceBackend:
    """模型推理后端抽象。"""

    def completion(self, call_kwargs: dict):
        raise NotImplementedError

    def stream_completion(self, call_kwargs: dict):
        raise NotImplementedError


class _LiteLLMBackend(_InferenceBackend):
    """默认后端：LiteLLM completion。"""

    def completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)

    def stream_completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)


class _ResponsesCompatBackend(_InferenceBackend):
    """
    Responses API 兼容占位后端。

    当前阶段保持与现有调用链兼容，内部仍代理到 LiteLLM completion。
    后续可在此类中替换为真正的 Responses API 调用实现。
    """

    def completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)

    def stream_completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)


class LiteLLMResponse:
    """
    响应包装类
    兼容 graph.py 中对 response.content / response.reasoning_content / response.usage 的访问
    同时支持原生 Tool Calling（兼容旧称 Function Calling）返回的 tool_calls
    """
    def __init__(self, content: str, reasoning_content: str = ""):
        self.content = content
        self.reasoning_content = reasoning_content
        self.additional_kwargs = {}
        self.response_metadata = {}
        self.usage = None  # 由调用方填充
        self.tool_calls = []  # 原生 Tool Calling 返回的工具调用列表


class LiteLLMChat:
    """
    模型适配器
    
    使用 LiteLLM 标准化接口调用各模型厂商，
    对外暴露与原来相同的 .invoke(messages) 和 .stream(messages) 接口。
    支持可选的原生 Tool Calling（兼容旧称 Function Calling，tools 参数）。
    """
    
    def __init__(self, model: str, api_key: str = None, api_base: str = None,
                 temperature: float = 0.1, backend: str = "litellm", **kwargs):
        """
        Args:
            model: 模型标识符 (如 "deepseek/deepseek-reasoner", "gpt-4o")
                   支持 litellm provider/model 格式
            api_key: API Key
            api_base: 自定义 API 端点 (可选)
            temperature: 生成温度
            **kwargs: 额外参数
        """
        self.model = model
        self.api_key = api_key
        self.api_base = api_base
        self.temperature = temperature
        self.max_context_tokens = 128000
        self.backend_name = str(backend or "litellm").strip().lower()
        self.backend = self._create_backend(self.backend_name)

        # 网络不稳定场景下的容错参数（默认开启重试）
        self.retry_enabled = self._parse_bool_option(
            kwargs.pop("retry_enabled", os.getenv("KLYNX_MODEL_RETRY_ENABLED", "1")),
            default=True,
        )
        self.retry_max_attempts = self._parse_int_option(
            kwargs.pop("retry_max_attempts", os.getenv("KLYNX_MODEL_RETRY_MAX_ATTEMPTS", "3")),
            default=3,
            minimum=1,
        )
        self.retry_initial_delay = self._parse_float_option(
            kwargs.pop("retry_initial_delay", os.getenv("KLYNX_MODEL_RETRY_INITIAL_DELAY", "1.0")),
            default=1.0,
            minimum=0.0,
        )
        self.retry_backoff_factor = self._parse_float_option(
            kwargs.pop("retry_backoff_factor", os.getenv("KLYNX_MODEL_RETRY_BACKOFF_FACTOR", "2.0")),
            default=2.0,
            minimum=1.0,
        )
        self.retry_max_delay = self._parse_float_option(
            kwargs.pop("retry_max_delay", os.getenv("KLYNX_MODEL_RETRY_MAX_DELAY", "8.0")),
            default=8.0,
            minimum=0.0,
        )
        self.retry_jitter = self._parse_float_option(
            kwargs.pop("retry_jitter", os.getenv("KLYNX_MODEL_RETRY_JITTER", "0.2")),
            default=0.2,
            minimum=0.0,
        )
        fallback_models_raw = kwargs.pop("fallback_models", os.getenv("KLYNX_FALLBACK_MODELS", ""))
        self.fallback_models = self._parse_fallback_models(fallback_models_raw)

        self.extra_kwargs = kwargs

        # 创建客户端（禁用代理，避免 SSL 错误）
        self.http_client = httpx.Client(trust_env=False)
    
    def invoke(self, messages, tools: Optional[list] = None):
        """
        调用模型，返回 LiteLLMResponse
        
        Args:
            messages: LangChain BaseMessage 列表 或 OpenAI 格式字典列表
            tools: 可选，OpenAI 格式的 tools JSON Schema 列表（用于原生 Tool Calling）
            
        Returns:
            LiteLLMResponse 对象
        """
        # 转换消息格式
        openai_messages = self._sanitize_jsonable(self._convert_messages(messages))
        
        # 构建调用参数
        call_kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "api_key": self.api_key,
        }
        
        if self.api_base:
            call_kwargs["api_base"] = self.api_base
            
        # 部分思考模型对 temperature 有严格限制
        if not any(k in self.model for k in ["reasoner", "o1", "kimi-k2.5"]):
            call_kwargs["temperature"] = self.temperature
        
        # 原生 Tool Calling: 注入 tools
        if tools:
            call_kwargs["tools"] = self._sanitize_jsonable(tools)
            call_kwargs["tool_choice"] = "auto"
        
        # 合并额外参数
        call_kwargs.update(self.extra_kwargs)
        call_kwargs = self._sanitize_jsonable(call_kwargs)
        
        # 调用 API（带重试与可选模型降级）
        response = self._completion_with_retry(call_kwargs)
        
        # 解析响应
        message = response.choices[0].message
        content = message.content or ""
        # 兼容各大厂商对 reasoning_content 的非标准输出名称
        reasoning_content = getattr(message, 'reasoning_content', '') or ""
        
        # 构建响应对象
        resp = LiteLLMResponse(content=content, reasoning_content=reasoning_content)
        
        # 解析原生 tool_calls（如果有）
        if hasattr(message, 'tool_calls') and message.tool_calls:
            resp.tool_calls = self._parse_native_tool_calls(message.tool_calls)
        
        # 填充 usage
        usage = response.usage
        if usage:
            resp.usage = {
                "prompt_tokens": getattr(usage, "prompt_tokens", 0),
                "completion_tokens": getattr(usage, "completion_tokens", 0),
                "total_tokens": getattr(usage, "total_tokens", 0),
            }
        
        return resp
    
    def stream(self, messages, tools: Optional[list] = None):
        """
        流式调用模型
        
        Args:
            messages: LangChain 消息列表
            tools: 可选，OpenAI 格式的 tools JSON Schema 列表（用于原生 Tool Calling）
            
        Yields:
            chunks with 'content', 'reasoning_content', 'usage', and 'tool_calls'
        """
        openai_messages = self._sanitize_jsonable(self._convert_messages(messages))
        
        call_kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "stream": True,
            "stream_options": {"include_usage": True},
            "api_key": self.api_key,
        }
        
        if self.api_base:
            call_kwargs["api_base"] = self.api_base
            
        if not any(k in self.model for k in ["reasoner", "o1", "kimi-k2.5"]):
            call_kwargs["temperature"] = self.temperature
        
        # 原生 Tool Calling: 注入 tools
        if tools:
            call_kwargs["tools"] = self._sanitize_jsonable(tools)
            call_kwargs["tool_choice"] = "auto"
            
        call_kwargs.update(self.extra_kwargs)
        call_kwargs = self._sanitize_jsonable(call_kwargs)
            
        response = self._stream_with_retry(call_kwargs)

        # 用于拼接流式 tool_calls 片段
        tool_call_accumulators = {}

        for chunk in response:
            # 处理 Usage 信息 (通常在最后一个 chunk)
            if hasattr(chunk, 'usage') and chunk.usage:
                yield {
                    "usage": {
                        "prompt_tokens": getattr(chunk.usage, "prompt_tokens", 0),
                        "completion_tokens": getattr(chunk.usage, "completion_tokens", 0),
                        "total_tokens": getattr(chunk.usage, "total_tokens", 0)
                    }
                }

            if not chunk.choices:
                continue

            delta = chunk.choices[0].delta
            content = delta.content or ""
            reasoning_content = getattr(delta, 'reasoning_content', '') or ""

            # 拼接流式 tool_calls 片段
            if hasattr(delta, 'tool_calls') and delta.tool_calls:
                for tc_delta in delta.tool_calls:
                    idx = tc_delta.index
                    if idx not in tool_call_accumulators:
                        tool_call_accumulators[idx] = {
                            "id": "",
                            "function": {"name": "", "arguments": ""}
                        }
                    acc = tool_call_accumulators[idx]
                    if tc_delta.id:
                        acc["id"] = tc_delta.id
                    if tc_delta.function:
                        if tc_delta.function.name:
                            acc["function"]["name"] += tc_delta.function.name
                        if tc_delta.function.arguments:
                            acc["function"]["arguments"] += tc_delta.function.arguments

            if content or reasoning_content or (hasattr(delta, 'tool_calls') and delta.tool_calls):
                yield {
                    "content": content,
                    "reasoning_content": reasoning_content
                }

        # 流结束后，如果有拼接好的 tool_calls，一次性 yield 出去
        if tool_call_accumulators:
            assembled = []
            for idx in sorted(tool_call_accumulators.keys()):
                acc = tool_call_accumulators[idx]
                assembled.append(self._parse_single_tool_call(acc))
            yield {"tool_calls": assembled}
    
    def _parse_native_tool_calls(self, raw_tool_calls) -> list:
        """
        将 LiteLLM 返回的原生 tool_calls 对象列表转换为统一的字典格式
        
        Returns:
            [{"tool": "read_file", "params": {"path": "..."}}, ...]
        """
        import json
        results = []
        for tc in raw_tool_calls:
            func = tc.function
            name = func.name
            try:
                args = json.loads(func.arguments) if isinstance(func.arguments, str) else func.arguments
            except json.JSONDecodeError:
                args = {"raw_arguments": func.arguments}
            results.append({
                "tool": name,
                "params": args
            })
        return results

    def _create_backend(self, backend: str) -> _InferenceBackend:
        if backend == "responses":
            return _ResponsesCompatBackend()
        return _LiteLLMBackend()
    
    def _parse_single_tool_call(self, acc: dict) -> dict:
        """
        解析单个拼接完成的 tool_call 累积器为统一格式
        """
        import json
        name = acc["function"]["name"]
        raw_args = acc["function"]["arguments"]
        try:
            args = json.loads(raw_args) if raw_args else {}
        except json.JSONDecodeError:
            args = {"raw_arguments": raw_args}
        return {
            "tool": name,
            "params": args
        }
    
    def _convert_messages(self, messages):
        """将 LangChain 消息列表转换为 OpenAI 格式"""
        openai_messages = []
        for msg in messages:
            if isinstance(msg, dict):
                openai_messages.append(msg)
            elif isinstance(msg, SystemMessage):
                openai_messages.append({"role": "system", "content": msg.content})
            elif isinstance(msg, HumanMessage):
                openai_messages.append({"role": "user", "content": msg.content})
            elif isinstance(msg, AIMessage):
                openai_messages.append({"role": "assistant", "content": msg.content})
            else:
                role = getattr(msg, 'type', 'user')
                openai_messages.append({"role": role, "content": str(msg.content)})
        return openai_messages

    def _sanitize_jsonable(self, value: Any) -> Any:
        """清洗 JSON 负载中的非法 Unicode 代理字符，避免请求序列化失败。"""
        if isinstance(value, str):
            # 例如 read_tui 输出中偶发的 lone surrogate（\ud800-\udfff）
            # 会在编码/JSON 序列化阶段触发 UTF-8 错误，这里统一替换为 U+FFFD。
            return value.encode("utf-8", errors="replace").decode("utf-8")
        if isinstance(value, list):
            return [self._sanitize_jsonable(item) for item in value]
        if isinstance(value, tuple):
            return tuple(self._sanitize_jsonable(item) for item in value)
        if isinstance(value, dict):
            return {
                self._sanitize_jsonable(key): self._sanitize_jsonable(item)
                for key, item in value.items()
            }
        return value

    def _completion_with_retry(self, call_kwargs: dict):
        model_candidates = self._get_model_candidates()
        for model_idx, model_name in enumerate(model_candidates):
            model_kwargs = dict(call_kwargs)
            model_kwargs["model"] = model_name
            for attempt in range(1, self.retry_max_attempts + 1):
                try:
                    return self.backend.completion(model_kwargs)
                except Exception as exc:
                    if (not self.retry_enabled) or (not self._is_retriable_error(exc)):
                        raise

                    has_next_attempt = attempt < self.retry_max_attempts
                    has_next_model = model_idx < len(model_candidates) - 1
                    if not has_next_attempt and not has_next_model:
                        raise

                    if has_next_attempt:
                        delay = self._retry_delay_seconds(attempt)
                        logger.warning(
                            "Model completion failed, retrying in %.2fs (model=%s, attempt=%s/%s): %s",
                            delay,
                            model_name,
                            attempt,
                            self.retry_max_attempts,
                            exc,
                        )
                        time.sleep(delay)
                    else:
                        logger.warning(
                            "Model completion failed after %s attempts, switching to fallback model: %s -> %s; error=%s",
                            self.retry_max_attempts,
                            model_name,
                            model_candidates[model_idx + 1],
                            exc,
                        )
                    continue

        raise RuntimeError("Model completion failed unexpectedly without captured exception")

    def _stream_with_retry(self, call_kwargs: dict):
        model_candidates = self._get_model_candidates()
        for model_idx, model_name in enumerate(model_candidates):
            model_kwargs = dict(call_kwargs)
            model_kwargs["model"] = model_name
            for attempt in range(1, self.retry_max_attempts + 1):
                emitted_any = False
                try:
                    response = self.backend.stream_completion(model_kwargs)
                    for chunk in response:
                        emitted_any = True
                        yield chunk
                    return
                except Exception as exc:
                    # 流中途已经输出内容时，不做重试避免重复输出
                    if emitted_any:
                        raise
                    if (not self.retry_enabled) or (not self._is_retriable_error(exc)):
                        raise

                    has_next_attempt = attempt < self.retry_max_attempts
                    has_next_model = model_idx < len(model_candidates) - 1
                    if not has_next_attempt and not has_next_model:
                        raise

                    if has_next_attempt:
                        delay = self._retry_delay_seconds(attempt)
                        logger.warning(
                            "Model stream failed, retrying in %.2fs (model=%s, attempt=%s/%s): %s",
                            delay,
                            model_name,
                            attempt,
                            self.retry_max_attempts,
                            exc,
                        )
                        time.sleep(delay)
                    else:
                        logger.warning(
                            "Model stream failed after %s attempts, switching to fallback model: %s -> %s; error=%s",
                            self.retry_max_attempts,
                            model_name,
                            model_candidates[model_idx + 1],
                            exc,
                        )
                    continue

        raise RuntimeError("Model stream failed unexpectedly without captured exception")

    def _get_model_candidates(self) -> list:
        names = [str(self.model).strip()]
        for item in self.fallback_models:
            candidate = str(item or "").strip()
            if candidate and candidate not in names:
                names.append(candidate)
        return names

    def _retry_delay_seconds(self, attempt: int) -> float:
        # attempt 从 1 开始，对应首次重试
        base_delay = self.retry_initial_delay * (self.retry_backoff_factor ** max(0, attempt - 1))
        jitter = random.uniform(0.0, self.retry_jitter) if self.retry_jitter > 0 else 0.0
        return min(self.retry_max_delay, base_delay + jitter)

    def _is_retriable_error(self, exc: Exception) -> bool:
        if isinstance(exc, (httpx.TimeoutException, httpx.TransportError, TimeoutError, ConnectionError)):
            return True

        text = f"{type(exc).__name__}: {exc}".lower()
        non_retriable_markers = (
            "invalid api key",
            "incorrect api key",
            "authentication",
            "unauthorized",
            "forbidden",
            "permission",
            "invalid_request",
            "badrequest",
            "context_length_exceeded",
            "max context length",
            "unsupported",
            "not found",
            "does not exist",
        )
        if any(marker in text for marker in non_retriable_markers):
            return False

        retriable_markers = (
            "connection error",
            "connection aborted",
            "connection reset",
            "remote protocol error",
            "temporarily unavailable",
            "service unavailable",
            "timeout",
            "timed out",
            "eof occurred in violation of protocol",
            "ssl",
            "429",
            "502",
            "503",
            "504",
            "rate limit",
            "try again",
        )
        return any(marker in text for marker in retriable_markers)

    @staticmethod
    def _parse_bool_option(value: Any, default: bool) -> bool:
        if value is None:
            return default
        if isinstance(value, bool):
            return value
        text = str(value).strip().lower()
        if text in {"1", "true", "yes", "y", "on"}:
            return True
        if text in {"0", "false", "no", "n", "off"}:
            return False
        return default

    @staticmethod
    def _parse_int_option(value: Any, default: int, minimum: int) -> int:
        try:
            parsed = int(value)
        except Exception:
            return default
        return max(minimum, parsed)

    @staticmethod
    def _parse_float_option(value: Any, default: float, minimum: float) -> float:
        try:
            parsed = float(value)
        except Exception:
            return default
        return max(minimum, parsed)

    @staticmethod
    def _parse_fallback_models(value: Any) -> list:
        if value is None:
            return []
        if isinstance(value, (list, tuple)):
            return [str(item).strip() for item in value if str(item or "").strip()]
        text = str(value).strip()
        if not text:
            return []
        return [segment.strip() for segment in text.split(",") if segment.strip()]

    def __repr__(self):
        return f"LiteLLMChat(model={self.model!r}, max_context_tokens={self.max_context_tokens})"


# 为了兼容历史代码依赖，提供原名的别名，建议未来直接使用 LiteLLMChat
DeepSeekReasonerChat = LiteLLMChat
DeepSeekReasonerResponse = LiteLLMResponse
