PlantAI-dev/notify_integration.py

#!/usr/bin/env python3
"""
微信通知集成 - 用于 bbav20311 任务系统

使用方法:
    from notify_integration import notify_task_start, notify_task_complete
    
    # 任务开始时
    await notify_task_start(user_wechat_id, task_id, task_type)
    
    # 任务完成时
    await notify_task_complete(user_wechat_id, task_id, task_type, output_files)
"""

import asyncio
import json
import subprocess
import select
from pathlib import Path
from typing import List, Optional
from datetime import datetime

# 你的微信ID（用于接收通知）
# 需要替换为你的真实微信ID
YOUR_WECHAT_ID = "wxid_xxxxxxxxxxxxxxxx"  # 需要替换


def call_mcp_tool(tool_name: str, params: dict) -> Optional[dict]:
    """调用 MCP 工具"""
    mcp_request = {
        "jsonrpc": "2.0",
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": params
        },
        "id": 1
    }
    
    try:
        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()
        
        # 等待响应（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:
                    result = response["result"]
                    if "content" in result and len(result["content"]) > 0:
                        text = result["content"][0].get("text", "")
                        return {"success": True, "text": text}
            except:
                pass
        else:
            proc.terminate()
    except Exception as e:
        print(f"MCP 调用失败: {e}")
    
    return None


def check_login_status() -> bool:
    """检查是否已登录"""
    result = call_mcp_tool("check_qrcode_status", {})
    if result:
        text = result.get("text", "").lower()
        return "logged in" in text or "success" in text
    return False


def send_text_message(user_id: str, message: str) -> bool:
    """发送文本消息"""
    result = call_mcp_tool("send_text_message", {
        "user_id": user_id,
        "message": message
    })
    return result is not None and result.get("success", False)


def notify_task_start(user_wechat_id: str, task_id: str, task_type: str) -> bool:
    """
    通知任务开始
    
    使用示例:
        notify_task_start("wxid_xxx", "task_001", "identify")
    """
    if not check_login_status():
        print("⚠️  微信未登录，跳过通知")
        return False
    
    message = f"""🚀 任务已启动

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

正在处理中，完成后将通知您..."""
    
    return send_text_message(user_wechat_id, message)


def notify_task_complete(
    user_wechat_id: str, 
    task_id: str, 
    task_type: str,
    output_files: List[str],
    duration: str = "未知",
    error: str = ""
) -> bool:
    """
    通知任务完成
    
    使用示例:
        notify_task_complete(
            "wxid_xxx", 
            "task_001", 
            "identify",
            ["result.txt", "plot.png"],
            duration="5分30秒"
        )
    """
    if not check_login_status():
        print("⚠️  微信未登录，跳过通知")
        return False
    
    if error:
        # 失败通知
        message = f"""❌ 任务失败通知

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

请检查输入数据或联系管理员"""
    else:
        # 成功通知
        files_str = "\n".join([f"  • {f}" for f in output_files[:5]])
        if len(output_files) > 5:
            files_str += f"\n  ...还有 {len(output_files) - 5} 个文件"
        
        message = f"""🎉 任务完成通知

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

📁 输出文件:
{files_str}

💡 提示: 回复"详情 {task_id}"查看完整结果"""
    
    return send_text_message(user_wechat_id, message)


# ==================== 测试 ====================

def test_notification():
    """测试通知功能"""
    print("="*60)
    print("🧪 测试微信通知功能")
    print("="*60)
    
    # 检查登录状态
    print("\n🔍 检查登录状态...")
    if check_login_status():
        print("✅ 已登录")
    else:
        print("❌ 未登录")
        print("请先运行 wechat_integration.py 登录")
        return
    
    # 测试发送消息
    # 注意：需要替换为你的真实微信ID
    test_user_id = YOUR_WECHAT_ID
    
    if test_user_id == "wxid_xxxxxxxxxxxxxxxx":
        print("\n⚠️  请修改 YOUR_WECHAT_ID 为你的真实微信ID")
        print("   文件位置: notify_integration.py")
        return
    
    print(f"\n💬 发送测试消息到 {test_user_id}...")
    success = send_text_message(test_user_id, "🧪 这是来自 bbav20311 的测试消息！")
    
    if success:
        print("✅ 消息发送成功")
    else:
        print("❌ 消息发送失败")


if __name__ == "__main__":
    test_notification()