#!/usr/bin/env python3
"""
WeChat Integration for bbav20311
基于 mcp-wechat-server 的微信集成模块

功能：
- 任务完成通知
- 接收用户指令
- 状态查询

作者: xialab的幸运星
日期: 2026-04-10
"""

import json
import os
import subprocess
import time
import asyncio
from pathlib import Path
from typing import Optional, Dict, Any, Callable
from dataclasses import dataclass
from datetime import datetime


@dataclass
class WeChatConfig:
    """微信配置"""
    mcp_server_command: str = "bunx"
    mcp_server_args: str = "mcp-wechat-server"
    qrcode_path: Path = Path.home() / ".mcp-wechat-server" / "qrcode.txt"
    qrcode_png_path: Path = Path.home() / ".mcp-wechat-server" / "qrcode.png"
    state_file: Path = Path.home() / ".mcp-wechat-server" / "state.json"
    
    # 任务通知模板
    task_complete_template: str = """🎉 任务完成通知

📋 任务ID: {task_id}
🔬 分析类型: {task_type}
⏱️ 耗时: {duration}
📊 状态: {status}

📁 输出文件:
{output_files}

💡 提示: 回复"详情 {task_id}"查看完整结果"""
    
    task_failed_template: str = """❌ 任务失败通知

📋 任务ID: {task_id}
🔬 分析类型: {task_type}
⏱️ 耗时: {duration}
❗ 错误: {error}

请检查输入数据或联系管理员"""


