plugin.py

"""麦麦失忆插件（MaiBot SDK 2.x）。

提供受权限控制的聊天上下文清理能力：
- /失忆 全部：清理当前聊天流全部消息记录
- /失忆 最近 [数量]：清理当前聊天流最近 N 条消息记录
- /失忆 之前 [小时]：清理当前聊天流指定小时数以前的消息记录
- /失忆 完全：可选的全局记忆清理，默认关闭，并要求二次确认
"""

import asyncio
import time
from datetime import datetime, timedelta
from typing import Any, ClassVar, Dict, List, Set, Tuple, Type

from maibot_sdk import Command, EventHandler, MaiBotPlugin
from maibot_sdk.types import EventType

from .config_models import ContextClearConfig


COMMAND_PATTERN = r"(?:.*，说：\s*)?/(?P<verb>失忆|忘记|断片|amnesia|forget|clear|清除上下文|清空上下文)(?:\s+(?P<args>.*))?$"
TOTAL_MODE_ALIASES = {"total", "完全", "彻底"}
ALL_MODE_ALIASES = {"all", "全部"}
RECENT_MODE_ALIASES = {"recent", "最近"}
BEFORE_MODE_ALIASES = {"before", "之前"}
HELP_MODE_ALIASES = {"help", "帮助"}
CONFIRM_WORD = "确认"


