#!/usr/bin/env python3
"""
微信服务 - 独立运行的微信机器人服务

功能：
1. 接收用户消息
2. 调用 bbav20311 API 执行任务
3. 发送任务完成通知

运行方式：
    python wechat_service.py

或后台运行：
    nohup python wechat_service.py > wechat_service.log 2>&1 &
"""

import json
import subprocess
import select
import time
import sys
import os
import requests
from pathlib import Path
from datetime import datetime

# 配置
YOUR_USER_ID = "o9cq801A7b5SQiVmMTKpOFN5DX4g@im.wechat"
BBAV20311_API = "http://localhost:7865"  # bbav20311 API 地址


def call_mcp_tool(tool_name: str, params: dict):
    """调用 MCP 工具"""
    mcp_request = {
        "jsonrpc": "2.0",
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": params
        },
        "id": 1
    }
    
    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()
    
    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:
                    return result["content"][0].get("text", "")
        except:
            pass
    else:
        proc.terminate()
    return None


def send_message(to_user: str, text: str):
    """发送消息"""
    return call_mcp_tool("send_text_message", {
        "to": to_user,
        "text": text
    })


def get_messages():
    """获取消息"""
    return call_mcp_tool("get_messages", {})


# 设置日志
import logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('/home/xia-lab/wutx/00test/bbav5/wechat_service.log', mode='a')
    ]
)
logger = logging.getLogger(__name__)

# 导入持久化存储
from session_storage import get_storage, init_storage

# 初始化持久化存储
storage = init_storage()

# 内存缓存（从数据库加载）
_running_tasks_cache = {}
_user_sessions_cache = {}
_user_active_jobs_cache = {}

def _load_from_storage():
    """从持久化存储加载数据到内存缓存"""
    global _running_tasks_cache
    _running_tasks_cache = storage.get_running_tasks()


# 启动时加载数据
_load_from_storage()

# 保持向后兼容的引用
running_tasks = _running_tasks_cache

def get_user_history(user_id: str) -> list:
    """获取用户历史任务"""
    return storage.get_user_history(user_id, limit=10)

def add_user_task(user_id: str, task_info: dict):
    """添加任务到用户历史"""
    storage.add_user_task(user_id, task_info)

def get_latest_task(user_id: str) -> dict:
    """获取用户最近一个任务"""
    return storage.get_latest_task(user_id)

def format_task_status(task: dict) -> str:
    """格式化任务状态"""
    status_emoji = {
        'pending': '⏳',
        'running': '🔄',
        'completed': '✅',
        'failed': '❌'
    }
    emoji = status_emoji.get(task.get('status', 'unknown'), '❓')
    return f"{emoji} {task.get('prefix', 'Unknown')} ({task.get('stage', 'unknown')}) - {task.get('status', 'unknown')}"


def get_user_active_job(user_id: str) -> dict:
    """获取用户当前活跃 job"""
    return storage.get_user_active_job(user_id)


def set_user_active_job(user_id: str, job_info: dict):
    """设置用户当前活跃 job"""
    storage.set_user_active_job(user_id, job_info)


def update_job_stage(user_id: str, stage: str, status: str):
    """更新 job 的阶段状态"""
    storage.update_job_stage(user_id, stage, status)


def get_completed_stages(user_id: str) -> list:
    """获取已完成的阶段列表"""
    job = get_user_active_job(user_id)
    if job and 'stages' in job:
        stages = job['stages']
        return [stage for stage, status in stages.items() if status == 'completed']
    return []


def format_job_pipeline(user_id: str) -> str:
    """格式化分析流程图"""
    job = get_user_active_job(user_id)
    if not job:
        return "暂无活跃任务"
    
    stages_order = ['identify', 'phylogeny', 'motif', 'domain', 'structure', 'promoter', 'synteny', 'ppi', 'kaks', 'expression']
    stages_name = {
        'identify': '鉴定',
        'phylogeny': '进化树',
        'motif': 'Motif',
        'domain': '结构域',
        'structure': '基因结构',
        'promoter': '启动子',
        'synteny': '共线性',
        'ppi': 'PPI网络',
        'kaks': 'KaKs分析',
        'expression': '表达分析'
    }
    
    pipeline = ""
    stages = job.get('stages', {})
    for stage in stages_order:
        status = stages.get(stage, 'pending')
        emoji = {'completed': '✅', 'running': '🔄', 'failed': '❌', 'pending': '⬜'}.get(status, '⬜')
        name = stages_name.get(stage, stage)
        pipeline += f"{emoji} {name} → "
    
    return pipeline.rstrip(" → ")

