zhenxun_bot/zhenxun/utils/depends/__init__.py

from typing import Any, Literal

from nonebot.adapters import Bot, Event
from nonebot.internal.params import Depends
from nonebot.matcher import Matcher
from nonebot.params import Command
from nonebot.permission import SUPERUSER
from nonebot_plugin_session import EventSession
from nonebot_plugin_uninfo import Uninfo

from zhenxun.configs.config import Config
from zhenxun.utils.limiters import ConcurrencyLimiter, FreqLimiter, RateLimiter
from zhenxun.utils.message import MessageUtils
from zhenxun.utils.time_utils import TimeUtils

_coolers: dict[str, FreqLimiter] = {}
_rate_limiters: dict[str, RateLimiter] = {}
_concurrency_limiters: dict[str, ConcurrencyLimiter] = {}


def _create_limiter_dependency(
    limiter_class: type,
    limiter_storage: dict,
    limiter_init_args: dict[str, Any],
    scope: Literal["user", "group", "global"],
    prompt: str,
    **kwargs,
):
    """
    一个高阶函数，用于创建不同类型的限制器依赖。

    参数:
        limiter_class: 限制器类 (FreqLimiter, RateLimiter, etc.).
        limiter_storage: 用于存储限制器实例的字典.
        limiter_init_args: 限制器类的初始化参数.
        scope: 限制作用域.
        prompt: 触发限制时的提示信息.
        **kwargs: 传递给特定限制器逻辑的额外参数.
    """

    async def dependency(
        matcher: Matcher, session: EventSession, bot: Bot, event: Event
    ) -> bool:
        if await SUPERUSER(bot, event):
            return True

        handler_id = (
            f"{matcher.plugin_name}:{matcher.handlers[0].call.__code__.co_firstlineno}"
        )

        key: str | None = None
        if scope == "user":
            key = session.id1
        elif scope == "group":
            key = session.id3 or session.id2 or session.id1
        elif scope == "global":
            key = f"global_{handler_id}"

        if not key:
            return True

        if handler_id not in limiter_storage:
            limiter_storage[handler_id] = limiter_class(**limiter_init_args)
        limiter = limiter_storage[handler_id]

        if isinstance(limiter, ConcurrencyLimiter):
            await limiter.acquire(key)
            matcher.state["_concurrency_limiter_info"] = {
                "limiter": limiter,
                "key": key,
            }
            return True
        else:
            if limiter.check(key):
                if isinstance(limiter, FreqLimiter):
                    limiter.start_cd(
                        key, kwargs.get("duration_sec", limiter.default_cd)
                    )
                return True
            else:
                left_time = limiter.left_time(key)
                format_kwargs = {
                    "cd_str": TimeUtils.format_duration(left_time),
                    **(kwargs.get("prompt_format_kwargs", {})),
                }
                message = prompt.format(**format_kwargs)
                await matcher.finish(message)

    return Depends(dependency)


def Cooldown(
    duration: str,
    *,
    scope: Literal["user", "group", "global"] = "user",
    prompt: str = "操作过于频繁，请等待 {cd_str}",
) -> bool:
    """声明式冷却检查依赖，限制用户操作频率

    参数:
        duration: 冷却时间字符串 (e.g., "30s", "10m", "1h")
        scope: 冷却作用域
        prompt: 自定义的冷却提示消息，可使用 {cd_str} 占位符

    返回:
        bool: 是否允许执行
    """
    try:
        parsed_seconds = TimeUtils.parse_time_string(duration)
    except ValueError as e:
        raise ValueError(f"Cooldown装饰器中的duration格式错误: {e}")

    return _create_limiter_dependency(
        limiter_class=FreqLimiter,
        limiter_storage=_coolers,
        limiter_init_args={"default_cd_seconds": parsed_seconds},
        scope=scope,
        prompt=prompt,
        duration_sec=parsed_seconds,
    )


