from typing import Optional, Dict
from dataclasses import dataclass
import os
from config import settings
from openai import OpenAI
from utils import logger


@dataclass
class LLMConfig:
    """LLM配置类"""
    api_key: Optional[str] = None
    base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    model: str = settings.LLM_MODEL
    enable_thinking: bool = True
    temperature: float = 0.7
    max_tokens: int = settings.LLM_MAX_TOKENS


class LLMThinkingEngine:
    """LLM驱动的思考引擎实现"""
    
    def __init__(self, system_prompt_file: str = "system_prompt.txt", config: Optional[LLMConfig] = None):
        """
        初始化LLMThinkingEngine
        
        Args:
            config: LLM配置对象，如果为None则使用默认配置
        """
        self.system_prompt_file = system_prompt_file
        self.config = config or LLMConfig()
        self._init_client()
    
    def _init_client(self):
        """初始化OpenAI客户端"""
        api_key = self.config.api_key or settings.LLM_API_KEY
        self.client = OpenAI(
            api_key=api_key,
            base_url=self.config.base_url,
        )
    
    def think(self, user_input: str) -> str:
        """
        基于LLM进行思考，返回下一步的行动
        
        Args:
            user_input: 用户输入内容
            
        Returns:
            Thought: 思考结果，包含行动类型和内容
        """
        # 构建适用于LLM的消息
        messages = self._build_messages(user_input)
        logger.info(f"LLM构建的消息: {messages}")
        # 调用LLM进行思考
        thinking_content, response_content = self._call_llm(messages)
        
        # logger.info(f"LLM思考结果: thinking_content={thinking_content}, response_content={response_content}")
        return response_content
    
    def _build_messages(self, user_input: str) -> list[Dict[str, str]]:
        """
        构建发送给LLM的消息

        Args:
            user_input: 用户输入内容

        Returns:
            消息列表，包含系统提示、历史和当前输入
        """
        messages = []
        
        # 系统提示
        system_prompt = self._get_system_prompt()
        messages.append({"role": "system", "content": system_prompt})
    
        # 用户输入
        messages.append({
            "role": "user",
            "content": user_input
        })
        
        return messages
    
    def _get_system_prompt(self) -> str:
        """
        获取系统提示词
        
        Returns:
            系统提示词
        """
        prompt_path = os.path.join(
            os.path.dirname(__file__), 
            "prompts", 
            self.system_prompt_file
        )
        with open(prompt_path, "r", encoding="utf-8") as f:
            return f.read()
    
    def _call_llm(self, messages: list[Dict[str, str]]) -> tuple[str, str]:
        """
        调用LLM API
        
        Args:
            messages: 消息列表
            
        Returns:
            (thinking_content, response_content): 思考内容和响应内容
        """
        thinking_content = ""
        response_content = ""
        
        try:
            completion = self.client.chat.completions.create(
                model=self.config.model,
                messages=messages,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                extra_body={"enable_thinking": self.config.enable_thinking},
                stream=True
            )
            
            # 流式处理响应
            for chunk in completion:
                delta = chunk.choices[0].delta
                
                # 收集思考内容
                if hasattr(delta, "reasoning_content") and delta.reasoning_content:
                    thinking_content += delta.reasoning_content
                
                # 收集响应内容
                if hasattr(delta, "content") and delta.content:
                    response_content += delta.content
            
        except Exception as e:
            # 错误处理
            response_content = f"调用LLM时出错：{str(e)}"
        
        return thinking_content, response_content
    
    def set_config(self, config: LLMConfig):
        """更新LLM配置"""
        self.config = config
        self._init_client()