def poll_task_status(task_id: str, job_id: str, prefix: str, stage: str, user_id: str):
    """轮询任务状态并发送完成通知"""
    logger.info(f"开始轮询任务状态: task_id={task_id}")
    max_polls = 60  # 最多轮询60次（5分钟）
    poll_count = 0
    
    while poll_count < max_polls:
        try:
            response = requests.get(
                f"{BBAV20311_API}/api/task/{task_id}",
                timeout=10
            )
            data = response.json()
            status = data.get('status')
            
            logger.info(f"任务 {task_id} 状态: {status}")
            
            # 更新任务状态（内存和持久化）
            if task_id in running_tasks:
                running_tasks[task_id]['status'] = status
                storage.update_running_task(task_id, status)
            
            if status == 'completed' or status == 'success':
                # 任务完成，发送通知
                result = data.get('result', {})
                manifest = result.get('manifest', {})
                metrics = manifest.get('metrics', {})
                final_members = metrics.get('final_members', 0)
                elapsed = manifest.get('elapsed_sec', 0)
                
                reply = f"""✅ 分析完成！

📋 Job ID: {job_id}
🧬 基因家族: {prefix}
🔬 分析类型: {stage}
📊 找到成员: {final_members} 个
⏱️ 耗时: {elapsed:.1f} 秒

结果文件位置:
/workspace/{job_id}/final/"""
                send_message(user_id, reply)
                logger.info(f"任务完成通知已发送: {task_id}")
                
                # 更新持久化存储中的任务状态
                storage.update_task_status(user_id, task_id, 'completed', final_members)
                storage.update_running_task(task_id, 'completed', final_members)
                
                # 从运行任务列表中移除
                storage.remove_running_task(task_id)
                if task_id in running_tasks:
                    del running_tasks[task_id]
                return
                
            elif status == 'failed':
                # 任务失败
                error_msg = data.get('message', '未知错误')
                reply = f"""❌ 分析失败

📋 Job ID: {job_id}
🧬 基因家族: {prefix}
🔬 分析类型: {stage}
❌ 错误: {error_msg}"""
                send_message(user_id, reply)
                logger.error(f"任务失败: {task_id}, 错误: {error_msg}")
                
                # 更新持久化存储中的任务状态
                storage.update_task_status(user_id, task_id, 'failed')
                storage.update_running_task(task_id, 'failed')
                
                storage.remove_running_task(task_id)
                if task_id in running_tasks:
                    del running_tasks[task_id]
                return
            
            # 继续轮询
            time.sleep(5)
            poll_count += 1
            
        except Exception as e:
            logger.error(f"轮询任务状态失败: {str(e)}")
            time.sleep(5)
            poll_count += 1
    
    # 超时
    logger.warning(f"任务轮询超时: {task_id}")
    reply = f"⏰ 任务 {job_id} 执行时间较长，请稍后通过前端查看结果"
    send_message(user_id, reply)
    storage.remove_running_task(task_id)
    if task_id in running_tasks:
        del running_tasks[task_id]


