WechatHookBot/utils/decorators.py

"""
消息处理装饰器模块

提供插件消息处理和定时任务的装饰器
使用工厂模式消除重复代码
"""

from functools import wraps
from typing import Callable, Dict, Union

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger

scheduler = AsyncIOScheduler()


# ==================== 定时任务装饰器 ====================

def schedule(
        trigger: Union[str, CronTrigger, IntervalTrigger],
        **trigger_args
) -> Callable:
    """
    定时任务装饰器

    例子:
    - @schedule('interval', seconds=30)
    - @schedule('cron', hour=8, minute=30, second=30)
    - @schedule('date', run_date='2024-01-01 00:00:00')
    """

    def decorator(func: Callable):
        job_id = f"{func.__module__}.{func.__qualname__}"

        @wraps(func)
        async def wrapper(self, *args, **kwargs):
            return await func(self, *args, **kwargs)

        setattr(wrapper, '_is_scheduled', True)
        setattr(wrapper, '_schedule_trigger', trigger)
        setattr(wrapper, '_schedule_args', trigger_args)
        setattr(wrapper, '_job_id', job_id)

        return wrapper

    return decorator


def add_job_safe(scheduler: AsyncIOScheduler, job_id: str, func: Callable, bot,
                 trigger: Union[str, CronTrigger, IntervalTrigger], **trigger_args):
    """添加函数到定时任务中，如果存在则先删除现有的任务"""
    try:
        scheduler.remove_job(job_id)
    except Exception:
        pass

    # 使用统一配置管理器读取调度器配置
    try:
        from utils.config_manager import get_scheduler_config
        scheduler_config = get_scheduler_config()
    except Exception:
        scheduler_config = {}

    # 应用调度器配置
    job_kwargs = {
        "coalesce": scheduler_config.get("coalesce", True),
        "max_instances": scheduler_config.get("max_instances", 1),
        "misfire_grace_time": scheduler_config.get("misfire_grace_time", 30)
    }
    job_kwargs.update(trigger_args)

    scheduler.add_job(func, trigger, args=[bot], id=job_id, **job_kwargs)


def remove_job_safe(scheduler: AsyncIOScheduler, job_id: str):
    """从定时任务中移除任务"""
    try:
        scheduler.remove_job(job_id)
    except Exception:
        pass


# ==================== 消息装饰器工厂 ====================

def _create_message_decorator(event_type: str, description: str):
    """
    消息装饰器工厂函数

    生成支持两种调用方式的装饰器:
    - @on_xxx_message       (无参数，使用默认优先级50)
    - @on_xxx_message(priority=80)  (有参数，自定义优先级)

    Args:
        event_type: 事件类型字符串，如 'text_message'
        description: 装饰器描述，用于生成文档字符串

    Returns:
        装饰器函数
    """
    def decorator_factory(priority=50):
        def decorator(func):
            # 处理无参数调用: @on_xxx_message 时 priority 实际是被装饰的函数
            if callable(priority):
                target_func = priority
                actual_priority = 50
            else:
                target_func = func
                actual_priority = min(max(priority, 0), 99)

            setattr(target_func, '_event_type', event_type)
            setattr(target_func, '_priority', actual_priority)
            return target_func

        # 判断调用方式
        if callable(priority):
            return decorator(priority)
        return decorator

    decorator_factory.__doc__ = f"{description}装饰器"
    decorator_factory.__name__ = f"on_{event_type}"
    return decorator_factory


# ==================== 消息类型定义 ====================

# 事件类型 -> 中文描述 映射表
MESSAGE_DECORATOR_TYPES: Dict[str, str] = {
    'text_message': '文本消息',
    'image_message': '图片消息',
    'voice_message': '语音消息',
    'video_message': '视频消息',
    'emoji_message': '表情消息',
    'file_message': '文件消息',
    'quote_message': '引用消息',
    'pat_message': '拍一拍',
    'at_message': '@消息',
    'system_message': '系统消息',
    'other_message': '其他消息',
}


# ==================== 生成所有消息装饰器 ====================

# 使用工厂函数生成装饰器
on_text_message = _create_message_decorator('text_message', '文本消息')
on_image_message = _create_message_decorator('image_message', '图片消息')
on_voice_message = _create_message_decorator('voice_message', '语音消息')
on_video_message = _create_message_decorator('video_message', '视频消息')
on_emoji_message = _create_message_decorator('emoji_message', '表情消息')
on_file_message = _create_message_decorator('file_message', '文件消息')
on_quote_message = _create_message_decorator('quote_message', '引用消息')
on_pat_message = _create_message_decorator('pat_message', '拍一拍')
on_at_message = _create_message_decorator('at_message', '@消息')
on_system_message = _create_message_decorator('system_message', '系统消息')
on_other_message = _create_message_decorator('other_message', '其他消息')


# ==================== 导出列表 ====================

__all__ = [
    # 定时任务
    'scheduler',
    'schedule',
    'add_job_safe',
    'remove_job_safe',
    # 消息装饰器
    'on_text_message',
    'on_image_message',
    'on_voice_message',
    'on_video_message',
    'on_emoji_message',
    'on_file_message',
    'on_quote_message',
    'on_pat_message',
    'on_at_message',
    'on_system_message',
    'on_other_message',
    # 工具
    'MESSAGE_DECORATOR_TYPES',
    '_create_message_decorator',
]