signal-bot/signal_bot/store.py

# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""SQLite persistence for tallies, cooldowns, daily limits and votes.

Everything the bot remembers lives here so a restart before the daily report
does not lose the day. The original bot kept all of it in module globals.
"""

from __future__ import annotations

import contextlib
import sqlite3
from dataclasses import dataclass
from typing import TYPE_CHECKING, Self

from signal_bot.schema import SCHEMA, SCHEMA_VERSION, TABLES

if TYPE_CHECKING:
    from collections.abc import Iterator
    from types import TracebackType


@dataclass(frozen=True, slots=True)
class Tally:
    """One person's message count for the current reporting period."""

    uuid: str
    display_name: str
    count: int


class Store:
    """Thin, synchronous SQLite wrapper. All state the bot keeps lives here."""

    def __init__(self, path: str) -> None:
        """Open (or create) the database at path and apply the schema."""
        self._conn = sqlite3.connect(path)
        self._conn.row_factory = sqlite3.Row
        self._apply_schema()

    def _apply_schema(self) -> None:
        """Create the tables, rebuilding them when the file predates them.

        ``CREATE TABLE IF NOT EXISTS`` is a no-op against a table that
        already exists with the wrong columns, so a database written before
        the group_id key would survive startup untouched and then fail on
        every single INSERT -- a bot that starts clean and answers nothing.
        PRAGMA user_version is the only thing that notices, because it is
        the only part of the file that records which shape it holds.

        The old rows are dropped rather than migrated: they carry no group,
        and inventing one would silently attribute a whole day's counts to
        the wrong conversation.
        """
        version = int(self._conn.execute("PRAGMA user_version").fetchone()[0])
        if version < SCHEMA_VERSION:
            for table in TABLES:
                self._conn.execute(f"DROP TABLE IF EXISTS {table}")
        self._conn.executescript(SCHEMA)
        self._conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
        self._conn.commit()

    def __enter__(self) -> Self:
        """Enter a context manager that closes the connection on exit."""
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        """Close the underlying connection."""
        self.close()

    def close(self) -> None:
        """Close the underlying connection."""
        self._conn.close()

    def forget_group(self, group_id: str) -> None:
        """Drop every row belonging to one group.

        Called when the bot is removed from a group, so that being added
        back starts a fresh tally and a clean cooldown rather than
        resurrecting whatever was true whenever it was last in there.
        """
        # Spelled out rather than looped over _TABLES: an f-string here is a
        # SQL-injection shape even when the input is a module constant, and
        # the repo does not take lint suppressions. test_forget_group_clears
        # _every_table is what keeps this list in step with the schema.
        self._conn.execute("DELETE FROM message_counts WHERE group_id = ?", (group_id,))
        self._conn.execute("DELETE FROM command_uses WHERE group_id = ?", (group_id,))
        self._conn.execute("DELETE FROM daily_uses WHERE group_id = ?", (group_id,))
        self._conn.execute("DELETE FROM votes WHERE group_id = ?", (group_id,))
        self._conn.execute("DELETE FROM group_meta WHERE group_id = ?", (group_id,))
        self._conn.commit()

    # -- per-group key/value state -----------------------------------------

    def meta(self, group_id: str, key: str) -> str | None:
        """Return a stored value for one group, or None if never set."""
        row = self._conn.execute(
            "SELECT value FROM group_meta WHERE group_id = ? AND key = ?",
            (group_id, key),
        ).fetchone()
        return None if row is None else str(row["value"])

    def set_meta(self, group_id: str, key: str, value: str) -> None:
        """Store a value for one group, replacing any previous one."""
        self._conn.execute(
            "INSERT INTO group_meta (group_id, key, value) VALUES (?, ?, ?) "
            "ON CONFLICT(group_id, key) DO UPDATE SET value = excluded.value",
            (group_id, key, value),
        )
        self._conn.commit()

    # -- message tallies ---------------------------------------------------

    def increment(self, group_id: str, uuid: str, display_name: str) -> None:
        """Add one to a person's tally in one group, creating the row when new."""
        self._conn.execute(
            """
            INSERT INTO message_counts (group_id, uuid, display_name, count)
            VALUES (?, ?, ?, 1)
            ON CONFLICT(group_id, uuid) DO UPDATE SET
                count = count + 1,
                display_name = excluded.display_name
            """,
            (group_id, uuid, display_name),
        )
        self._conn.commit()

    def tallies(self, group_id: str) -> list[Tally]:
        """Return one group's tallies, highest count first."""
        rows = self._conn.execute(
            "SELECT uuid, display_name, count FROM message_counts "
            "WHERE group_id = ? ORDER BY count DESC, display_name ASC",
            (group_id,),
        ).fetchall()
        return [Tally(r["uuid"], r["display_name"], r["count"]) for r in rows]

    def reset_tallies(self, group_id: str) -> None:
        """Clear one group's tallies, as its daily report does after publishing."""
        self._conn.execute("DELETE FROM message_counts WHERE group_id = ?", (group_id,))
        self._conn.commit()

    # -- per-group, per-user, per-command rate limiting --------------------

    def record_use(self, group_id: str, user_uuid: str, command: str, now: float, day: str) -> None:
        """Record one successful use for both cooldown and daily-limit checks."""
        self._conn.execute(
            "INSERT INTO command_uses (group_id, user_uuid, command, used_at) "
            "VALUES (?, ?, ?, ?) "
            "ON CONFLICT(group_id, user_uuid, command) DO UPDATE SET used_at = excluded.used_at",
            (group_id, user_uuid, command, now),
        )
        self._conn.execute(
            """
            INSERT INTO daily_uses (group_id, user_uuid, command, day, uses)
            VALUES (?, ?, ?, ?, 1)
            ON CONFLICT(group_id, user_uuid, command, day) DO UPDATE SET uses = uses + 1
            """,
            (group_id, user_uuid, command, day),
        )
        self._conn.commit()

    def last_used(self, group_id: str, user_uuid: str, command: str) -> float | None:
        """Return when this user last ran this command in this group, or None."""
        row = self._conn.execute(
            "SELECT used_at FROM command_uses WHERE group_id = ? AND user_uuid = ? AND command = ?",
            (group_id, user_uuid, command),
        ).fetchone()
        return None if row is None else float(row["used_at"])

    def uses_today(self, group_id: str, user_uuid: str, command: str, day: str) -> int:
        """Return how often this user ran this command in this group that day."""
        row = self._conn.execute(
            "SELECT uses FROM daily_uses "
            "WHERE group_id = ? AND user_uuid = ? AND command = ? AND day = ?",
            (group_id, user_uuid, command, day),
        ).fetchone()
        return 0 if row is None else int(row["uses"])

    # -- vote gating -------------------------------------------------------

    def add_vote(self, group_id: str, command: str, voter_uuid: str, now: float) -> None:
        """Record one vote, ignoring a repeat vote from the same person."""
        self._conn.execute(
            "INSERT INTO votes (group_id, command, voter_uuid, voted_at) VALUES (?, ?, ?, ?) "
            "ON CONFLICT(group_id, command, voter_uuid) DO NOTHING",
            (group_id, command, voter_uuid, now),
        )
        self._conn.commit()

    def count_votes(self, group_id: str, command: str, not_before: float) -> int:
        """Return how many votes for command in this group are still in the window."""
        row = self._conn.execute(
            "SELECT COUNT(*) AS n FROM votes WHERE group_id = ? AND command = ? AND voted_at >= ?",
            (group_id, command, not_before),
        ).fetchone()
        return int(row["n"])

    def clear_votes(self, group_id: str, command: str) -> None:
        """Drop every vote for a command in one group, after it fires or expires."""
        self._conn.execute(
            "DELETE FROM votes WHERE group_id = ? AND command = ?", (group_id, command)
        )
        self._conn.commit()


@contextlib.contextmanager
def open_store(path: str) -> Iterator[Store]:
    """Open a Store as a context manager."""
    store = Store(path)
    try:
        yield store
    finally:
        store.close()