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

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

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


class LiteLLMResponse:
    """
    响应包装类
    兼容 graph.py 中对 response.content / response.reasoning_content / response.usage 的访问
    """
    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  # 由调用方填充


class LiteLLMChat:
    """
    模型适配器
    
    使用 LiteLLM 标准化接口调用各模型厂商，
    对外暴露与原来相同的 .invoke(messages) 和 .stream(messages) 接口。
    """
    
    def __init__(self, model: str, api_key: str = None, api_base: str = None,
                 temperature: float = 0.1, **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.extra_kwargs = kwargs
        self.max_context_tokens = 128000
        
        # 创建客户端（禁用代理，避免 SSL 错误）
        self.http_client = httpx.Client(trust_env=False)
    
    def invoke(self, messages):
        """
        调用模型，返回 LiteLLMResponse
        
        Args:
            messages: LangChain BaseMessage 列表 或 OpenAI 格式字典列表
            
        Returns:
            LiteLLMResponse 对象
        """
        # 转换消息格式
        openai_messages = self._convert_messages(messages)
        
        # 构建调用参数
        call_kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "api_key": self.api_key,
            "client": self.http_client,  # 绕过代理
        }
        
        if self.api_base:
            call_kwargs["api_base"] = self.api_base
            
        # reasoner 模型/o1模型不支持 temperature，litellm 会自动 drop，但显式加不加都行
        if "reasoner" not in self.model and "o1" not in self.model:
            call_kwargs["temperature"] = self.temperature
        
        # 合并额外参数
        call_kwargs.update(self.extra_kwargs)
        
        # 调用 API
        response = litellm.completion(**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)
        
        # 填充 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 _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 stream(self, messages):
        """
        流式调用模型
        
        Args:
            messages: LangChain 消息列表
            
        Yields:
            chunks with 'content', 'reasoning_content', and 'usage'
        """
        openai_messages = self._convert_messages(messages)
        
        call_kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "stream": True,
            "stream_options": {"include_usage": True},
            "api_key": self.api_key,
            "client": self.http_client,
        }
        
        if self.api_base:
            call_kwargs["api_base"] = self.api_base
            
        if "reasoner" not in self.model and "o1" not in self.model:
            call_kwargs["temperature"] = self.temperature
            
        call_kwargs.update(self.extra_kwargs)
            
        try:
            response = litellm.completion(**call_kwargs)
            
            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 ""
                
                if content or reasoning_content:
                    yield {
                        "content": content,
                        "reasoning_content": reasoning_content
                    }
        except Exception as e:
            import traceback
            traceback.print_exc()
            raise

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


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