from Plugins.WPSAPI import *
from datetime import datetime
from llama_index.llms.ollama import Ollama
from llama_index.core.llms import ChatMessage
from typing import List, Dict, Optional

logger: ProjectConfig = Architecture.Get(ProjectConfig)
OLLAMA_URL = logger.FindItem("ollama_url", "http://ollama.liubai.site")
OLLAMA_MODEL = logger.FindItem("ollama_model", "qwen3:0.6b")
MAX_HISTORY = logger.FindItem("chat_ai_max_history", 20)
logger.SaveProperties()


class ChatAIAgent:
    """AI 对话智能体 - 维护全局会话历史"""
    
    _instance: Optional['ChatAIAgent'] = None
    
    def __init__(self, ollama_url: str = OLLAMA_URL, max_history: int = MAX_HISTORY):
        """初始化 AI 智能体
        
        Args:
            ollama_url: Ollama 服务地址
            max_history: 最大历史消息数量
        """
        self.ollama_url = ollama_url
        self.max_history = max_history
        self.llm = Ollama(model=OLLAMA_MODEL, base_url=ollama_url, request_timeout=600.0)
        
        # 全局消息历史列表
        self.global_history: List[Dict[str, any]] = []
        
        # 系统提示词
        self.system_prompt = """你是一个友好的AI助手，能够同时在多个群聊中与不同用户对话。

重要规则：
- 你能看到所有群聊的对话历史，每条消息都标注了时间、群聊ID和用户ID
- 当你需要特指某个用户或专门回复某个用户时，必须在回复中使用以下格式：<at user_id="用户ID"></at>
- 禁止直接输出用户ID数字，始终使用 <at> 标签包裹
- 注意区分不同群聊和不同用户的对话上下文

示例对话：
[2025-11-20 10:00:00] [群聊#12345] [用户#11111]: 你好
你的回复: 你好！有什么可以帮助你的吗？

[2025-11-20 10:01:00] [群聊#67890] [用户#22222]: 刚刚谁在说话, 说了什么？
你的回复: 刚刚群聊12345的用户<at user_id="11111"></at>在说话, 他说了你好。"""
        
        logger.Log("Info", f"{ConsoleFrontColor.GREEN}ChatAIAgent 初始化完成{ConsoleFrontColor.RESET}")
        logger.Log("Info", f"模型: {OLLAMA_MODEL}, 最大历史: {max_history}")
    
    @classmethod
    def get_instance(cls) -> 'ChatAIAgent':
        """获取单例实例"""
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance
    
    def _add_message(self, chat_id: int, user_id: int, role: str, content: str) -> None:
        """添加消息到全局历史
        
        Args:
            chat_id: 群聊ID
            user_id: 用户ID
            role: 角色 (user/assistant)
            content: 消息内容
        """
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        message = {
            "timestamp": timestamp,
            "chat_id": chat_id,
            "user_id": user_id,
            "role": role,
            "content": content
        }
        self.global_history.append(message)
        self._trim_history()
    
    def _trim_history(self) -> None:
        """裁剪历史消息到最大长度"""
        if len(self.global_history) > self.max_history:
            # 保留最新的 max_history 条消息
            self.global_history = self.global_history[-self.max_history:]
            logger.Log("Info", f"历史消息已裁剪到 {self.max_history} 条")
    
    def clear_history(self) -> bool:
        """清空全局历史
        
        Returns:
            是否成功
        """
        try:
            old_count = len(self.global_history)
            self.global_history.clear()
            logger.Log("Info", f"已清空全局历史，共删除 {old_count} 条消息")
            return True
        except Exception as e:
            logger.Log("Error", f"清空历史失败: {e}")
            return False
    
    def _format_message_with_context(self, msg: Dict[str, any]) -> str:
        """格式化单条消息，包含上下文信息
        
        Args:
            msg: 消息字典
            
        Returns:
            格式化后的消息字符串
        """
        timestamp = msg.get("timestamp", "")
        chat_id = msg.get("chat_id", "")
        user_id = msg.get("user_id", "")
        role = msg.get("role", "")
        content = msg.get("content", "")
        
        if role == "user":
            return f"[{timestamp}] [群聊#{chat_id}] [用户#{user_id}]: {content}"
        else:  # assistant
            return f"[{timestamp}] [AI助手]: {content}"
    
    def _format_history_for_llm(self) -> str:
        """格式化全局历史为 LLM 输入
        
        Returns:
            格式化后的历史字符串
        """
        if not self.global_history:
            return ""
        
        formatted_messages = [
            self._format_message_with_context(msg)
            for msg in self.global_history
        ]
        return "\n".join(formatted_messages)
    
    async def chat(self, message: str, chat_id: int, user_id: int) -> str:
        """处理对话
        
        Args:
            message: 用户消息
            chat_id: 群聊ID
            user_id: 用户ID
            
        Returns:
            AI 回复
        """
        try:
            # 添加用户消息到历史
            self._add_message(chat_id, user_id, "user", message)
            
            # 格式化历史消息（不包括刚添加的用户消息）
            history_without_last = self.global_history[:-1] if len(self.global_history) > 1 else []
            history_text = "\n".join([
                self._format_message_with_context(msg)
                for msg in history_without_last
            ]) if history_without_last else "（这是第一条对话）"
            
            # 构建当前用户消息的格式化版本
            current_msg = f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] [群聊#{chat_id}] [用户#{user_id}]: {message}"
            
            # 构建对话提示
            conversation_prompt = f"""以下是历史对话记录：
{history_text}

当前消息：
{current_msg}

请作为AI助手回复当前用户。记住：
1. 使用 <at user_id="{user_id}"></at> 格式来提到任意你想提及的用户
2. 根据对话历史和当前消息，给出有意义的回复
3. 如果历史中有其他群聊或用户的信息，你可以引用它们"""
            
            logger.Log("Info", f"处理对话 - 群聊#{chat_id} 用户#{user_id}")
            logger.Log("Info", f"当前历史消息数: {len(self.global_history)}")
            
            # 使用 achat 方法进行对话
            messages = [
                ChatMessage(role="system", content=self.system_prompt),
                ChatMessage(role="user", content=conversation_prompt)
            ]
            
            response = await self.llm.achat(messages)
            answer = str(response.message.content)
            
            # 添加 AI 回复到历史
            self._add_message(chat_id, user_id, "assistant", answer)
            
            logger.Log("Info", f"AI 回复长度: {len(answer)} 字符")
            
            return answer
            
        except Exception as e:
            logger.Log("Error", f"对话处理失败: {e}")
            import traceback
            error_trace = traceback.format_exc()
            logger.Log("Error", f"详细错误:\n{error_trace}")
            return f"处理对话时出错: {str(e)}"


class ChatAIPlugin(WPSAPI):
    """AI 对话插件"""
    
    def __init__(self):
        super().__init__()
        self.ai_agent = ChatAIAgent.get_instance()
    
    @override
    def dependencies(self) -> List[Type]:
        return [WPSAPI]
    
    @override
    def is_enable_plugin(self) -> bool:
        return True
    
    def get_guide_title(self) -> str:
        return "AI 智能对话"
    
    def get_guide_subtitle(self) -> str:
        return "基于 LlamaIndex + Ollama 的多群聊 AI 对话系统"
    
    def get_guide_metadata(self) -> Dict[str, str]:
        return {
            "AI模型": OLLAMA_MODEL,
            "最大历史": str(MAX_HISTORY),
            "功能": "跨群聊智能对话",
        }
    
    def collect_command_entries(self) -> Sequence[GuideEntry]:
        return (
            {
                "title": "AI对话",
                "identifier": "ai_chat",
                "description": "与AI助手对话，AI能看到所有群聊的历史消息并智能回复。",
                "metadata": {"模型": OLLAMA_MODEL},
                "icon": "🤖",
                "badge": "AI",
                "details": [
                    {
                        "type": "list",
                        "items": [
                            "AI 维护全局会话历史，能看到所有群聊的对话",
                            "每条消息包含时间戳、群聊ID、用户ID",
                            f"最多保留最近 {MAX_HISTORY} 条消息",
                            "AI 会使用 @用户 格式回复",
                            "示例：ai_chat 你好",
                            "示例：ai_chat 刚才其他群聊说了什么？",
                        ]
                    }
                ]
            },
        )
    
    def collect_guide_entries(self) -> Sequence[GuideEntry]:
        return (
            {
                "title": "全局会话",
                "description": (
                    "AI 维护一个全局的对话历史，能够跨群聊、跨用户理解上下文。"
                    "所有群聊的对话都在同一个历史中，AI 能够关联不同群聊的信息。"
                ),
                "icon": "🌐",
            },
            {
                "title": "智能回复",
                "description": (
                    "AI 能够识别消息来源（群聊、用户、时间），并使用 @ 格式回复特定用户。"
                ),
                "icon": "💬",
            },
            {
                "title": "历史管理",
                "description": (
                    f"自动保留最近 {MAX_HISTORY} 条消息，超出部分自动清理。"
                    "可使用 ai_chat_clear 指令手动清空所有历史。"
                ),
                "icon": "📝",
            },
        )
    
    @override
    def wake_up(self) -> None:
        logger.Log("Info", f"{ConsoleFrontColor.GREEN}ChatAIPlugin AI对话插件已加载{ConsoleFrontColor.RESET}")
        self.register_plugin("ai_chat")
        self.register_plugin("default")
    
    @override
    async def callback(self, message: str, chat_id: int, user_id: int) -> str|None:
        """处理用户对话"""
        try:
            if not message or message.strip() == "":
                help_text = f"""# 🤖 AI 智能对话使用帮助

**直接发送消息即可与 AI 对话**

**特性：**
- ✨ AI 能看到所有群聊的对话历史
- 🌐 跨群聊智能理解上下文
- 💬 自动使用 @ 格式回复用户
- 📝 保留最近 {MAX_HISTORY} 条消息

**示例：**
- `ai_chat 你好`
- `ai_chat 今天天气怎么样？`
- `ai_chat 刚才其他群聊说了什么？`
- `ai_chat 请总结一下最近的对话`

**清空历史：**
使用 `ai_chat_clear` 指令清空所有历史消息

**技术：**
基于 LlamaIndex + Ollama ({OLLAMA_MODEL})"""
                return await self.send_markdown_message(help_text, chat_id, user_id)
            
            # 使用智能体处理对话
            answer = await self.ai_agent.chat(message, chat_id, user_id)
            
            # 格式化返回结果
            formatted_answer = answer
            
            return await self.send_markdown_message(formatted_answer, chat_id, user_id)
            
        except Exception as e:
            logger.Log("Error", f"AI对话异常: {e}")
            import traceback
            error_detail = traceback.format_exc()
            logger.Log("Error", f"详细错误: {error_detail}")
            error_msg = f"""❌ **处理对话时出错**

错误信息：{str(e)}

请稍后重试或联系管理员。"""
            return await self.send_markdown_message(error_msg, chat_id, user_id)


class ChatAIClearPlugin(WPSAPI):
    """清空 AI 对话历史插件"""
    
    def __init__(self):
        super().__init__()
    
    @override
    def dependencies(self) -> List[Type]:
        return [ChatAIPlugin]
    
    @override
    def is_enable_plugin(self) -> bool:
        return True
    
    def get_guide_title(self) -> str:
        return "清空 AI 对话历史"
    
    def get_guide_subtitle(self) -> str:
        return "清空全局 AI 对话历史记录"
    
    def get_guide_metadata(self) -> Dict[str, str]:
        return {
            "功能": "历史管理",
            "作用范围": "全局",
        }
    
    def collect_command_entries(self) -> Sequence[GuideEntry]:
        return (
            {
                "title": "清空历史",
                "identifier": "ai_chat_clear",
                "description": "清空 AI 的全局对话历史记录。",
                "icon": "🗑️",
                "badge": "管理",
                "details": [
                    {
                        "type": "list",
                        "items": [
                            "清空所有群聊的对话历史",
                            "AI 将不再记得之前的对话",
                            "不影响 AI 的基础功能",
                            "示例：ai_chat_clear",
                        ]
                    }
                ]
            },
        )
    
    @override
    def wake_up(self) -> None:
        logger.Log("Info", f"{ConsoleFrontColor.GREEN}ChatAIClearPlugin 清空历史插件已加载{ConsoleFrontColor.RESET}")
        self.register_plugin("ai_chat_clear")
    
    @override
    async def callback(self, message: str, chat_id: int, user_id: int) -> str|None:
        """清空对话历史"""
        try:
            # 获取 AI Agent 实例
            ai_agent = ChatAIAgent.get_instance()
            
            # 清空历史
            success = ai_agent.clear_history()
            
            if success:
                result_msg = """✅ **历史已清空**

AI 的全局对话历史已成功清空。

AI 将不再记得之前的对话内容。"""
            else:
                result_msg = """❌ **清空失败**

清空历史时出现错误，请稍后重试。"""
            
            return await self.send_markdown_message(result_msg, chat_id, user_id)
            
        except Exception as e:
            logger.Log("Error", f"清空历史异常: {e}")
            error_msg = f"""❌ **清空历史时出错**

错误信息：{str(e)}

请稍后重试或联系管理员。"""
            return await self.send_markdown_message(error_msg, chat_id, user_id)