class WeChatIntegration:
    """
    微信集成类
    
    使用方法:
    1. 初始化: wechat = WeChatIntegration()
    2. 登录: await wechat.login()
    3. 发送通知: await wechat.notify_task_complete(user_id, task_info)
    4. 监听消息: await wechat.start_message_listener(callback)
    """
    
    def __init__(self, config: Optional[WeChatConfig] = None):
        self.config = config or WeChatConfig()
        self.is_logged_in = False
        self.message_cursor = 0
        self.use_mock = False  # 是否使用模拟模式
        
    async def login(self, timeout: int = 300) -> bool:
        """
        登录微信
        
        流程:
        1. 调用 login_qrcode 获取二维码
        2. 显示二维码供用户扫描
        3. 轮询 check_qrcode_status 等待登录
        
        Args:
            timeout: 超时时间（秒）
            
        Returns:
            bool: 是否登录成功
        """
        print("🔐 正在启动微信登录...")
        
        # 步骤1: 获取二维码
        print("📱 正在生成登录二维码...")
        qrcode_result = await self._call_mcp_tool("login_qrcode", {})
        
        if not qrcode_result:
            print("❌ 获取二维码失败")
            return False
            
        # 显示二维码
        await self._display_qrcode()
        
        # 步骤2: 等待扫码
        print("⏳ 请使用微信扫描二维码...")
        print(f"   二维码文本: {self.config.qrcode_path}")
        print(f"   二维码图片: {self.config.qrcode_png_path}")
        
        # 如果是模拟模式，直接返回成功
        if self.use_mock:
            print("✅ 模拟登录成功！")
            self.is_logged_in = True
            return True
        
        start_time = time.time()
        while time.time() - start_time < timeout:
            status = await self._call_mcp_tool("check_qrcode_status", {})
            
            if status.get("status") == "logged_in":
                print("✅ 登录成功！")
                self.is_logged_in = True
                return True
            elif status.get("status") == "expired":
                print("⚠️ 二维码已过期，正在重新生成...")
                await self._call_mcp_tool("login_qrcode", {})
                await self._display_qrcode()
            
            await asyncio.sleep(2)
        
        print("⏰ 登录超时")
        return False
    
    async def _display_qrcode(self):
        """显示二维码"""
        # 尝试显示文本二维码
        if self.config.qrcode_path.exists():
            try:
                qrcode_text = self.config.qrcode_path.read_text()
                print("\n" + "="*40)
                print("微信扫码登录:")
                print(qrcode_text)
                print("="*40 + "\n")
            except Exception as e:
                print(f"   无法读取文本二维码: {e}")
        
        # 提示图片路径
        if self.config.qrcode_png_path.exists():
            print(f"   图片二维码: {self.config.qrcode_png_path}")
    
    async def send_text_message(self, user_id: str, message: str) -> bool:
        """
        发送文本消息
        
        Args:
            user_id: 用户微信ID
            message: 消息内容
            
        Returns:
            bool: 是否发送成功
        """
        if not self.is_logged_in:
            print("❌ 未登录，请先调用 login()")
            return False
        
        result = await self._call_mcp_tool("send_text_message", {
            "user_id": user_id,
            "message": message
        })
        
        return result.get("success", False)
    
    async def send_typing(self, user_id: str) -> bool:
        """
        发送"正在输入..."状态
        
        Args:
            user_id: 用户微信ID
            
        Returns:
            bool: 是否发送成功
        """
        result = await self._call_mcp_tool("send_typing", {
            "user_id": user_id
        })
        return result.get("success", False)
    
    async def notify_task_complete(self, user_id: str, task_info: Dict[str, Any]) -> bool:
        """
        发送任务完成通知
        
        Args:
            user_id: 用户微信ID
            task_info: 任务信息字典
                - task_id: 任务ID
                - task_type: 任务类型
                - duration: 耗时
                - status: 状态 (success/failed)
                - output_files: 输出文件列表
                - error: 错误信息（如果失败）
                
        Returns:
            bool: 是否发送成功
        """
        if task_info.get("status") == "success":
            template = self.config.task_complete_template
        else:
            template = self.config.task_failed_template
        
        # 格式化输出文件列表
        output_files = task_info.get("output_files", [])
        if output_files:
            files_str = "\n".join([f"  • {f}" for f in output_files[:5]])  # 最多显示5个
            if len(output_files) > 5:
                files_str += f"\n  ...还有 {len(output_files) - 5} 个文件"
        else:
            files_str = "  (无输出文件)"
        
        message = template.format(
            task_id=task_info.get("task_id", "未知"),
            task_type=task_info.get("task_type", "未知"),
            duration=task_info.get("duration", "未知"),
            status=task_info.get("status", "未知"),
            output_files=files_str,
            error=task_info.get("error", "")
        )
        
        return await self.send_text_message(user_id, message)
    
    async def get_messages(self, timeout: int = 30) -> list:
        """
        获取新消息（长轮询）
        
        Args:
            timeout: 轮询超时时间（秒）
            
        Returns:
            list: 新消息列表
        """
        if not self.is_logged_in:
            print("❌ 未登录，请先调用 login()")
            return []
        
        result = await self._call_mcp_tool("get_messages", {
            "cursor": self.message_cursor,
            "timeout": timeout
        })
        
        messages = result.get("messages", [])
        if messages:
            # 更新游标
            self.message_cursor = result.get("next_cursor", self.message_cursor)
        
        return messages
    
    async def start_message_listener(self, callback: Callable[[Dict], None], poll_interval: int = 5):
        """
        启动消息监听循环
        
        Args:
            callback: 消息处理回调函数
            poll_interval: 轮询间隔（秒）
        """
        print("🎧 启动消息监听...")
        print(f"   轮询间隔: {poll_interval}秒")
        print("   按 Ctrl+C 停止\n")
        
        try:
            while True:
                messages = await self.get_messages(timeout=poll_interval)
                
                for msg in messages:
                    print(f"📨 收到消息: {msg}")
                    await callback(msg)
                
                await asyncio.sleep(poll_interval)
                
        except KeyboardInterrupt:
            print("\n🛑 消息监听已停止")
    
    async def _call_mcp_tool(self, tool_name: str, params: Dict) -> Dict:
        """
        调用 MCP 工具
        
        首先尝试使用真实的 mcp-wechat-server，如果失败则回退到模拟模式
        """
        try:
            # 尝试真实调用
            return await self._real_mcp_call(tool_name, params)
        except Exception as e:
            # 真实调用失败，使用模拟模式
            if not self.use_mock:
                print(f"   ⚠️  MCP 服务器不可用: {e}")
                print("   切换到模拟模式")
                self.use_mock = True
            return self._mock_call(tool_name, params)
    
    async def _real_mcp_call(self, tool_name: str, params: Dict) -> Dict:
        """真实的 MCP 调用"""
        import subprocess
        import select
        
        # 构建 MCP 请求 (JSON-RPC 格式)
        mcp_request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": params
            },
            "id": 1
        }
        
        # 调用 mcp-wechat-server
        proc = subprocess.Popen(
            ["bunx", "mcp-wechat-server"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        # 发送请求
        request_json = json.dumps(mcp_request) + "\n"
        proc.stdin.write(request_json)
        proc.stdin.flush()
        proc.stdin.close()  # 关闭 stdin 表示请求发送完毕
        
        # 读取响应（增加超时到 30 秒）
        ready, _, _ = select.select([proc.stdout], [], [], 30)
        
        if ready:
            response_line = proc.stdout.readline()
            proc.terminate()
            
            try:
                response = json.loads(response_line)
                if "result" in response:
                    # 解析 MCP 响应格式
                    result = response["result"]
                    if "content" in result and len(result["content"]) > 0:
                        text_content = result["content"][0].get("text", "")
                        # 根据工具类型返回不同格式
                        if tool_name == "login_qrcode":
                            return {"success": True, "message": text_content}
                        elif tool_name == "check_qrcode_status":
                            # 检查是否已登录
                            if "logged in" in text_content.lower() or "success" in text_content.lower():
                                return {"status": "logged_in", "message": text_content}
                            else:
                                return {"status": "waiting", "message": text_content}
                        elif tool_name == "send_text_message":
                            return {"success": "success" in text_content.lower(), "message": text_content}
                        else:
                            return {"success": True, "message": text_content}
                    return result
                elif "error" in response:
                    return {"success": False, "error": response["error"]}
            except json.JSONDecodeError:
                return {"success": False, "error": "Invalid JSON response"}
        else:
            proc.terminate()
            raise TimeoutError("Timeout waiting for MCP response")
    
    def _mock_call(self, tool_name: str, params: Dict) -> Dict:
        """模拟 MCP 调用（用于测试）"""
        if tool_name == "login_qrcode":
            # 创建模拟的二维码文件
            qrcode_dir = Path.home() / ".mcp-wechat-server"
            qrcode_dir.mkdir(exist_ok=True)
            
            # 模拟二维码文本
            qrcode_text = "[模拟二维码] 请想象这里有一个微信登录二维码"
            (qrcode_dir / "qrcode.txt").write_text(qrcode_text)
            
            return {"success": True, "qrcode_url": "weixin://login/mock"}
            
        elif tool_name == "check_qrcode_status":
            # 模拟登录成功
            return {"status": "logged_in"}
            
        elif tool_name == "send_text_message":
            user_id = params.get('user_id', 'unknown')
            message = params.get('message', '')[:50]
            print(f"   [模拟发送消息给 {user_id}]: {message}...")
            return {"success": True}
            
        elif tool_name == "send_typing":
            return {"success": True}
            
        elif tool_name == "get_messages":
            # 模拟返回消息
            return {"messages": [], "next_cursor": self.message_cursor}
        
        return {"success": False, "error": "Unknown tool"}


class WeChatTaskNotifier:
    """
    任务通知器 - 与 bbav20311 任务系统集成
    
    使用示例:
        notifier = WeChatTaskNotifier(user_wechat_id="wxid_xxx")
        await notifier.notify_task_started(task_id="task_001", task_type="identify")
        # ... 任务执行 ...
        await notifier.notify_task_completed(task_id="task_001", result=...)
    """
    
    def __init__(self, user_wechat_id: str, wechat: Optional[WeChatIntegration] = None):
        self.user_wechat_id = user_wechat_id
        self.wechat = wechat or WeChatIntegration()
        self.task_start_times: Dict[str, datetime] = {}
    
    async def notify_task_started(self, task_id: str, task_type: str, **kwargs):
        """通知任务开始"""
        self.task_start_times[task_id] = datetime.now()
        
        message = f"""🚀 任务已启动

📋 任务ID: {task_id}
🔬 分析类型: {task_type}
⏱️ 开始时间: {self.task_start_times[task_id].strftime('%Y-%m-%d %H:%M:%S')}

正在处理中，完成后将通知您..."""
        
        await self.wechat.send_text_message(self.user_wechat_id, message)
    
    async def notify_task_completed(self, task_id: str, output_files: list, **kwargs):
        """通知任务完成"""
        start_time = self.task_start_times.get(task_id)
        duration = "未知"
        if start_time:
            delta = datetime.now() - start_time
            duration = f"{delta.total_seconds() // 60:.0f}分{delta.total_seconds() % 60:.0f}秒"
        
        task_info = {
            "task_id": task_id,
            "task_type": kwargs.get("task_type", "未知"),
            "duration": duration,
            "status": "success",
            "output_files": output_files
        }
        
        await self.wechat.notify_task_complete(self.user_wechat_id, task_info)
    
    async def notify_task_failed(self, task_id: str, error: str, **kwargs):
        """通知任务失败"""
        start_time = self.task_start_times.get(task_id)
        duration = "未知"
        if start_time:
            delta = datetime.now() - start_time
            duration = f"{delta.total_seconds() // 60:.0f}分{delta.total_seconds() % 60:.0f}秒"
        
        task_info = {
            "task_id": task_id,
            "task_type": kwargs.get("task_type", "未知"),
            "duration": duration,
            "status": "failed",
            "output_files": [],
            "error": error
        }
        
        await self.wechat.notify_task_complete(self.user_wechat_id, task_info)


# ==================== 测试代码 ====================

async def test_wechat_integration():
    """测试微信集成"""
    print("="*60)
    print("🧪 bbav20311 微信集成测试")
    print("="*60)
    
    # 1. 初始化
    print("\n📦 步骤1: 初始化微信集成模块")
    wechat = WeChatIntegration()
    
    # 2. 登录
    print("\n🔐 步骤2: 登录测试")
    success = await wechat.login(timeout=10)
    if not success:
        print("❌ 登录失败，跳过后续测试")
        return
    
    # 3. 发送测试消息
    print("\n💬 步骤3: 发送测试消息")
    test_user_id = "test_user_001"
    test_message = "你好！这是来自 bbav20311 的测试消息 🎉"
    
    success = await wechat.send_text_message(test_user_id, test_message)
    print(f"   发送结果: {'✅ 成功' if success else '❌ 失败'}")
    
    # 4. 发送任务通知
    print("\n📊 步骤4: 发送任务完成通知")
    task_info = {
        "task_id": "task_20260410_001",
        "task_type": "基因家族鉴定 (identify)",
        "duration": "5分32秒",
        "status": "success",
        "output_files": [
            "Ft.SPL.identify.gene.list.txt",
            "Ft.SPL.identify.protein.fa",
            "Ft.SPL.identify.report.pdf"
        ]
    }
    
    success = await wechat.notify_task_complete(test_user_id, task_info)
    print(f"   通知结果: {'✅ 成功' if success else '❌ 失败'}")
    
    # 5. 发送失败通知
    print("\n❌ 步骤5: 发送任务失败通知")
    failed_task_info = {
        "task_id": "task_20260410_002",
        "task_type": "进化树分析 (phylogeny)",
        "duration": "2分15秒",
        "status": "failed",
        "output_files": [],
        "error": "输入文件格式错误: 缺少必需的序列数据"
    }
    
    success = await wechat.notify_task_complete(test_user_id, failed_task_info)
    print(f"   通知结果: {'✅ 成功' if success else '❌ 失败'}")
    
    # 6. 测试任务通知器
    print("\n🤖 步骤6: 测试任务通知器集成")
    notifier = WeChatTaskNotifier(user_wechat_id=test_user_id, wechat=wechat)
    
    task_id = "task_20260410_003"
    await notifier.notify_task_started(task_id=task_id, task_type="Motif分析")
    
    # 模拟任务执行时间
    print("   模拟任务执行中...")
    await asyncio.sleep(2)
    
    await notifier.notify_task_completed(
        task_id=task_id,
        task_type="Motif分析",
        output_files=[
            "Ft.SPL.motif.meme.txt",
            "Ft.SPL.motif.logo.png"
        ]
    )
    
    print("\n" + "="*60)
    print("✅ 所有测试完成！")
    print("="*60)


async def test_message_listener():
    """测试消息监听"""
    print("\n🎧 测试消息监听（按 Ctrl+C 停止）")
    
    wechat = WeChatIntegration()
    await wechat.login(timeout=10)
    
    async def message_handler(msg: Dict):
        """处理收到的消息"""
        print(f"   处理消息: {msg}")
        
        # 示例：根据消息内容执行不同操作
        content = msg.get("content", "").lower()
        
        if "状态" in content or "status" in content:
            # 查询任务状态
            reply = "📊 当前系统状态正常，所有服务运行中"
            await wechat.send_text_message(msg.get("from"), reply)
            
        elif "帮助" in content or "help" in content:
            reply = """🤖 bbav20311 基因家族分析助手

可用命令:
• 状态 - 查询系统状态
• 帮助 - 显示此帮助信息
• 任务列表 - 查看进行中的任务
• 详情 <任务ID> - 查看任务详情

直接发送分析需求即可开始任务！"""
            await wechat.send_text_message(msg.get("from"), reply)
    
    await wechat.start_message_listener(callback=message_handler, poll_interval=5)


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "listen":
        # 运行消息监听测试
        asyncio.run(test_message_listener())
    else:
        # 运行基础功能测试
        asyncio.run(test_wechat_integration())