def create_task(prefix: str, stage: str = "identify", user_id: str = None, reuse_job: bool = True):
    """创建 bbav20311 任务"""
    try:
        user_id = user_id or YOUR_USER_ID
        
        # 检查是否应该复用现有job
        active_job = get_user_active_job(user_id) if reuse_job else None
        
        if active_job and active_job.get('prefix') == prefix:
            # 复用现有job
            job_id = active_job.get('job_id')
            logger.info(f"复用现有job: {job_id}")
            
            # 检查该阶段是否已完成
            if active_job.get('stages', {}).get(stage) == 'completed':
                return {
                    "error": f"该阶段已完成",
                    "job_id": job_id,
                    "stage": stage,
                    "already_completed": True
                }
        else:
            # 1. 创建新任务
            logger.info(f"创建新任务...")
            response = requests.post(
                f"{BBAV20311_API}/api/new_job",
                timeout=10
            )
            job_data = response.json()
            job_id = job_data.get("job_id")
            logger.info(f"新任务ID: {job_id}")

            if not job_id:
                return {"error": "Failed to create job"}

            # 2. 创建符号链接到测试数据
            source_input = "/home/xia-lab/wutx/00test/bbav5/input/job_1775821345494"
            target_input = f"/home/xia-lab/wutx/00test/bbav5/input/{job_id}"

            if os.path.exists(source_input):
                if not os.path.exists(target_input):
                    os.symlink(source_input, target_input)
                    logger.info(f"已创建符号链接")
                elif len(os.listdir(target_input)) == 0:
                    os.rmdir(target_input)
                    os.symlink(source_input, target_input)
            
            # 设置用户活跃job
            set_user_active_job(user_id, {
                'job_id': job_id,
                'prefix': prefix,
                'created_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'stages': {}
            })

        # 3. 执行分析任务
        logger.info(f"执行分析任务: job_id={job_id}, stage={stage}, prefix={prefix}")
        
        # 更新阶段状态为运行中
        update_job_stage(user_id, stage, 'running')
        response = requests.post(
            f"{BBAV20311_API}/api/execute",
            json={
                "job_id": job_id,
                "stage": stage,
                "prefix": prefix,
                "wechat_user_id": user_id or YOUR_USER_ID
            },
            timeout=30
        )
        result = response.json()
        logger.info(f"执行结果: {result}")
        
        # 4. 如果是异步任务，启动后台轮询
        if result.get('async') and result.get('task_id'):
            task_id = result.get('task_id')
            task_info = {
                'task_id': task_id,
                'job_id': job_id,
                'prefix': prefix,
                'stage': stage,
                'status': 'running',
                'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'user_id': user_id or YOUR_USER_ID
            }
            running_tasks[task_id] = task_info
            
            # 添加到持久化存储
            storage.add_running_task(task_id, task_info)
            add_user_task(user_id or YOUR_USER_ID, task_info)
            
            # 在后台线程中轮询任务状态
            import threading
            poll_thread = threading.Thread(
                target=poll_task_status,
                args=(task_id, job_id, prefix, stage, user_id or YOUR_USER_ID),
                daemon=True
            )
            poll_thread.start()
            logger.info(f"已启动任务状态轮询线程: {task_id}")
        
        return result
    except Exception as e:
        logger.error(f"创建任务失败: {str(e)}")
        import traceback
        traceback.print


