from functools import wraps
from telegram import Update
from telegram.ext import ContextTypes
from config import ADMIN_IDS
from database.db import DB
import time, collections

# Simple anti-flood: max 20 messages per 10s per user
_buckets: dict[int, collections.deque] = {}
FLOOD_LIMIT, FLOOD_WINDOW = 20, 10

def admin_only(func):
    @wraps(func)
    async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *a, **kw):
        uid = update.effective_user.id if update.effective_user else 0
        if uid not in ADMIN_IDS:
            if update.callback_query:
                await update.callback_query.answer("Admin only.", show_alert=True)
            elif update.message:
                await update.message.reply_text("⛔ Admin only.")
            return
        return await func(update, context, *a, **kw)
    return wrapper

def anti_flood(func):
    @wraps(func)
    async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *a, **kw):
        uid = update.effective_user.id if update.effective_user else 0
        now = time.time()
        bucket = _buckets.setdefault(uid, collections.deque())
        while bucket and now - bucket[0] > FLOOD_WINDOW:
            bucket.popleft()
        if len(bucket) >= FLOOD_LIMIT:
            return
        bucket.append(now)
        # Ban check
        row = await DB.fetchone("SELECT is_banned FROM users WHERE user_id=?", (uid,))
        if row and row["is_banned"]:
            if update.message:
                await update.message.reply_text("⛔ Your account is banned.")
            return
        return await func(update, context, *a, **kw)
    return wrapper