perefouras/database.py

# database.py
import sqlite3
import random
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, Tuple

DB_FILE = "data/database.db"

CREATE_GUILD_STATE_QUERY = """
INSERT INTO guild_state (guild_id, cooldown_until, cooldown_ratio, self_control, last_updated)
VALUES (?, datetime('now'), 1.0, 1.0, datetime('now'))
"""

GET_GUILD_STATE_QUERY = """
SELECT cooldown_until, cooldown_ratio, self_control, last_updated
FROM guild_state
WHERE guild_id = ?
"""

CONFIGURE_COOLDOWN_RATIO_QUERY = """
UPDATE guild_state
SET cooldown_ratio = ?
WHERE guild_id = ?
"""

DAMAGE_SELF_CONTROL_QUERY = """
UPDATE guild_state
SET self_control = self_control * 0.9
WHERE guild_id = ?
"""

RESET_COOLDOWN_QUERY = """
UPDATE guild_state
SET cooldown_until = ?,
    self_control = self_control + 1.0,
    last_updated = datetime('now')
    WHERE guild_id = ?
"""

SAVE_RIDDLE_QUERY = """
INSERT OR REPLACE INTO active_riddles
(channel_id, riddle_index, nb_clues, message_id, solver_id, updated_at)
VALUES (?, ?, ?, ?, ?, datetime('now'))
"""

INCREMENT_CLUE_QUERY = """
UPDATE active_riddles
SET nb_clues = nb_clues + 1
    WHERE channel_id = ?
"""

DELETE_RIDDLE_QUERY = "DELETE FROM active_riddles WHERE channel_id = ?"

def ensure_db() -> None:
    """Initialize SQLite database and create tables if needed."""
    Path(DB_FILE).parent.mkdir(parents=True, exist_ok=True)

    with sqlite3.connect(DB_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS guild_state (
                guild_id TEXT PRIMARY KEY,
                cooldown_until TEXT NOT NULL DEFAULT '1970-01-01T00:00:00',
                cooldown_ratio REAL NOT NULL DEFAULT 1.0,
                self_control REAL NOT NULL DEFAULT 1.0,
                last_updated TEXT NOT NULL DEFAULT (datetime('now'))
            )
        """)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS active_riddles (
                channel_id TEXT PRIMARY KEY,
                riddle_index INTEGER NOT NULL,
                nb_clues INTEGER NOT NULL DEFAULT -1,
                message_id INTEGER NOT NULL,
                solver_id TEXT,
                created_at TEXT NOT NULL DEFAULT (datetime('now')),
                updated_at TEXT NOT NULL DEFAULT (datetime('now'))
            )
        """)
        conn.commit()


def get_connection() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_FILE)
    conn.row_factory = sqlite3.Row
    return conn


def get_guild_state(guild_id: str) -> Dict[str, Any]:
    with get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(GET_GUILD_STATE_QUERY,(guild_id,))
        row = cursor.fetchone()

        if not row:
            cursor.execute(CREATE_GUILD_STATE_QUERY, (guild_id,))
            conn.commit()

            cursor.execute(GET_GUILD_STATE_QUERY,(guild_id,))
            row = cursor.fetchone()

        if row:
            return {
                "cooldown_until": row["cooldown_until"],
                "cooldown_ratio": row["cooldown_ratio"],
                "self_control": row["self_control"],
                "last_updated": row["last_updated"],
            }
    return {}


def rhyme_configure_cooldown_ratio(guild_id: str, cooldown_ratio: float) -> bool:
    with get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(CONFIGURE_COOLDOWN_RATIO_QUERY, (cooldown_ratio, guild_id,))
        conn.commit()


def rhyme_damage_self_control(guild_id: str) -> None:
    with get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(DAMAGE_SELF_CONTROL_QUERY, (guild_id,))
        conn.commit()


def rhyme_reset_cooldown(
    guild_id: str,
    ratio: float
) -> None:
    with get_connection() as conn:
        # Calculate new cooldown duration
        wait_time = random.randint(900, 10800) if bool(random.getrandbits(1)) else random.randint(0, 900)
        wait_time = 0 if ratio < 0 else int(wait_time * ratio)
        cooldown_until = datetime.now().replace(second=0, microsecond=0) + timedelta(seconds=wait_time)

        cursor = conn.cursor()
        cursor.execute(RESET_COOLDOWN_QUERY, (cooldown_until.isoformat(), guild_id,))
        conn.commit()


def get_active_riddle(channel_id: str) -> Any:
    with get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM active_riddles WHERE channel_id = ?", (str(channel_id),))
        row = cursor.fetchone()

        if row:
            return {
                "channel_id": row["channel_id"],
                "riddle_index": row["riddle_index"],
                "nb_clues": row["nb_clues"],
                "message_id": row["message_id"],
                "solver_id": row["solver_id"],
                "created_at": row["created_at"],
                "updated_at": row["updated_at"],
            }
        return None


def save_active_riddle(
    channel_id: str,
    riddle_index: int,
    nb_clues: int,
    message_id: int,
    solver_id: str = "",
) -> None:
    with get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(SAVE_RIDDLE_QUERY, (channel_id, riddle_index, nb_clues, message_id, solver_id))
        conn.commit()


def increment_clue(channel_id: str) -> None:
    with get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(INCREMENT_CLUE_QUERY, (channel_id,),)
        conn.commit()


def remove_active_riddle(channel_id: str) -> None:
    with get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(DELETE_RIDDLE_QUERY, (str(channel_id),))
        conn.commit()