NeoBot/plugins/bot_status.py

"""
Bot 状态查询插件

提供 /status 指令，以图片形式展示机器人当前的综合运行状态。
"""
import os
import psutil
import time
import asyncio
from datetime import datetime, timedelta

from core.bot import Bot
from core.managers.command_manager import matcher
from core.managers.image_manager import image_manager
from core.managers.redis_manager import redis_manager
from core.utils.executor import run_in_thread_pool
from core.utils.logger import logger
from models.events.message import MessageEvent, MessageSegment
from models.objects import LoginInfo, Status, VersionInfo

__plugin_meta__ = {
    "name": "bot_status",
    "description": "以图片形式展示机器人当前的综合运行状态",
    "usage": "/status 或 /状态",
}

# 记录机器人启动时间
START_TIME = time.time()
# 获取当前进程
PROCESS = psutil.Process(os.getpid())

def _get_system_info():
    """
    同步函数：使用 psutil 获取系统信息，避免阻塞事件循环。
    """
    try:
        # interval=1 会阻塞1秒，必须在线程池中运行
        cpu_percent = psutil.cpu_percent(interval=1)
        mem_info = psutil.virtual_memory()
        bot_mem_mb = PROCESS.memory_info().rss / (1024 * 1024)
        
        # 磁盘信息
        disk_usage = psutil.disk_usage('/')
        
        # 网络信息
        net_io = psutil.net_io_counters()
        
        return {
            "cpu_percent": f"{cpu_percent:.1f}",
            "mem_percent": f"{mem_info.percent:.1f}",
            "mem_total": f"{mem_info.total / (1024**3):.1f}",
            "mem_used": f"{mem_info.used / (1024**3):.1f}",
            "bot_mem_mb": f"{bot_mem_mb:.2f}",
            "disk_percent": f"{disk_usage.percent:.1f}",
            "disk_total": f"{disk_usage.total / (1024**3):.1f}",
            "disk_used": f"{disk_usage.used / (1024**3):.1f}",
            "net_sent": f"{net_io.bytes_sent / (1024**2):.1f}",
            "net_recv": f"{net_io.bytes_recv / (1024**2):.1f}",
        }
    except Exception as e:
        logger.error(f"获取系统信息失败: {e}")
        return {
            "cpu_percent": "N/A",
            "mem_percent": "N/A",
            "mem_total": "N/A",
            "mem_used": "N/A",
            "bot_mem_mb": "N/A",
            "disk_percent": "N/A",
            "disk_total": "N/A",
            "disk_used": "N/A",
            "net_sent": "N/A",
            "net_recv": "N/A",
        }

