signal-bot/signal_bot/decorators.py

# 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