NeoBot/plugins/broadcast.py

# -*- coding: utf-8 -*-
"""
管理员专用的广播插件
功能：
- 仅限管理员在私聊中调用。
- 通过回复一条消息并发送指令，将该消息转发给机器人所在的所有群聊。
- 支持跨机器人广播：当任意机器人接收到广播消息时，会通过 Redis 发布消息，
  所有其他机器人订阅后也会转发给它们各自的群聊。
- 使用通用消息格式，不使用合并转发（聊天记录）格式。
"""
import asyncio
import json
from core.managers.command_manager import matcher
from models.events.message import MessageEvent, PrivateMessageEvent
from core.permission import Permission
from core.utils.logger import logger
from core.managers.redis_manager import redis_manager

# --- 会话状态管理 ---
# 结构: {user_id: asyncio.TimerHandle}
broadcast_sessions: dict[int, asyncio.TimerHandle] = {}

# 广播消息订阅任务
_broadcast_subscription_task = None

def cleanup_session(user_id: int):
    """
    清理超时的广播会话。
    """
    if user_id in broadcast_sessions:
        del broadcast_sessions[user_id]
        logger.info(f"[Broadcast] 会话 {user_id} 已超时，自动取消。")


async def broadcast_message_to_groups(bot, message, source_robot_id: str = "unknown"):
    """
    将消息广播到所有群聊
    
    Args:
        bot: 机器人实例
        message: 要发送的消息
        source_robot_id: 消息来源机器人ID（用于日志）
    """
    try:
        group_list = await bot.get_group_list()
        if not group_list:
            logger.warning(f"[Broadcast] 机器人 {source_robot_id} 目前没有加入任何群聊")
            return
        
        success_count, failed_count = 0, 0
        total_groups = len(group_list)
        
        for group in group_list:
            try:
                await bot.send_group_msg(group.group_id, message)
                success_count += 1
            except Exception as e:
                failed_count += 1
                logger.error(f"[Broadcast] 机器人 {source_robot_id} 发送至群聊 {group.group_id} 失败: {e}")
        
        logger.success(f"[Broadcast] 机器人 {source_robot_id} 广播完成: {total_groups} 个群聊, 成功 {success_count}, 失败 {failed_count}")
        
    except Exception as e:
        logger.error(f"[Broadcast] 机器人 {source_robot_id} 获取群聊列表失败: {e}")


async def start_broadcast_subscription():
    """
    启动 Redis 广播消息订阅
    """
    global _broadcast_subscription_task
    
    if _broadcast_subscription_task is None:
        _broadcast_subscription_task = asyncio.create_task(broadcast_subscription_loop())
        logger.success("[Broadcast] Redis 广播订阅已启动")


async def stop_broadcast_subscription():
    """
    停止 Redis 广播消息订阅
    """
    global _broadcast_subscription_task
    
    if _broadcast_subscription_task:
        _broadcast_subscription_task.cancel()
        try:
            await _broadcast_subscription_task
        except asyncio.CancelledError:
            pass
        _broadcast_subscription_task = None
        logger.info("[Broadcast] Redis 广播订阅已停止")


async def broadcast_subscription_loop():
    """
    Redis 广播消息订阅循环
    """
    if redis_manager.redis is None:
        logger.warning("[Broadcast] Redis 未初始化，无法启动广播订阅")
        return
    
    try:
        pubsub = redis_manager.redis.pubsub()
        await pubsub.subscribe("neobot_broadcast")
        
        logger.success("[Broadcast] 已订阅 Redis 广播频道")
        
        async for message in pubsub.listen():
            if message["type"] == "message":
                try:
                    data = json.loads(message["data"])
                    robot_id = data.get("robot_id", "unknown")
                    message_data = data.get("message")
                    
                    logger.info(f"[Broadcast] 收到跨机器人广播消息: 来源 {robot_id}")
                    
                    # 获取当前机器人的实例
                    from core.ws import WS
                    if WS.instance:
                        await broadcast_message_to_groups(WS.instance, message_data, robot_id)
                    
                except json.JSONDecodeError as e:
                    logger.error(f"[Broadcast] 解析广播消息失败: {e}")
                except Exception as e:
                    logger.error(f"[Broadcast] 处理广播消息失败: {e}")
                    
    except Exception as e:
        logger.error(f"[Broadcast] 广播订阅循环异常: {e}")


@matcher.command("broadcast", "广播", permission=Permission.ADMIN)
async def broadcast_start(event: MessageEvent):
    """
    广播指令的入口，启动一个等待用户消息的会话。
    """
    # 1. 仅限私聊
    if not isinstance(event, PrivateMessageEvent):
        return

    user_id = event.user_id
    
    # 如果上一个会话的超时任务还在，先取消它
    if user_id in broadcast_sessions:
        broadcast_sessions[user_id].cancel()

    await event.reply("已进入广播模式，请在 60 秒内发送您想要广播的消息内容。")
    
    # 设置 60 秒超时
    loop = asyncio.get_running_loop()
    timeout_handler = loop.call_later(
        60, 
        cleanup_session, 
        user_id
    )
    broadcast_sessions[user_id] = timeout_handler
    
    # 确保广播订阅已启动
    await start_broadcast_subscription()

@matcher.on_message()
async def handle_broadcast_content(event: MessageEvent):
    """
    通用消息处理器，用于捕获广播模式下的消息输入。
    将捕获到的消息直接发送给机器人所在的所有群聊，并通过 Redis 发布给其他机器人。
    """
    # 仅处理私聊消息，且用户在广播会话中
    if not isinstance(event, PrivateMessageEvent) or event.user_id not in broadcast_sessions:
        return

    user_id = event.user_id
    
    # 成功捕获到消息，取消超时任务并清理会话
    broadcast_sessions[user_id].cancel()
    del broadcast_sessions[user_id]

    message_to_broadcast = event.message
    if not message_to_broadcast:
        await event.reply("捕获到的消息为空，已取消广播。")
        return True

    # 获取当前机器人ID（使用反向WS的机器人ID）
    from core.ws import WS
    robot_id = "unknown"
    if WS.instance and hasattr(WS.instance, 'self_id'):
        robot_id = str(WS.instance.self_id)
    
    # --- 执行本地广播 ---
    await broadcast_message_to_groups(event.bot, message_to_broadcast, robot_id)
    
    # --- 通过 Redis 发布消息给其他机器人 ---
    try:
        if redis_manager.redis:
            broadcast_data = {
                "robot_id": robot_id,
                "message": message_to_broadcast
            }
            await redis_manager.redis.publish("neobot_broadcast", json.dumps(broadcast_data))
            logger.success(f"[Broadcast] 已通过 Redis 发布广播消息: 来源 {robot_id}")
    except Exception as e:
        logger.error(f"[Broadcast] 发布 Redis 消息失败: {e}")
    
    await event.reply("广播已完成！")
    
    return True # 消费事件，防止其他处理器响应