class ContextClearPlugin(MaiBotPlugin):
    """新版 SDK 失忆插件。"""

    plugin_name: ClassVar[str] = "context_clear_plugin"
    plugin_version: ClassVar[str] = "2.0.0"
    config_model: ClassVar[Type[ContextClearConfig]] = ContextClearConfig

    def __init__(self) -> None:
        """初始化插件状态。"""
        super().__init__()
        self._pending_total_confirmations: Dict[str, Dict[str, Any]] = {}

    async def on_load(self) -> None:
        """插件加载。"""
        if not self.config.plugin.enabled:
            self.ctx.logger.info("[context_clear] 已禁用（config.plugin.enabled=False）")
            return
        if not self.config.plugin.permission:
            self.ctx.logger.warning("[context_clear] 未配置 plugin.permission，所有失忆命令都会被拒绝")
        self.ctx.logger.info("[context_clear] 已加载 v%s", self.plugin_version)

    async def on_unload(self) -> None:
        """插件卸载。"""
        self._pending_total_confirmations.clear()
        self.ctx.logger.info("[context_clear] 已卸载")

    async def on_config_update(self, scope: str, config_data: Dict[str, Any], version: str) -> None:
        """配置热更新。"""
        del config_data
        self.ctx.logger.info("[context_clear] 配置已更新 scope=%s version=%s", scope, version)

    def _is_allowed(self, user_id: str) -> bool:
        """检查用户是否有执行权限。"""
        allowed_users = {str(item).strip() for item in self.config.plugin.permission if str(item).strip()}
        return bool(user_id and user_id in allowed_users)

    @staticmethod
    def _message_id_from_message(message: Any) -> str:
        """从 SDK 消息字典中提取消息 ID。"""
        if not isinstance(message, dict):
            return ""
        return str(message.get("message_id") or "").strip()

    @staticmethod
    def _plain_text_from_message(message: Any) -> str:
        """从 SDK 消息字典中提取纯文本。"""
        if not isinstance(message, dict):
            return ""
        return str(message.get("processed_plain_text") or message.get("plain_text") or "").strip()

    @staticmethod
    def _user_id_from_message(message: Any) -> str:
        """从 SDK 消息字典中提取用户 ID。"""
        if not isinstance(message, dict):
            return ""
        message_info = message.get("message_info")
        if not isinstance(message_info, dict):
            return ""
        user_info = message_info.get("user_info")
        if not isinstance(user_info, dict):
            return ""
        return str(user_info.get("user_id") or "").strip()

    @staticmethod
    def _as_int(value: Any) -> int:
        """将能力返回值转换为删除数量。"""
        if isinstance(value, bool):
            return int(value)
        if isinstance(value, int):
            return value
        if isinstance(value, dict):
            for key in ("count", "deleted", "result"):
                if key in value:
                    return ContextClearPlugin._as_int(value.get(key))
        try:
            return int(value)
        except (TypeError, ValueError):
            return 0

    async def _send_notice(self, stream_id: str, text: str) -> None:
        """发送不写入长期消息库的提示。"""
        await self.ctx.send.text(text, stream_id, storage_message=False)

    async def _delete_message_id(self, message_id: str) -> int:
        """按消息 ID 删除一条消息记录。"""
        if not message_id:
            return 0
        result = await self.ctx.db.delete("Messages", {"message_id": message_id})
        return self._as_int(result)

    async def _delete_message_ids(self, message_ids: List[str]) -> int:
        """逐条删除消息记录，避免依赖数据库能力的批量 in 操作。"""
        deleted = 0
        seen: Set[str] = set()
        for message_id in message_ids:
            normalized_id = str(message_id or "").strip()
            if not normalized_id or normalized_id in seen:
                continue
            seen.add(normalized_id)
            deleted += await self._delete_message_id(normalized_id)
        return deleted

    def _cleanup_command_message_later(self, command_message_id: str) -> None:
        """延迟清理触发命令消息。"""
        if not self.config.safety.cleanup_command_message or not command_message_id:
            return

        async def cleanup() -> None:
            await asyncio.sleep(3)
            deleted = await self._delete_message_id(command_message_id)
            if deleted:
                self.ctx.logger.info("[context_clear] 已清理命令消息 message_id=%s", command_message_id)

        asyncio.create_task(cleanup())

    async def _forget_all(self, stream_id: str, command_message_id: str = "") -> int:
        """清理当前聊天流全部消息。"""
        count = await self.ctx.db.count("Messages", {"session_id": stream_id})
        if count <= 0:
            await self._send_notice(stream_id, "🤔 好像...我们之前没聊过天吧？")
            return 0
        deleted = self._as_int(await self.ctx.db.delete("Messages", {"session_id": stream_id}))
        self._cleanup_command_message_later(command_message_id)
        await self._send_notice(
            stream_id,
            f"💫 *咚* 诶...你们是谁？我怎么在这里？\n\n（麦麦忘记了当前聊天的 {deleted} 条记忆）",
        )
        self.ctx.logger.info("[context_clear] 已清理当前聊天全部消息 stream_id=%s deleted=%s", stream_id, deleted)
        return deleted

    async def _forget_recent(self, stream_id: str, count: int, command_message_id: str = "") -> int:
        """清理当前聊天流最近 N 条消息。"""
        max_count = self.config.limits.max_recent_count
        normalized_count = max(1, min(count, max_count))
        messages = await self.ctx.db.get(
            "Messages",
            filters={"session_id": stream_id},
            limit=normalized_count,
            order_by="-timestamp",
        )
        if not isinstance(messages, list) or not messages:
            await self._send_notice(stream_id, "🤔 咦...我们刚才说了什么吗？")
            return 0
        message_ids = [str(item.get("message_id") or "") for item in messages if isinstance(item, dict)]
        deleted = await self._delete_message_ids(message_ids)
        self._cleanup_command_message_later(command_message_id)
        await self._send_notice(stream_id, f"😵 诶？刚才发生了什么？\n\n（麦麦忘记了最近 {deleted} 条记忆）")
        self.ctx.logger.info(
            "[context_clear] 已清理最近消息 stream_id=%s requested=%s deleted=%s",
            stream_id,
            normalized_count,
            deleted,
        )
        return deleted

    async def _forget_before_hours(self, stream_id: str, hours: int, command_message_id: str = "") -> int:
        """清理当前聊天流指定小时数以前的消息。"""
        normalized_hours = max(1, hours)
        threshold_time = (datetime.now() - timedelta(hours=normalized_hours)).timestamp()
        messages = await self.ctx.message.get_by_time_in_chat(
            stream_id,
            "0",
            str(threshold_time),
            limit=self.config.limits.max_time_scan_messages,
            limit_mode="earliest",
        )
        if not isinstance(messages, list) or not messages:
            await self._send_notice(stream_id, f"🤔 {normalized_hours} 小时前？那时候我们有聊过吗...")
            return 0
        message_ids = [str(item.get("message_id") or "") for item in messages if isinstance(item, dict)]
        deleted = await self._delete_message_ids(message_ids)
        self._cleanup_command_message_later(command_message_id)
        await self._send_notice(
            stream_id,
            f"😌 嗯...{normalized_hours} 小时前的事都是浮云~\n\n（麦麦忘记了 {deleted} 条久远的记忆）",
        )
        self.ctx.logger.info("[context_clear] 已按时间清理 stream_id=%s hours=%s deleted=%s", stream_id, normalized_hours, deleted)
        return deleted

    def _cleanup_expired_confirmations(self) -> None:
        """清理过期完全失忆确认状态。"""
        now = time.time()
        timeout = self.config.safety.confirm_timeout
        expired_users = [
            user_id
            for user_id, pending in self._pending_total_confirmations.items()
            if now - float(pending.get("timestamp", 0)) > timeout
        ]
        for user_id in expired_users:
            self._pending_total_confirmations.pop(user_id, None)

    async def _request_total_confirmation(self, user_id: str, stream_id: str) -> None:
        """发起完全失忆确认。"""
        if not self.config.safety.allow_total_amnesia:
            await self._send_notice(
                stream_id,
                "⚠️ 完全失忆当前未启用。\n\n如确需清理全局记忆，请先在插件配置中开启 safety.allow_total_amnesia。",
            )
            return
        self._cleanup_expired_confirmations()
        timeout = self.config.safety.confirm_timeout
        self._pending_total_confirmations[user_id] = {
            "timestamp": time.time(),
            "stream_id": stream_id,
        }
        await self._send_notice(
            stream_id,
            "⚠️ **危险警告！**\n\n"
            "这将清理麦麦的全局记忆数据，包括聊天记录、聊天流、人物印象、表达学习、长期聊天历史、黑话和工具记录。\n"
            "不会修改本地文件、统计数据或配置文件。\n\n"
            "💥 **这是不可逆操作。**\n\n"
            f"如果确认，请在 {timeout} 秒内发送 `/失忆 完全 确认`，或直接回复 `确认`。",
        )

    async def _confirm_total(self, user_id: str, stream_id: str) -> Tuple[bool, str]:
        """确认并执行完全失忆。"""
        if not self.config.safety.allow_total_amnesia:
            self._pending_total_confirmations.pop(user_id, None)
            return False, "⚠️ 完全失忆当前未启用。"
        self._cleanup_expired_confirmations()
        pending = self._pending_total_confirmations.get(user_id)
        if not pending:
            return False, "❌ 没有待确认的完全失忆请求。请先发送 `/失忆 完全`。"
        if str(pending.get("stream_id") or "") != stream_id:
            return False, "❌ 请在发起请求的同一个聊天中确认。"
        elapsed = time.time() - float(pending.get("timestamp", 0))
        if elapsed > self.config.safety.confirm_timeout:
            self._pending_total_confirmations.pop(user_id, None)
            return False, "⏰ 确认已超时，请重新发送 `/失忆 完全`。"
        self._pending_total_confirmations.pop(user_id, None)
        deleted_stats = await self._execute_total_amnesia()
        total_deleted = sum(deleted_stats.values())
        await self._send_notice(
            stream_id,
            f"💫✨ *完全失忆完成* ✨💫\n\n我...我是谁？这里是哪里？\n\n📊 清除了 {total_deleted} 项记忆数据",
        )
        return True, f"已执行完全失忆，清除 {total_deleted} 项数据"

    async def _execute_total_amnesia(self) -> Dict[str, int]:
        """执行受控的全局记忆清理。"""
        model_names = [
            "Messages",
            "ChatSession",
            "PersonInfo",
            "Expression",
            "ChatHistory",
            "Jargon",
            "ToolRecord",
        ]
        stats: Dict[str, int] = {}
        for model_name in model_names:
            result = await self.ctx.db.query(model_name, query_type="delete", filters={})
            stats[model_name] = self._as_int(result)
        self.ctx.logger.warning("[context_clear] 已执行完全失忆 stats=%s", stats)
        return stats

    async def _show_help(self, stream_id: str) -> None:
        """显示帮助。"""
        await self._send_notice(
            stream_id,
            "💫 麦麦失忆插件\n\n"
            "用法：\n"
            "/失忆 全部 - 忘记当前聊天的所有消息记录\n"
            "/失忆 最近 [数量] - 忘记当前聊天最近 N 条消息（默认 10）\n"
            "/失忆 之前 [小时] - 忘记当前聊天 N 小时以前的消息\n"
            "/失忆 完全 - 全局记忆清理（默认禁用，需配置开启并二次确认）\n"
            "/失忆 帮助 - 显示此帮助\n\n"
            "别名：/忘记 /断片 /amnesia /forget /clear /清除上下文 /清空上下文",
        )

    @Command(
        "context_clear_command",
        description="让麦麦选择性失忆：/失忆 全部|最近|之前|完全|帮助",
        pattern=COMMAND_PATTERN,
    )
    async def handle_context_clear(
        self,
        stream_id: str = "",
        user_id: str = "",
        matched_groups: Any = None,
        message: Any = None,
        **kwargs: Any,
    ) -> Tuple[bool, str, bool]:
        """处理失忆命令。"""
        del kwargs
        if not self.config.plugin.enabled:
            return False, "插件已禁用", True
        if not stream_id:
            return False, "缺少 stream_id", True
        if not user_id:
            user_id = self._user_id_from_message(message)
        if not self._is_allowed(user_id):
            return False, "没有权限", True

        groups = matched_groups if isinstance(matched_groups, dict) else {}
        args_text = str(groups.get("args") or "").strip()
        parts = args_text.split()
        subcommand = parts[0] if parts else "help"
        command_message_id = self._message_id_from_message(message)

        try:
            if subcommand in HELP_MODE_ALIASES:
                await self._show_help(stream_id)
                return True, "显示帮助", True
            if subcommand in ALL_MODE_ALIASES:
                deleted = await self._forget_all(stream_id, command_message_id)
                return True, f"已清理当前聊天 {deleted} 条消息", True
            if subcommand in RECENT_MODE_ALIASES:
                count = self.config.limits.default_recent_count
                if len(parts) > 1:
                    count = int(parts[1])
                deleted = await self._forget_recent(stream_id, count, command_message_id)
                return True, f"已清理最近 {deleted} 条消息", True
            if subcommand in BEFORE_MODE_ALIASES:
                hours = int(parts[1]) if len(parts) > 1 else 24
                deleted = await self._forget_before_hours(stream_id, hours, command_message_id)
                return True, f"已清理 {hours} 小时前的 {deleted} 条消息", True
            if subcommand in TOTAL_MODE_ALIASES:
                if len(parts) > 1 and parts[1] == CONFIRM_WORD:
                    success, result = await self._confirm_total(user_id, stream_id)
                    return success, result, True
                await self._request_total_confirmation(user_id, stream_id)
                return True, "等待完全失忆确认", True
        except ValueError:
            await self._send_notice(stream_id, "参数必须是整数。使用 `/失忆 帮助` 查看用法。")
            return False, "参数错误", True
        except Exception as exc:  # noqa: BLE001 - 命令入口需要给用户统一回执
            self.ctx.logger.exception("[context_clear] 命令执行失败")
            await self._send_notice(stream_id, f"😖 失忆失败了...\n\n错误：{exc}")
            return False, str(exc), True

        await self._send_notice(stream_id, f"🤨 {subcommand}？这是什么意思...\n使用 `/失忆 帮助` 查看用法。")
        return False, "未知子命令", True

    @EventHandler(
        "context_clear_confirm_handler",
        description="监听完全失忆的直接确认消息",
        event_type=EventType.ON_MESSAGE,
        intercept_message=True,
        weight=100,
    )
    async def handle_confirm_message(self, message: Any = None, stream_id: str = "", **kwargs: Any) -> Dict[str, Any]:
        """处理直接回复“确认”的二次确认消息。"""
        del kwargs
        if not self.config.plugin.enabled:
            return {"continue_processing": True}
        text = self._plain_text_from_message(message)
        if text != CONFIRM_WORD:
            return {"continue_processing": True}
        user_id = self._user_id_from_message(message)
        target_stream_id = stream_id or (str(message.get("session_id") or "") if isinstance(message, dict) else "")
        pending = self._pending_total_confirmations.get(user_id)
        if not pending or str(pending.get("stream_id") or "") != target_stream_id:
            return {"continue_processing": True}
        if not self._is_allowed(user_id):
            return {"continue_processing": True}

        success, result = await self._confirm_total(user_id, target_stream_id)
        if success:
            self._cleanup_command_message_later(self._message_id_from_message(message))
        else:
            await self._send_notice(target_stream_id, result)
        return {
            "continue_processing": False,
            "custom_result": result,
        }


def create_plugin() -> ContextClearPlugin:
    """SDK 加载入口。"""
    return ContextClearPlugin()