def RateLimit(
    count: int,
    duration: str,
    *,
    scope: Literal["user", "group", "global"] = "user",
    prompt: str = "太快了，在 {duration_str} 内只能触发{limit}次，请等待 {cd_str}",
) -> bool:
    """声明式速率限制依赖，在指定时间窗口内限制操作次数

    参数:
        count: 在时间窗口内允许的最大调用次数
        duration: 时间窗口字符串 (e.g., "1m", "1h")
        scope: 限制作用域
        prompt: 自定义的提示消息，可使用 {cd_str}, {duration_str}, {limit} 占位符

    返回:
        bool: 是否允许执行
    """
    try:
        parsed_seconds = TimeUtils.parse_time_string(duration)
    except ValueError as e:
        raise ValueError(f"RateLimit装饰器中的duration格式错误: {e}")

    return _create_limiter_dependency(
        limiter_class=RateLimiter,
        limiter_storage=_rate_limiters,
        limiter_init_args={"max_calls": count, "time_window": parsed_seconds},
        scope=scope,
        prompt=prompt,
        prompt_format_kwargs={"duration_str": duration, "limit": count},
    )


def ConcurrencyLimit(
    count: int,
    *,
    scope: Literal["user", "group", "global"] = "global",
    prompt: str | None = "当前功能繁忙，请稍后再试...",
) -> bool:
    """声明式并发数限制依赖，控制某个功能同时执行的实例数量

    参数:
        count: 最大并发数
        scope: 限制作用域
        prompt: 提示消息（暂未使用，主要用于未来扩展超时功能）

    返回:
        bool: 是否允许执行
    """
    return _create_limiter_dependency(
        limiter_class=ConcurrencyLimiter,
        limiter_storage=_concurrency_limiters,
        limiter_init_args={"max_concurrent": count},
        scope=scope,
        prompt=prompt or "",
    )


def CheckUg(check_user: bool = True, check_group: bool = True):
    """检测群组id和用户id是否存在

    参数:
        check_user: 检查用户id.
        check_group: 检查群组id.
    """

    async def dependency(session: EventSession):
        if check_user:
            user_id = session.id1
            if not user_id:
                await MessageUtils.build_message("用户id为空").finish()
        if check_group:
            group_id = session.id3 or session.id2
            if not group_id:
                await MessageUtils.build_message("群组id为空").finish()

    return Depends(dependency)


def OneCommand():
    """
    获取单个命令Command
    """

    async def dependency(
        cmd: tuple[str, ...] = Command(),
    ):
        return cmd[0] if cmd else None

    return Depends(dependency)


def UserName():
    """
    用户名称
    """

    async def dependency(user_info: Uninfo):
        return user_info.user.nick or user_info.user.name or ""

    return Depends(dependency)


def GetConfig(
    module: str | None = None,
    config: str = "",
    default_value: Any = None,
    prompt: str | None = None,
):
    """获取配置项

    参数:
        module: 模块名，为空时默认使用当前插件模块名
        config: 配置项名称
        default_value: 默认值
        prompt: 为空时提示
    """

    async def dependency(matcher: Matcher):
        module_ = module or matcher.plugin_name
        if module_:
            value = Config.get_config(module_, config, default_value)
            if value is None and prompt:
                await matcher.finish(prompt)
            return value

    return Depends(dependency)


def CheckConfig(
    module: str | None = None,
    config: str | list[str] = "",
    prompt: str | None = None,
):
    """检测配置项在配置文件中是否填写

    参数:
        module: 模块名，为空时默认使用当前插件模块名
        config: 需要检查的配置项名称
        prompt: 为空时提示
    """

    async def dependency(matcher: Matcher):
        module_ = module or matcher.plugin_name
        if module_:
            config_list = [config] if isinstance(config, str) else config
            for c in config_list:
                if Config.get_config(module_, c) is None:
                    await matcher.finish(prompt or f"配置项 {c} 未填写！")

    return Depends(dependency)