@matcher.command("status", "状态")
async def handle_status(bot: Bot, event: MessageEvent, args: list[str]):
    """
    处理 status 指令，生成并回复机器人状态图片。
    """
    logger.info(f"收到用户 {event.user_id} 的状态查询指令，开始生成状态图...")

    try:
        # 1. 获取API信息 (增加独立错误处理)
        # 尝试获取或更新 bot.nickname
        if not hasattr(bot, "nickname") or not bot.nickname:
            try:
                # 优先使用 get_stranger_info 获取自身信息，比 get_login_info 更轻量
                stranger_info = await bot.get_stranger_info(user_id=bot.self_id)
                bot.nickname = stranger_info.nickname
            except Exception as e:
                logger.warning(f"获取 stranger_info 失败: {e}, 将回退到 login_info")
                try:
                    login_info = await bot.get_login_info()
                    bot.nickname = login_info.nickname
                except Exception as e2:
                    logger.warning(f"获取 login_info 也失败了: {e2}")
                    bot.nickname = "获取失败"
        
        nickname = bot.nickname

        # 状态信息：如果能响应此命令，说明机器人必然在线且状态良好
        # 这避免了依赖可能超时或未实现的 get_status API
        logger.debug("正在推断机器人状态...")
        status_info = Status(online=True, good=True)
        logger.debug(f"推断状态成功: online={status_info.online}, good={status_info.good}")

        try:
            version_info = await bot.get_version_info()
        except Exception as e:
            logger.warning(f"获取 version_info 失败: {e}")
            version_info = VersionInfo(app_name="获取失败", app_version="N/A", protocol_version="N/A")

        # 2. 计算运行时长
        uptime_seconds = int(time.time() - START_TIME)
        uptime_delta = timedelta(seconds=uptime_seconds)
        days = uptime_delta.days
        hours, remainder = divmod(uptime_delta.seconds, 3600)
        minutes, seconds = divmod(remainder, 60)
        uptime_str = f"{days}天 {hours:02}:{minutes:02}:{seconds:02}"

        bot_info_data = {
            "user_id": bot.self_id,
            "nickname": nickname,
            "avatar_url": f"https://q1.qlogo.cn/g?b=qq&nk={bot.self_id}&s=640",
            "start_time": datetime.fromtimestamp(START_TIME).strftime("%Y-%m-%d %H:%M:%S"),
            "uptime": uptime_str,
        }

        # 3. 获取统计数据
        try:
            msgs_recv = await redis_manager.get("neobot:stats:messages_received") or 0
            msgs_sent = await redis_manager.get("neobot:stats:messages_sent") or 0
            command_stats_raw = await redis_manager.redis.hgetall("neobot:command_stats")
            
            total_commands = sum(int(v) for v in command_stats_raw.values())
            
            stats_data = {
                "messages_received": int(msgs_recv),
                "messages_sent": int(msgs_sent),
                "total_commands": total_commands,
            }

            command_stats_data = sorted(
                [{"name": k, "count": int(v)} for k, v in command_stats_raw.items()],
                key=lambda x: x["count"],
                reverse=True
            )
        except Exception as e:
            logger.error(f"获取Redis统计数据失败: {e}")
            stats_data = {
                "messages_received": 0,
                "messages_sent": 0,
                "total_commands": 0,
            }
            command_stats_data = []

        # 4. 异步获取系统信息
        # 设置超时，防止 psutil 阻塞过久
        try:
            system_data = await asyncio.wait_for(run_in_thread_pool(_get_system_info), timeout=5.0)
        except asyncio.TimeoutError:
            logger.error("获取系统信息超时")
            system_data = {
                "cpu_percent": "Timeout",
                "mem_percent": "Timeout",
                "mem_total": "Timeout",
                "mem_used": "Timeout",
                "bot_mem_mb": "Timeout",
                "disk_percent": "Timeout",
                "disk_total": "Timeout",
                "disk_used": "Timeout",
                "net_sent": "Timeout",
                "net_recv": "Timeout",
            }
        except Exception as e:
            logger.error(f"获取系统信息异常: {e}")
            system_data = {
                "cpu_percent": "Error",
                "mem_percent": "Error",
                "mem_total": "Error",
                "mem_used": "Error",
                "bot_mem_mb": "Error",
                "disk_percent": "Error",
                "disk_total": "Error",
                "disk_used": "Error",
                "net_sent": "Error",
                "net_recv": "Error",
            }

        # 5. 准备模板所需的所有数据
        template_data = {
            "bot_info": bot_info_data,
            "status_info": status_info,
            "version_info": version_info,
            "stats": stats_data,
            "system": system_data,
            "command_stats": command_stats_data,
        }

        # 6. 渲染图片
        try:
            base64_str = await image_manager.render_template_to_base64(
                template_name="status.html",
                data=template_data,
                output_name="status.png",
                image_type="png"
            )

            if base64_str:
                await event.reply(MessageSegment.image(base64_str))
            else:
                # 如果渲染失败，image_manager 内部会记录错误，这里给用户一个通用提示
                await event.reply("状态图片生成失败，可能是渲染服务出现问题，请联系管理员。")
        except Exception as e:
            logger.error(f"渲染图片失败: {e}")
            await event.reply("状态图片渲染过程中发生错误。")

    except Exception as e:
        logger.exception(f"生成状态图时发生意外错误, 用户: {event.user_id}")
        await event.reply(f"获取状态信息时发生未知错误，请稍后再试或联系管理员。")