def handle_message(msg: dict):
    """处理收到的消息"""
    from_user = msg.get("from_user_id", "")
    text = msg.get("text", "").strip()
    
    print(f"[{datetime.now().strftime('%H:%M:%S')}] 📨 收到: {text}")
    
    # 只处理你的消息
    if from_user != YOUR_USER_ID:
        return
    
    # 处理命令
    if text == "你是谁":
        reply = """🤖 我是 bbav20311 基因家族分析助手

我可以帮你：
• 运行基因家族分析任务
• 查询任务状态
• 发送任务完成通知

发送"帮助"查看可用命令。"""
        send_message(from_user, reply)
        
    elif text in ["帮助", "help"]:
        reply = """📋 可用命令：

🧬 分析任务：
• "分析 Ft.SPL" - 运行基因家族鉴定
• "分析 Ft.WRKY 进化树" - 指定分析类型
• "分析 Ft.SPL PPI网络" - 蛋白互作网络分析
• "分析 Ft.SPL KaKs" - 选择压力分析
• "分析 Ft.SPL 表达" - RNA-seq表达分析
• "状态" - 查询系统状态

🔬 支持的分析类型：
• 鉴定 (identify)
• 进化树 (phylogeny)
• Motif分析 (motif)
• 结构域 (domain)
• 基因结构 (structure)
• 启动子 (promoter)
• 共线性 (synteny)
• PPI网络 (ppi)
• KaKs分析 (kaks)
• 表达分析 (expression)

💬 其他：
• "你是谁" - 查看介绍
• "帮助" - 显示此帮助
• "流程" - 查看当前分析进度
• "继续分析" - 查看可继续的阶段

直接发送基因家族名称即可开始分析！"""
        send_message(from_user, reply)
        
    elif text == "状态":
        # 获取用户的任务统计
        stats = storage.get_user_stats(from_user)
        running_count = stats.get('running', 0)
        completed_count = stats.get('completed', 0)
        reply = f"""📊 系统状态

✅ 微信服务：在线
✅ 后端服务：运行中

📋 您的任务：
  • 进行中：{running_count} 个
  • 已完成：{completed_count} 个
  • 失败：{failed_count} 个

发送"历史"查看任务详情
发送"最新"查看最新任务状态"""
        send_message(from_user, reply)
        
    elif text in ["历史", "任务", "tasks"]:
        history = get_user_history(from_user)
        if not history:
            reply = "📭 您还没有执行过分析任务\n\n发送「分析 Ft.SPL」开始第一个任务！"
        else:
            reply = "📋 您的任务历史（最近10个）：\n\n"
            for i, task in enumerate(reversed(history[-10:]), 1):
                reply += f"{i}. {format_task_status(task)}\n"
                reply += f"   📁 {task.get('job_id', 'N/A')}\n"
                reply += f"   🕐 {task.get('timestamp', 'N/A')}\n\n"
        send_message(from_user, reply)
        
    elif text in ["最新", "最近", "last"]:
        latest = get_latest_task(from_user)
        if not latest:
            reply = "📭 您还没有执行过分析任务"
        else:
            reply = f"""📋 最新任务状态

🧬 基因家族: {latest.get('prefix', 'Unknown')}
🔬 分析类型: {latest.get('stage', 'unknown')}
📁 Job ID: {latest.get('job_id', 'N/A')}
{format_task_status(latest)}
🕐 提交时间: {latest.get('timestamp', 'N/A')}"""
            
            # 如果任务还在运行，显示进度
            if latest.get('status') == 'running':
                reply += "\n\n⏳ 任务正在执行中，请稍后..."
            elif latest.get('status') == 'completed':
                final_members = latest.get('final_members', 'N/A')
                reply += f"\n📊 找到成员: {final_members} 个"
        send_message(from_user, reply)
        
    elif text in ["分析完了吗", "完成了吗", "好了吗"]:
        # 查询最新任务状态
        latest = get_latest_task(from_user)
        if not latest:
            reply = "📭 您还没有执行过分析任务\n\n发送「分析 Ft.SPL」开始分析！"
        elif latest.get('status') == 'running':
            reply = f"""⏳ 任务执行中...

🧬 基因家族: {latest.get('prefix', 'Unknown')}
🔬 分析类型: {latest.get('stage', 'unknown')}
📁 Job ID: {latest.get('job_id', 'N/A')}
🕐 提交时间: {latest.get('timestamp', 'N/A')}

请稍后再查询，或发送"最新"查看状态"""
        elif latest.get('status') == 'completed':
            final_members = latest.get('final_members', 'N/A')
            reply = f"""✅ 分析已完成！

🧬 基因家族: {latest.get('prefix', 'Unknown')}
🔬 分析类型: {latest.get('stage', 'unknown')}
📁 Job ID: {latest.get('job_id', 'N/A')}
📊 找到成员: {final_members} 个

结果文件位置:
/workspace/{latest.get('job_id', 'N/A')}/final/"""
        elif latest.get('status') == 'failed':
            reply = f"""❌ 分析失败

🧬 基因家族: {latest.get('prefix', 'Unknown')}
🔬 分析类型: {latest.get('stage', 'unknown')}
📁 Job ID: {latest.get('job_id', 'N/A')}

请检查参数后重试"""
        send_message(from_user, reply)
        
    elif text in ["流程", "pipeline", "进度"]:
        # 显示当前分析流程图
        pipeline = format_job_pipeline(from_user)
        active_job = get_user_active_job(from_user)
        
        if active_job:
            completed = get_completed_stages(from_user)
            reply = f"""📊 当前分析流程

🧬 基因家族: {active_job.get('prefix', 'Unknown')}
📁 Job ID: {active_job.get('job_id', 'N/A')}

{pipeline}

✅ 已完成: {', '.join(completed) if completed else '无'}

发送「继续分析」查看可执行的阶段"""
        else:
            reply = "📭 暂无活跃任务\n\n发送「分析 Ft.SPL」开始分析！"
        send_message(from_user, reply)
        
    elif text in ["继续分析", "下一步", "下一阶段"]:
        # 显示可以继续分析的阶段
        active_job = get_user_active_job(from_user)
        
        if not active_job:
            reply = "📭 暂无活跃任务\n\n发送「分析 Ft.SPL」开始分析！"
        else:
            completed = get_completed_stages(from_user)
            prefix = active_job.get('prefix', 'Unknown')
            
            # 定义阶段依赖关系
            stage_deps = {
                'phylogeny': ['identify'],
                'motif': ['phylogeny'],
                'domain': ['phylogeny'],
                'structure': ['phylogeny'],
                'promoter': ['phylogeny'],
                'synteny': ['identify'],
                'ppi': ['identify'],
                'kaks': ['identify'],
                'expression': ['identify']
            }
            
            # 找出可以执行的阶段
            available_stages = []
            for stage, deps in stage_deps.items():
                if stage not in completed:
                    if all(d in completed for d in deps):
                        available_stages.append(stage)
            
            if not available_stages:
                reply = f"""✅ 所有阶段已完成！

🧬 基因家族: {prefix}
📁 Job ID: {active_job.get('job_id', 'N/A')}

已完成阶段:
{', '.join(completed)}

可以发送「分析 Ft.XXX」开始新的基因家族分析"""
            else:
                stages_name = {
                    'phylogeny': '进化树分析',
                    'motif': 'Motif分析',
                    'domain': '结构域分析',
                    'structure': '基因结构分析',
                    'promoter': '启动子分析',
                    'synteny': '共线性分析',
                    'ppi': 'PPI网络分析',
                    'kaks': 'KaKs选择压力分析',
                    'expression': 'RNA-seq表达分析'
                }
                
                reply = f"""📋 可继续分析的阶段

🧬 基因家族: {prefix}
📁 Job ID: {active_job.get('job_id', 'N/A')}
✅ 已完成: {', '.join(completed) if completed else '无'}

可执行阶段:
"""
                for i, stage in enumerate(available_stages, 1):
                    name = stages_name.get(stage, stage)
                    reply += f"{i}. 「分析 {prefix} {name}」\n"
                
                reply += "\n直接发送上述指令即可继续分析"
        
        send_message(from_user, reply)
        
    elif "分析" in text:
        # 提取基因家族名称
        import re
        match = re.search(r'Ft\.[A-Za-z]+', text)
        if match:
            gene_family = match.group(0)
            
            # 提取分析类型
            stage = "identify"  # 默认
            text_lower = text.lower()
            if "进化树" in text or "phylogeny" in text:
                stage = "phylogeny"
            elif "motif" in text:
                stage = "motif"
            elif "domain" in text:
                stage = "domain"
            elif "结构" in text:
                stage = "structure"
            elif "启动子" in text or "promoter" in text:
                stage = "promoter"
            elif "共线性" in text or "synteny" in text:
                stage = "synteny"
            elif "ppi" in text_lower or "蛋白互作" in text or "互作网络" in text:
                stage = "ppi"
            elif "kaks" in text_lower or "选择压力" in text or "ka/ks" in text_lower:
                stage = "kaks"
            elif "表达" in text or "expression" in text_lower or "rna" in text_lower or "rna-seq" in text_lower:
                stage = "expression"
            
            # 检查是否有活跃job且基因家族相同
            active_job = get_user_active_job(from_user)
            if active_job and active_job.get('prefix') == gene_family:
                reply = f"""🚀 继续分析：{gene_family}
🔬 分析类型：{stage}
📁 使用现有 Job: {active_job.get('job_id')}

任务已提交，正在处理中...
完成后将通知您。"""
            else:
                reply = f"""🚀 新建分析：{gene_family}
🔬 分析类型：{stage}

任务已创建，正在处理中...
完成后将通知您。"""
            
            send_message(from_user, reply)

            # 创建实际任务（会自动复用现有job）
            result = create_task(gene_family, stage, from_user, reuse_job=True)
            print(f"任务创建结果: {result}")

            # 发送任务创建结果
            if "error" in result:
                if result.get('already_completed'):
                    send_message(from_user, f"⚠️ 该阶段已完成！\n📁 Job ID: {result.get('job_id')}\n发送「继续分析」查看其他可执行阶段")
                else:
                    send_message(from_user, f"❌ 任务创建失败: {result['error']}")
            else:
                job_id = result.get('job_id', 'unknown')
                active_job = get_user_active_job(from_user)
                if active_job and len(active_job.get('stages', {})) > 1:
                    completed = get_completed_stages(from_user)
                    send_message(from_user, f"✅ 任务已提交\n📋 Job ID: {job_id}\n✅ 已完成: {', '.join(completed)}\n⏳ 请稍后查看结果")
                else:
                    send_message(from_user, f"✅ 任务已提交\n📋 Job ID: {job_id}\n⏳ 请稍后查看结果")
        else:
            reply = "❓ 无法识别基因家族名称，请使用格式如：Ft.SPL、Ft.WRKY"
            send_message(from_user, reply)
    else:
        reply = f"🤖 收到：{text}\n\n发送'帮助'查看可用命令，或直接发送'分析 Ft.SPL'开始分析。"
        send_message(from_user, reply)


def main():
    """主循环"""
    print("="*60)
    print("🤖 bbav20311 微信服务")
    print("="*60)
    print(f"监听用户: {YOUR_USER_ID}")
    print(f"API地址: {BBAV20311_API}")
    print("按 Ctrl+C 停止\n")
    
    last_message_id = 0
    
    try:
        while True:
            result = get_messages()
            
            if result:
                try:
                    data = json.loads(result)
                    messages = data.get("messages", [])
                    
                    for msg in messages:
                        msg_id = msg.get("message_id", 0)
                        
                        # 只处理新消息
                        if msg_id > last_message_id:
                            last_message_id = msg_id
                            handle_message(msg)
                            
                except json.JSONDecodeError:
                    pass
            
            time.sleep(2)
            
    except KeyboardInterrupt:
        print("\n\n🛑 服务已停止")
        sys.exit(0)


if __name__ == "__main__":
    main()