# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. """Reusable gates that wrap a command handler. These are the mechanisms salvaged from the unmerged vote and rate-limit branches, detached from the command they were originally written for. They only decide whether a handler may run; recording a successful use is the dispatcher's job, so stacking two gates cannot double-count one invocation. A gate that refuses raises ``CommandRefusedError`` rather than returning, which is what tells the dispatcher no use took place. Returning quietly used to look identical to a successful run. """ from __future__ import annotations import functools import logging import math from collections.abc import Callable from typing import TYPE_CHECKING from signal_bot.registry import CommandRefusedError, Handler from signal_bot.stats.text import plural if TYPE_CHECKING: from signal_bot.context import CommandContext logger = logging.getLogger(__name__) Decorator = Callable[[Handler], Handler] # Padding on the *displayed* wait only, never on the gate itself. A number that # is a second short sends somebody back for another refusal, which reads as the # bot lying; a number that is a second long costs nothing. DISPLAY_MARGIN_SECONDS = 2 def wait_display(remaining: float) -> tuple[int, str]: """Round a remaining wait up (plus margin) and pick its Polish noun form. Three forms, not two: 1 sekundę, 2-4 sekundy, 5+ sekund -- and the rule reads the last two digits, so 12-14 take the many-form and 22-24 do not. """ seconds = math.ceil(remaining) + DISPLAY_MARGIN_SECONDS return seconds, plural(seconds, "SEKUNDĘ", "SEKUNDY", "SEKUND") def cooldown(seconds: float, message: str) -> Decorator: """Refuse a command until `seconds` have passed since that user's last use. The original bot shared one cooldown window across every user and every command, so one person's cat picture silenced everybody. """ def decorate(handler: Handler) -> Handler: @functools.wraps(handler) async def wrapper(ctx: CommandContext) -> None: last = ctx.store.last_used(ctx.group_id, ctx.message.source_uuid, ctx.command_name) if last is not None and ctx.now - last < seconds: shown, unit = wait_display(seconds - (ctx.now - last)) # Logged with the number the group was given: when this said # one thing and the gate did another, nothing recorded either. logger.info( "cooldown refused %s for %s: told them %d %s", ctx.command_name, ctx.message.source_name, shown, unit, ) await ctx.reply(message.format(seconds=shown, unit=unit)) raise CommandRefusedError await handler(ctx) return wrapper return decorate def daily_limit(max_uses: int, message: str) -> Decorator: """Allow a command at most `max_uses` times per user per calendar day.""" def decorate(handler: Handler) -> Handler: @functools.wraps(handler) async def wrapper(ctx: CommandContext) -> None: used = ctx.store.uses_today( ctx.group_id, ctx.message.source_uuid, ctx.command_name, ctx.day ) if used >= max_uses: await ctx.reply(message.format(max_uses=max_uses, used=used)) raise CommandRefusedError await handler(ctx) return wrapper return decorate def vote_gated(required: int, window_seconds: float, message: str) -> Decorator: """Run a command only once `required` distinct people ask within a window. A repeat vote from the same person is ignored, and the tally is cleared once the command fires so the next run starts from zero. """ def decorate(handler: Handler) -> Handler: @functools.wraps(handler) async def wrapper(ctx: CommandContext) -> None: ctx.store.add_vote(ctx.group_id, ctx.command_name, ctx.message.source_uuid, ctx.now) votes = ctx.store.count_votes(ctx.group_id, ctx.command_name, ctx.now - window_seconds) if votes < required: await ctx.reply(message.format(votes=votes, required=required)) raise CommandRefusedError ctx.store.clear_votes(ctx.group_id, ctx.command_name) await handler(ctx) return wrapper return decorate