"""
用户鉴权中间件

检查用户是否在白名单中
"""

from functools import wraps
from typing import Callable, Any

from telegram import Update
from telegram.ext import ContextTypes


class AuthMiddleware:
    """
    鉴权中间件类
    """

    def __init__(self, allowed_users: set, admin_users: set = None):
        """
        初始化鉴权中间件

        参数:
            allowed_users: 允许使用的用户 ID 集合
            admin_users: 管理员用户 ID 集合
        """
        self.allowed_users = allowed_users or set()
        self.admin_users = admin_users or set()

    def is_allowed(self, user_id: int) -> bool:
        """检查用户是否允许使用"""
        # 如果白名单为空，允许所有人
        if not self.allowed_users:
            return True
        return user_id in self.allowed_users or user_id in self.admin_users

    def is_admin(self, user_id: int) -> bool:
        """检查用户是否是管理员"""
        return user_id in self.admin_users


def require_auth(func: Callable) -> Callable:
    """
    鉴权装饰器

    用于命令处理函数，检查用户是否在白名单中

    用法:
        @require_auth
        async def my_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
            ...
    """

    @wraps(func)
    async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs) -> Any:
        user = update.effective_user
        if not user:
            return

        user_id = user.id
        allowed_users = context.bot_data.get("allowed_users", set())
        admin_users = context.bot_data.get("admin_users", set())

        # 如果白名单为空，允许所有人使用
        if not allowed_users and not admin_users:
            return await func(update, context, *args, **kwargs)

        # 检查是否在白名单或管理员列表中
        if user_id in allowed_users or user_id in admin_users:
            return await func(update, context, *args, **kwargs)

        # 未授权
        await update.message.reply_text(
            f"⛔ 你没有权限使用此 Bot\n\n"
            f"你的用户 ID: `{user_id}`\n\n"
            f"请联系管理员将你的 ID 添加到白名单",
            parse_mode="Markdown"
        )
        return None

    return wrapper


def require_admin(func: Callable) -> Callable:
    """
    管理员权限装饰器

    用于需要管理员权限的命令

    用法:
        @require_admin
        async def admin_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
            ...
    """

    @wraps(func)
    async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs) -> Any:
        user = update.effective_user
        if not user:
            return

        user_id = user.id
        admin_users = context.bot_data.get("admin_users", set())

        if user_id in admin_users:
            return await func(update, context, *args, **kwargs)

        await update.message.reply_text("⛔ 此命令需要管理员权限")
        return None

    return wrapper