"""
LLM 客户端工具

提供统一的 LLM 接口,支持多种模型提供商。
使用 LangChain 的 ChatModel 抽象,简化不同厂商 API 的调用。
"""
from typing import Optional, List, Dict, Any, AsyncGenerator
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage
from langchain_core.callbacks import AsyncCallbackHandler
from core.config import settings
import time


def get_llm(
    provider: str = "openai",
    model: Optional[str] = None,
    temperature: float = 0.7,
    streaming: bool = False,
    **kwargs
) -> BaseChatModel:
    """
    获取 LLM 实例
    
    Args:
        provider: 模型提供商 ("openai", "anthropic", "ollama")
        model: 模型名称 (如果不指定则使用配置中的默认值)
        temperature: 温度参数 (0-2)
        streaming: 是否启用流式输出
        **kwargs: 其他参数传递给模型
        
    Returns:
        BaseChatModel: LangChain 的聊天模型实例
    """
    
    if provider == "openai":
        from langchain_openai import ChatOpenAI
        
        return ChatOpenAI(
            model=model or settings.OPENAI_MODEL or "gpt-4",
            temperature=temperature,
            api_key=settings.OPENAI_API_KEY,
            streaming=streaming,
            **kwargs
        )
    
    elif provider == "anthropic":
        from langchain_anthropic import ChatAnthropic
        
        return ChatAnthropic(
            model=model or settings.ANTHROPIC_MODEL or "claude-3-opus-20240229",
            temperature=temperature,
            api_key=settings.ANTHROPIC_API_KEY,
            max_tokens=kwargs.pop("max_tokens", 4096),
            streaming=streaming,
            **kwargs
        )
    
    elif provider == "ollama":
        try:
            from langchain_ollama import ChatOllama
            
            return ChatOllama(
                model=model or settings.OLLAMA_MODEL or "llama3",
                base_url=settings.OLLAMA_BASE_URL or "http://localhost:11434",
                temperature=temperature,
                **kwargs
            )
        except ImportError:
            raise ImportError(
                "langchain-ollama not installed. Run: pip install langchain-ollama"
            )
    
    else:
        raise ValueError(f"Unsupported provider: {provider}. Use 'openai', 'anthropic', or 'ollama'")


# 便捷函数 - 常用配置
def get_fast_llm(provider: str = "openai") -> BaseChatModel:
    """获取快速响应的 LLM (低温度,适合事实性问题)"""
    return get_llm(provider, temperature=0.3)


def get_creative_llm(provider: str = "openai") -> BaseChatModel:
    """获取创造性 LLM (高温度,适合创意写作)"""
    return get_llm(provider, temperature=0.9)


def get_streaming_llm(provider: str = "openai") -> BaseChatModel:
    """获取支持流式输出的 LLM"""
    return get_llm(provider, streaming=True)


class TokenUsageCallbackHandler(AsyncCallbackHandler):
    """
    Token 使用回调处理器
    
    用于捕获 LLM 调用的 token 使用情况
    """
    def __init__(self):
        self.prompt_tokens = 0
        self.completion_tokens = 0
        self.total_tokens = 0
        self.response_content = ""
    
    async def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str] = None, **kwargs):
        """LLM 开始时的回调"""
        pass
    
    async def on_llm_end(self, response, **kwargs):
        """LLM 结束时的回调，获取 token 统计"""
        try:
            # 从 response 中提取 token 信息
            if hasattr(response, 'llm_output') and response.llm_output:
                token_usage = response.llm_output.get('token_usage', {})
                self.prompt_tokens = token_usage.get('prompt_tokens', 0)
                self.completion_tokens = token_usage.get('completion_tokens', 0)
                self.total_tokens = token_usage.get('total_tokens', 0)
        except Exception as e:
            print(f"[TokenUsageCallback] 提取 token 信息失败: {e}")
    
    async def on_llm_new_token(self, token: str, **kwargs):
        """每个新 token 的回调（流式输出）"""
        self.response_content += token


