shirobana-bot/bot-utility/agent/modes.py

"""SQLite-backed command/agent mode storage."""

from __future__ import annotations

import os
import json
import sqlite3
import threading
import time
from pathlib import Path
from typing import Optional, Union


DEFAULT_DB_PATH = Path(__file__).resolve().parents[1] / "data" / "agent.db"
_VALID_MODES = frozenset(("command", "agent", "chat"))


class ModeStore:
    """Persist conversation mode with private and group scopes kept separate.

    A private conversation is keyed by ``(channel, uid)``. A group conversation
    is keyed by ``(channel, group_id)`` and therefore shared by its members.
    Every operation uses a short-lived SQLite connection, making one instance
    safe to share between the event loop and worker threads.
    """

    def __init__(self, db_path: Optional[Union[str, os.PathLike]] = None) -> None:
        self.db_path = Path(db_path) if db_path is not None else DEFAULT_DB_PATH
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        with self._connect() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS agent_modes (
                    channel TEXT NOT NULL,
                    scope TEXT NOT NULL,
                    subject_id TEXT NOT NULL,
                    mode TEXT NOT NULL CHECK(mode IN ('command', 'agent', 'chat')),
                    updated_at REAL NOT NULL,
                    PRIMARY KEY (channel, scope, subject_id)
                )
                """
            )

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(str(self.db_path), timeout=10)
        conn.execute("PRAGMA busy_timeout = 10000")
        return conn

    @staticmethod
    def _key(channel: str, uid: int, group_id: Optional[int]) -> tuple:
        if not isinstance(channel, str) or not channel.strip():
            raise ValueError("channel must be a non-empty string")
        if group_id is not None:
            return channel.strip(), "group", str(group_id)
        return channel.strip(), "private", str(uid)

    def _configured_default(self, channel: str, uid: int, group_id: Optional[int]) -> str:
        """Return the settings.json default without overriding saved preferences.

        Example: {"chat_mode": {"default": "command", "admin": "chat",
        "group": "command"}}.  Administrator chat is intended for private
        conversations; groups stay in command unless explicitly configured.
        """
        if group_id is not None:
            scope = "group"
        else:
            try:
                import identity
                scope = "admin" if identity.is_admin(uid) else "default"
            except Exception:
                scope = "default"
        try:
            settings_path = Path(__file__).resolve().parents[1] / "settings.json"
            with settings_path.open("r", encoding="utf-8") as fh:
                config = json.load(fh).get("chat_mode", {})
            mode = config.get(scope, config.get("default", "command"))
            return mode if mode in ("chat", "agent") else "command"
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            return "command"

    def get_mode(self, channel: str, uid: int, group_id: Optional[int] = None) -> str:
        """Return saved mode, then the config default; ultimate fallback is command."""
        key = self._key(channel, uid, group_id)
        with self._lock, self._connect() as conn:
            row = conn.execute(
                "SELECT mode FROM agent_modes WHERE channel=? AND scope=? AND subject_id=?",
                key,
            ).fetchone()
        if row and row[0] in _VALID_MODES:
            return row[0]
        return self._configured_default(channel, uid, group_id)

    def set_mode(
        self,
        channel: str,
        uid: int,
        mode: str,
        group_id: Optional[int] = None,
    ) -> str:
        """Set and return the mode for a private user or an entire group."""
        if mode not in _VALID_MODES:
            raise ValueError("mode must be 'command' or 'chat'")
        key = self._key(channel, uid, group_id)
        with self._lock, self._connect() as conn:
            conn.execute(
                """
                INSERT INTO agent_modes(channel, scope, subject_id, mode, updated_at)
                VALUES (?, ?, ?, ?, ?)
                ON CONFLICT(channel, scope, subject_id) DO UPDATE SET
                    mode=excluded.mode, updated_at=excluded.updated_at
                """,
                (*key, mode, time.time()),
            )
        return mode


_default_store: Optional[ModeStore] = None
_default_lock = threading.Lock()


def _default() -> ModeStore:
    global _default_store
    if _default_store is None:
        with _default_lock:
            if _default_store is None:
                _default_store = ModeStore()
    return _default_store


def get_mode(channel: str, uid: int, group_id: Optional[int] = None) -> str:
    """Module-level convenience wrapper using ``bot-utility/data/agent.db``."""
    return _default().get_mode(channel, uid, group_id)


def set_mode(
    channel: str,
    uid: int,
    mode: str,
    group_id: Optional[int] = None,
) -> str:
    """Module-level convenience wrapper using ``bot-utility/data/agent.db``."""
    return _default().set_mode(channel, uid, mode, group_id)