zhenxun_bot/zhenxun/models/bot_console.py

from typing import Literal, overload

from tortoise import fields

from zhenxun.services.cache.runtime_cache import BotMemoryCache
from zhenxun.services.db_context import Model
from zhenxun.services.db_context.schema_ops import AlterColumnType, RenameColumn
from zhenxun.utils.enum import CacheType


class BotConsole(Model):
    id = fields.IntField(pk=True, generated=True, auto_increment=True)
    """自增id"""
    bot_id = fields.CharField(255, unique=True, description="bot_id")
    """bot_id"""
    status = fields.BooleanField(default=True, description="Bot状态")
    """Bot状态"""
    create_time = fields.DatetimeField(auto_now_add=True, description="创建时间")
    """创建时间"""
    platform = fields.CharField(255, null=True, description="平台")
    """平台"""
    block_plugins = fields.TextField(default="", description="禁用插件")
    """禁用插件"""
    block_tasks = fields.TextField(default="", description="禁用被动技能")
    """禁用被动技能"""
    available_plugins = fields.TextField(default="", description="可用插件")
    """可用插件"""
    available_tasks = fields.TextField(default="", description="可用被动技能")
    """可用被动技能管理镜像，不作为运行白名单。"""

    class Meta:  # pyright: ignore [reportIncompatibleVariableOverride]
        table = "bot_console"
        table_description = "Bot数据表"

    cache_type = CacheType.BOT
    """缓存类型"""
    cache_key_field = "bot_id"
    """缓存键字段"""

    @staticmethod
    def format(name: str) -> str:
        return f"<{name},"

    @overload
    @classmethod
    async def get_bot_status(cls) -> list[tuple[str, bool]]: ...

    @overload
    @classmethod
    async def get_bot_status(cls, bot_id: str) -> bool: ...

    @classmethod
    async def get_bot_status(
        cls, bot_id: str | None = None
    ) -> list[tuple[str, bool]] | bool:
        """
        获取bot状态

        参数:
            bot_id (str, optional): bot_id. Defaults to None.

        返回:
            list[tuple[str, bool]] | bool: bot状态
        """
        if not bot_id:
            data = await BotMemoryCache.get_all()
            return [(bot_id, snapshot.status) for bot_id, snapshot in data.items()]
        result = await BotMemoryCache.get(bot_id)
        return result.status if result else False

    @overload
    @classmethod
    async def get_tasks(cls) -> list[tuple[str, list[str]]]: ...

    @overload
    @classmethod
    async def get_tasks(cls, bot_id: str) -> list[str]: ...

    @overload
    @classmethod
    async def get_tasks(cls, *, status: bool) -> dict[str, list[str]]: ...

    @overload
    @classmethod
    async def get_tasks(cls, bot_id: str, status: bool = True) -> list[str]: ...

    @classmethod
    async def get_tasks(cls, bot_id: str | None = None, status: bool | None = True):
        """
        获取bot被动技能管理镜像。

        available_tasks 只服务管理命令和展示，不参与运行白名单判断。
        被动运行真源是 TaskInfo.status/load_status、BotConsole.block_tasks、
        GroupConsole.block_task/superuser_block_task。

        参数:
            bot_id (str | None, optional): bot_id. Defaults to None.
            status (bool | None, optional): 被动状态. Defaults to True.

        返回:
            list[tuple[str, str]] | str: 被动技能
        """
        if not bot_id:
            data = await BotMemoryCache.get_all()
            task_attr = "available_tasks" if status else "block_tasks"
            return {
                bot_id: cls.convert_module_format(getattr(snapshot, task_attr))
                for bot_id, snapshot in data.items()
            }
        result = await BotMemoryCache.get(bot_id)
        if result:
            tasks = result.available_tasks if status else result.block_tasks
            return cls.convert_module_format(tasks)
        return []

    @overload
    @classmethod
    async def get_plugins(cls) -> dict[str, list[str]]: ...

    @overload
    @classmethod
    async def get_plugins(cls, bot_id: str) -> list[str]: ...

    @overload
    @classmethod
    async def get_plugins(cls, *, status: bool) -> dict[str, list[str]]: ...

    @overload
    @classmethod
    async def get_plugins(cls, bot_id: str, status: bool = True) -> list[str]: ...

    @classmethod
    async def get_plugins(cls, bot_id: str | None = None, status: bool = True):
        """
        获取bot插件

        参数:
            bot_id (str | None, optional): bot_id. Defaults to None.
            status (bool, optional): 插件状态. Defaults to True.

        返回:
            list[tuple[str, str]] | str: 插件
        """
        if not bot_id:
            data = await BotMemoryCache.get_all()
            plugin_attr = "available_plugins" if status else "block_plugins"
            return {
                bot_id: cls.convert_module_format(getattr(snapshot, plugin_attr))
                for bot_id, snapshot in data.items()
            }

        result = await BotMemoryCache.get(bot_id)
        if result:
            plugins = result.available_plugins if status else result.block_plugins
            return cls.convert_module_format(plugins)
        return []

    @classmethod
    async def set_bot_status(cls, status: bool, bot_id: str | None = None) -> None:
        """
        设置bot状态

        参数:
            status (bool): 状态
            bot_id (str, optional): bot_id. Defaults to None.

        Raises:
            ValueError: 未找到 bot_id
        """
        if bot_id:
            affected_rows = await cls.filter(bot_id=bot_id).update(status=status)
            if not affected_rows:
                raise ValueError(f"未找到 bot_id: {bot_id}")
            await BotMemoryCache.update_status(bot_id, status)
        else:
            await cls.all().update(status=status)
            await BotMemoryCache.refresh()

    @overload
    @classmethod
    def convert_module_format(cls, data: str) -> list[str]: ...

    @overload
    @classmethod
    def convert_module_format(cls, data: list[str]) -> str: ...

    @classmethod
    def convert_module_format(cls, data: str | list[str]) -> str | list[str]:
        """
        在 `<aaa,<bbb,<ccc,` 和 `["aaa", "bbb", "ccc"]` 之间进行相互转换。

        参数:
            data (str | list[str]): 输入数据，可能是格式化字符串或字符串列表。

        返回:
            str | list[str]: 根据输入类型返回转换后的数据。
        """
        if isinstance(data, str):
            return [item.strip(",") for item in data.split("<") if item]
        elif isinstance(data, list):
            return "".join(cls.format(item) for item in data)

    @classmethod
    async def _toggle_field(
        cls,
        bot_id: str,
        from_field: str,
        to_field: str,
        data: str,
    ) -> None:
        """
        在 from_field 和 to_field 之间移动指定的 data

        参数:
            bot_id (str): 目标 bot 的 ID
            from_field (str): 源字段名称
            to_field (str): 目标字段名称
            data (str): 要插入的内容

        Raises:
            ValueError: 如果 data 不在 from_field 和 to_field 中
        """
        bot_data, _ = await cls.get_or_create(bot_id=bot_id)
        formatted_data = cls.format(data)

        from_list: str = getattr(bot_data, from_field)
        to_list: str = getattr(bot_data, to_field)

        if formatted_data not in (from_list + to_list):
            raise ValueError(f"{data} 不在源字段和目标字段中")

        if formatted_data in from_list:
            from_list = from_list.replace(formatted_data, "", 1)
            if formatted_data not in to_list:
                to_list += formatted_data

        setattr(bot_data, from_field, from_list)
        setattr(bot_data, to_field, to_list)

        await bot_data.save(update_fields=[from_field, to_field])

    @classmethod
    async def disable_plugin(cls, bot_id: str | None, plugin_name: str) -> None:
        """
        禁用插件

        参数:
            bot_id (str | None): bot_id
            plugin_name (str): 插件名称
        """
        if bot_id:
            await cls._toggle_field(
                bot_id,
                "available_plugins",
                "block_plugins",
                plugin_name,
            )
        else:
            bot_list = await cls.all()
            for bot in bot_list:
                await cls._toggle_field(
                    bot.bot_id,
                    "available_plugins",
                    "block_plugins",
                    plugin_name,
                )

    @classmethod
    async def enable_plugin(cls, bot_id: str | None, plugin_name: str) -> None:
        """
        启用插件

        参数:
            bot_id (str | None): bot_id
            plugin_name (str): 插件名称
        """
        if bot_id:
            await cls._toggle_field(
                bot_id,
                "block_plugins",
                "available_plugins",
                plugin_name,
            )
        else:
            bot_list = await cls.all()
            for bot in bot_list:
                await cls._toggle_field(
                    bot.bot_id,
                    "block_plugins",
                    "available_plugins",
                    plugin_name,
                )

    @classmethod
    async def disable_task(cls, bot_id: str | None, task_name: str) -> None:
        """
        禁用被动技能

        参数:
            bot_id (str | None): bot_id
            task_name (str): 被动技能名称
        """
        if bot_id:
            await cls._toggle_field(
                bot_id,
                "available_tasks",
                "block_tasks",
                task_name,
            )
        else:
            bot_list = await cls.all()
            for bot in bot_list:
                await cls._toggle_field(
                    bot.bot_id,
                    "available_tasks",
                    "block_tasks",
                    task_name,
                )

    @classmethod
    async def enable_task(cls, bot_id: str | None, task_name: str) -> None:
        """
        启用被动技能

        参数:
            bot_id (str | None): bot_id
            task_name (str): 被动技能名称
        """
        if bot_id:
            await cls._toggle_field(
                bot_id,
                "block_tasks",
                "available_tasks",
                task_name,
            )
        else:
            bot_list = await cls.all()
            for bot in bot_list:
                await cls._toggle_field(
                    bot.bot_id,
                    "block_tasks",
                    "available_tasks",
                    task_name,
                )

    @classmethod
    async def disable_all(
        cls,
        bot_id: str,
        feat: Literal["plugins", "tasks"],
    ) -> None:
        """
        禁用全部插件或被动技能

        参数:
            bot_id (str): bot_id
            feat (Literal["plugins", "tasks"]): 插件或被动技能
        """
        bot_data, _ = await cls.get_or_create(bot_id=bot_id)
        if feat == "plugins":
            available_plugins = cls.convert_module_format(bot_data.available_plugins)
            block_plugins = cls.convert_module_format(bot_data.block_plugins)
            bot_data.block_plugins = cls.convert_module_format(
                available_plugins + block_plugins
            )
            bot_data.available_plugins = ""
        elif feat == "tasks":
            available_tasks = cls.convert_module_format(bot_data.available_tasks)
            block_tasks = cls.convert_module_format(bot_data.block_tasks)
            bot_data.block_tasks = cls.convert_module_format(
                available_tasks + block_tasks
            )
            bot_data.available_tasks = ""
        await bot_data.save(
            update_fields=[
                "available_tasks",
                "block_tasks",
                "available_plugins",
                "block_plugins",
            ]
        )

    @classmethod
    async def enable_all(
        cls,
        bot_id: str,
        feat: Literal["plugins", "tasks"],
    ) -> None:
        """
        启用全部插件或被动技能

        参数:
            bot_id (str): bot_id
            feat (Literal["plugins", "tasks"]): 插件或被动技能
        """
        bot_data, _ = await cls.get_or_create(bot_id=bot_id)
        if feat == "plugins":
            available_plugins = cls.convert_module_format(bot_data.available_plugins)
            block_plugins = cls.convert_module_format(bot_data.block_plugins)
            bot_data.available_plugins = cls.convert_module_format(
                available_plugins + block_plugins
            )
            bot_data.block_plugins = ""
        elif feat == "tasks":
            available_tasks = cls.convert_module_format(bot_data.available_tasks)
            block_tasks = cls.convert_module_format(bot_data.block_tasks)
            bot_data.available_tasks = cls.convert_module_format(
                available_tasks + block_tasks
            )
            bot_data.block_tasks = ""
        await bot_data.save(
            update_fields=[
                "available_tasks",
                "block_tasks",
                "available_plugins",
                "block_plugins",
            ]
        )

    @classmethod
    async def is_block_plugin(cls, bot_id: str, plugin_name: str) -> bool:
        """
        检查插件是否被禁用

        参数:
            bot_id (str): bot_id
            plugin_name (str): 插件某款

        返回:
            bool: 是否被禁用
        """
        bot_data = await BotMemoryCache.get(bot_id)
        if not bot_data:
            return False
        return cls.format(plugin_name) in bot_data.block_plugins

    @classmethod
    async def is_block_task(cls, bot_id: str, task_name: str) -> bool:
        """
        检查被动技能是否被禁用

        参数:
            bot_id (str): bot_id
            task_name (str): 被动技能名称

        返回:
            bool: 是否被禁用
        """
        bot_data = await BotMemoryCache.get(bot_id)
        if not bot_data:
            return False
        return cls.format(task_name) in bot_data.block_tasks

    @classmethod
    async def create(cls, *args, **kwargs):
        result = await super().create(*args, **kwargs)
        await BotMemoryCache.upsert_from_model(result)
        return result

    @classmethod
    async def update_or_create(cls, *args, **kwargs):
        result = await super().update_or_create(*args, **kwargs)
        await BotMemoryCache.upsert_from_model(result[0])
        return result

    async def save(self, *args, **kwargs):
        await super().save(*args, **kwargs)
        await BotMemoryCache.upsert_from_model(self)

    async def delete(self, *args, **kwargs):
        bot_id = self.bot_id
        await super().delete(*args, **kwargs)
        await BotMemoryCache.remove(bot_id)

    @classmethod
    async def _run_script(cls):
        return [
            RenameColumn("bot_console", "block_plugin", "block_plugins"),
            RenameColumn("bot_console", "block_task", "block_tasks"),
            AlterColumnType("bot_console", "block_plugins", "TEXT"),
            AlterColumnType("bot_console", "block_tasks", "TEXT"),
        ]