class LLMClient:
    """
    LLM 客户端封装类
    
    提供统一的异步接口，支持自定义API配置、流式输出和 token 统计
    """
    
    async def chat_completion(
        self,
        messages: List[BaseMessage],
        api_url: str,
        api_key: str,
        model: str = "gpt-3.5-turbo",
        temperature: float = 1.0,
        max_tokens: int = 500,
        request_timeout: int = 60,
        stream: bool = False,
        **kwargs
    ) -> Dict[str, Any]:
        """
        调用 LLM API 生成回复
        
        Args:
            messages: LangChain 消息列表
            api_url: API 地址
            api_key: API 密钥
            model: 模型名称
            temperature: 温度参数
            max_tokens: 最大 token 数
            request_timeout: 请求超时时间（秒）
            stream: 是否启用流式输出
            **kwargs: 其他参数
            
        Returns:
            OpenAI 格式的响应字典，包含 token 使用信息
        """
        try:
            from langchain_openai import ChatOpenAI
            
            # 创建回调处理器
            callback_handler = TokenUsageCallbackHandler()
            
            # 创建自定义的 ChatOpenAI 实例
            llm = ChatOpenAI(
                model=model,
                temperature=temperature,
                api_key=api_key,
                base_url=api_url if api_url else None,
                max_tokens=max_tokens,
                streaming=stream,
                callbacks=[callback_handler],
                request_timeout=request_timeout,  # ✅ 设置超时时间
                **kwargs
            )
            
            start_time = time.time()
            
            if stream:
                # 流式模式
                full_content = ""
                async for chunk in llm.astream(messages):
                    if hasattr(chunk, 'content'):
                        full_content += chunk.content
                
                duration = time.time() - start_time
                
                return {
                    "choices": [
                        {
                            "message": {
                                "role": "assistant",
                                "content": full_content
                            }
                        }
                    ],
                    "usage": {
                        "prompt_tokens": callback_handler.prompt_tokens,
                        "completion_tokens": callback_handler.completion_tokens,
                        "total_tokens": callback_handler.total_tokens
                    },
                    "duration": duration
                }
            else:
                # 非流式模式
                response = await llm.ainvoke(messages)
                duration = time.time() - start_time
                
                return {
                    "choices": [
                        {
                            "message": {
                                "role": "assistant",
                                "content": response.content
                            }
                        }
                    ],
                    "usage": {
                        "prompt_tokens": callback_handler.prompt_tokens,
                        "completion_tokens": callback_handler.completion_tokens,
                        "total_tokens": callback_handler.total_tokens
                    },
                    "duration": duration
                }
            
        except Exception as e:
            print(f"[LLMClient] 调用失败: {e}")
            raise
    
    async def stream_chat(
        self,
        messages: List[BaseMessage],
        api_url: str,
        api_key: str,
        model: str = "gpt-3.5-turbo",
        temperature: float = 1.0,
        max_tokens: int = 500,
        request_timeout: int = 60,
        **kwargs
    ) -> AsyncGenerator[Dict[str, Any], None]:
        """
        流式调用 LLM API
        
        Args:
            messages: LangChain 消息列表
            api_url: API 地址
            api_key: API 密钥
            model: 模型名称
            temperature: 温度参数
            max_tokens: 最大 token 数
            request_timeout: 请求超时时间（秒）
            **kwargs: 其他参数
            
        Yields:
            包含 token 片段的字典
        """
        try:
            from langchain_openai import ChatOpenAI
            
            print(f"\n[LLMClient] 🔧 创建 ChatOpenAI 实例")
            print(f"  - Model: {model}")
            print(f"  - API URL: {api_url[:50]}..." if len(api_url) > 50 else f"  - API URL: {api_url}")
            print(f"  - Temperature: {temperature}")
            print(f"  - Max Tokens: {max_tokens}")
            print(f"  - Request Timeout: {request_timeout}s")
            
            # 创建回调处理器
            callback_handler = TokenUsageCallbackHandler()
            
            # 创建自定义的 ChatOpenAI 实例
            llm = ChatOpenAI(
                model=model,
                temperature=temperature,
                api_key=api_key,
                base_url=api_url if api_url else None,
                max_tokens=max_tokens,
                streaming=True,
                callbacks=[callback_handler],
                request_timeout=request_timeout,  # ✅ 设置超时时间
                **kwargs
            )
            
            start_time = time.time()
            
            print(f"[LLMClient] 🚀 开始流式请求...")
            print(f"  - Messages 数量: {len(messages)}")
            if messages:
                first_msg_role = getattr(messages[0], 'role', 'unknown')
                first_msg_preview = str(getattr(messages[0], 'content', ''))[:50]
                print(f"  - 第一条消息: [{first_msg_role}] {first_msg_preview}...")
            
            chunk_count = 0
            # 流式输出
            async for chunk in llm.astream(messages):
                if hasattr(chunk, 'content') and chunk.content:
                    chunk_count += 1
                    
                    # 第一个 chunk 时记录
                    if chunk_count == 1:
                        first_chunk_time = time.time()
                        print(f"[LLMClient] ✨ 收到第一个 chunk (耗时: {first_chunk_time - start_time:.2f}s)")
                    
                    yield {
                        "type": "chunk",
                        "content": chunk.content
                    }
            
            duration = time.time() - start_time
            
            print(f"[LLMClient] ✅ 流式请求完成")
            print(f"  - 总 Chunks: {chunk_count}")
            print(f"  - 耗时: {duration:.2f}秒")
            print(f"  - Prompt Tokens: {callback_handler.prompt_tokens}")
            print(f"  - Completion Tokens: {callback_handler.completion_tokens}")
            print(f"  - Total Tokens: {callback_handler.total_tokens}\n")
            
            # 最后发送 token 使用信息
            yield {
                "type": "usage",
                "usage": {
                    "prompt_tokens": callback_handler.prompt_tokens,
                    "completion_tokens": callback_handler.completion_tokens,
                    "total_tokens": callback_handler.total_tokens
                },
                "duration": duration
            }
            
        except Exception as e:
            print(f"\n[LLMClient] ❌ 流式调用失败: {e}")
            import traceback
            traceback.print_exc()
            raise
