db.py

from __future__ import annotations

import os
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Optional, Sequence

try:
    import aiomysql
except Exception:
    aiomysql = None


@dataclass(frozen=True)
class MariaDBConfig:
    host: str
    port: int
    user: str
    password: str
    database: str
    min_pool: int = 1
    max_pool: int = 10


def _env_truthy(value: Optional[str]) -> bool:
    return str(value or "").strip().lower() in {"1", "true", "yes", "y", "on"}


def load_mariadb_config() -> Optional[MariaDBConfig]:
    host = (os.getenv("MARIADB_HOST") or "").strip()
    user = (os.getenv("MARIADB_USER") or "").strip()
    password = os.getenv("MARIADB_PASSWORD") or ""
    database = (os.getenv("MARIADB_DATABASE") or "").strip()

    if not (host and user and database):
        return None

    try:
        port = int(os.getenv("MARIADB_PORT", "3306"))
    except ValueError:
        port = 3306

    try:
        min_pool = int(os.getenv("MARIADB_MIN_POOL", "1"))
    except ValueError:
        min_pool = 1

    try:
        max_pool = int(os.getenv("MARIADB_MAX_POOL", "10"))
    except ValueError:
        max_pool = 10

    if min_pool < 1:
        min_pool = 1
    if max_pool < min_pool:
        max_pool = min_pool

    return MariaDBConfig(
        host=host,
        port=port,
        user=user,
        password=password,
        database=database,
        min_pool=min_pool,
        max_pool=max_pool,
    )


def mariadb_enabled() -> bool:
    return load_mariadb_config() is not None


async def create_mariadb_pool(cfg: MariaDBConfig):
    if aiomysql is None:
        raise RuntimeError("aiomysql is not installed")

    return await aiomysql.create_pool(
        host=cfg.host,
        port=cfg.port,
        user=cfg.user,
        password=cfg.password,
        db=cfg.database,
        minsize=cfg.min_pool,
        maxsize=cfg.max_pool,
        autocommit=True,
    )


async def ensure_schema(pool) -> None:
    """Create tables needed by JartonBot when MARIADB_AUTO_MIGRATE is enabled."""
    if not _env_truthy(os.getenv("MARIADB_AUTO_MIGRATE")):
        return

    tables: dict[str, str] = {
        "badword_offenses": (
            "CREATE TABLE IF NOT EXISTS badword_offenses ("
            "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
            "user_id BIGINT UNSIGNED NOT NULL,"
            "guild_id BIGINT UNSIGNED NOT NULL,"
            "channel_id BIGINT UNSIGNED NOT NULL,"
            "message_id BIGINT UNSIGNED NULL,"
            "content TEXT NOT NULL,"
            "detected_words TEXT NOT NULL,"
            "created_at DATETIME(6) NOT NULL,"
            "PRIMARY KEY (id),"
            "INDEX idx_badword_user_created (user_id, created_at),"
            "INDEX idx_badword_created (created_at)"
            ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
        ),
        "tickets": (
            "CREATE TABLE IF NOT EXISTS tickets ("
            "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
            "guild_id BIGINT UNSIGNED NOT NULL,"
            "channel_id BIGINT UNSIGNED NOT NULL,"
            "category VARCHAR(64) NOT NULL,"
            "opener_id BIGINT UNSIGNED NOT NULL,"
            "status VARCHAR(16) NOT NULL,"
            "created_at DATETIME(6) NOT NULL,"
            "closed_at DATETIME(6) NULL,"
            "closed_by BIGINT UNSIGNED NULL,"
            "close_reason VARCHAR(255) NULL,"
            "transcript_path VARCHAR(512) NULL,"
            "PRIMARY KEY (id),"
            "UNIQUE KEY uq_tickets_channel (channel_id),"
            "INDEX idx_tickets_guild_status (guild_id, status),"
            "INDEX idx_tickets_opener (opener_id, created_at)"
            ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
        ),
        "applications": (
            "CREATE TABLE IF NOT EXISTS applications ("
            "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
            "guild_id BIGINT UNSIGNED NULL,"
            "user_id BIGINT UNSIGNED NOT NULL,"
            "role_applied VARCHAR(64) NOT NULL,"
            "answers_json LONGTEXT NOT NULL,"
            "submitted_at DATETIME(6) NOT NULL,"
            "thread_id BIGINT UNSIGNED NULL,"
            "transcript_channel_id BIGINT UNSIGNED NULL,"
            "PRIMARY KEY (id),"
            "INDEX idx_apps_user (user_id, submitted_at),"
            "INDEX idx_apps_guild (guild_id, submitted_at)"
            ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
        ),
        "mod_warns": (
            "CREATE TABLE IF NOT EXISTS mod_warns ("
            "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
            "guild_id BIGINT UNSIGNED NOT NULL,"
            "user_id BIGINT UNSIGNED NOT NULL,"
            "moderator_id BIGINT UNSIGNED NOT NULL,"
            "reason VARCHAR(512) NOT NULL,"
            "created_at DATETIME(6) NOT NULL,"
            "PRIMARY KEY (id),"
            "INDEX idx_modwarn_user (guild_id, user_id, created_at)"
            ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
        ),
        "count_state": (
            "CREATE TABLE IF NOT EXISTS count_state ("
            "guild_id BIGINT UNSIGNED NOT NULL,"
            "state_json LONGTEXT NOT NULL,"
            "updated_at DATETIME(6) NOT NULL,"
            "PRIMARY KEY (guild_id)"
            ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
        ),
        "invite_joins": (
            "CREATE TABLE IF NOT EXISTS invite_joins ("
            "guild_id BIGINT UNSIGNED NOT NULL,"
            "member_id BIGINT UNSIGNED NOT NULL,"
            "inviter_id BIGINT UNSIGNED NULL,"
            "joined_at DATETIME(6) NOT NULL,"
            "PRIMARY KEY (guild_id, member_id),"
            "INDEX idx_invite_inviter (guild_id, inviter_id),"
            "INDEX idx_invite_joined (guild_id, joined_at)"
            ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
        ),
        "suggestions": (
            "CREATE TABLE IF NOT EXISTS suggestions ("
            "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
            "guild_id BIGINT UNSIGNED NOT NULL,"
            "channel_id BIGINT UNSIGNED NOT NULL,"
            "message_id BIGINT UNSIGNED NOT NULL,"
            "author_id BIGINT UNSIGNED NOT NULL,"
            "title VARCHAR(120) NOT NULL,"
            "description TEXT NOT NULL,"
            "status VARCHAR(16) NOT NULL,"
            "denial_reason VARCHAR(500) NULL,"
            "reviewed_by BIGINT UNSIGNED NULL,"
            "reviewed_at DATETIME(6) NULL,"
            "created_at DATETIME(6) NOT NULL,"
            "PRIMARY KEY (id),"
            "UNIQUE KEY uq_suggestions_message (message_id),"
            "INDEX idx_suggestions_guild_status (guild_id, status, created_at),"
            "INDEX idx_suggestions_author (author_id, created_at)"
            ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
        ),
    }

    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("SELECT DATABASE()")
            db_row = await cur.fetchone()
            db_name = db_row[0] if db_row else None

            created = 0
            existed = 0

            for table_name, ddl in tables.items():
                exists = False
                if db_name:
                    await cur.execute(
                        "SELECT 1 FROM information_schema.tables WHERE table_schema=%s AND table_name=%s LIMIT 1",
                        (db_name, table_name),
                    )
                    exists = (await cur.fetchone()) is not None

                if exists:
                    existed += 1
                    print(f"ℹ️ MariaDB schema: {table_name} already exists; skipping.")
                    continue

                await cur.execute(ddl)
                created += 1
                print(f"✅ MariaDB schema: created {table_name}.")

            print(f"✅ MariaDB schema ready (created {created}, existing {existed}).")


def utcnow() -> datetime:
    return datetime.now(timezone.utc)


async def record_suggestion_submit(
    pool,
    *,
    guild_id: int,
    channel_id: int,
    message_id: int,
    author_id: int,
    title: str,
    description: str,
    created_at: datetime,
) -> None:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "INSERT INTO suggestions (guild_id, channel_id, message_id, author_id, title, description, status, created_at) "
                "VALUES (%s,%s,%s,%s,%s,%s,%s,%s) "
                "ON DUPLICATE KEY UPDATE "
                "guild_id=VALUES(guild_id), channel_id=VALUES(channel_id), author_id=VALUES(author_id), "
                "title=VALUES(title), description=VALUES(description), status=VALUES(status), "
                "denial_reason=NULL, reviewed_by=NULL, reviewed_at=NULL",
                (
                    guild_id,
                    channel_id,
                    message_id,
                    author_id,
                    (title or "")[:120],
                    description or "",
                    "pending",
                    created_at.replace(tzinfo=None),
                ),
            )


async def get_suggestion_by_message_id(pool, *, message_id: int) -> Optional[dict[str, Any]]:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT id, guild_id, channel_id, message_id, author_id, title, description, status, denial_reason, reviewed_by, reviewed_at, created_at "
                "FROM suggestions WHERE message_id=%s LIMIT 1",
                (message_id,),
            )
            row = await cur.fetchone()
            if not row:
                return None

            return {
                "id": int(row[0]),
                "guild_id": int(row[1]),
                "channel_id": int(row[2]),
                "message_id": int(row[3]),
                "author_id": int(row[4]),
                "title": row[5],
                "description": row[6],
                "status": row[7],
                "denial_reason": row[8],
                "reviewed_by": int(row[9]) if row[9] is not None else None,
                "reviewed_at": row[10],
                "created_at": row[11],
            }


async def record_suggestion_review(
    pool,
    *,
    message_id: int,
    status: str,
    reviewed_by: int,
    reviewed_at: datetime,
    denial_reason: Optional[str] = None,
) -> bool:
    normalized_status = "approved" if status == "approved" else "denied"
    normalized_reason = None if normalized_status == "approved" else (denial_reason or "")[:500]

    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "UPDATE suggestions SET status=%s, denial_reason=%s, reviewed_by=%s, reviewed_at=%s "
                "WHERE message_id=%s AND status=%s",
                (
                    normalized_status,
                    normalized_reason,
                    reviewed_by,
                    reviewed_at.replace(tzinfo=None),
                    message_id,
                    "pending",
                ),
            )
            return bool(cur.rowcount)


async def record_badword_offense(
    pool,
    *,
    user_id: int,
    guild_id: int,
    channel_id: int,
    message_id: Optional[int],
    content: str,
    detected_words: Sequence[str],
    created_at: datetime,
) -> None:
    words = ",".join(sorted({w.lower() for w in detected_words if w}))

    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "INSERT INTO badword_offenses (user_id, guild_id, channel_id, message_id, content, detected_words, created_at) "
                "VALUES (%s,%s,%s,%s,%s,%s,%s)",
                (user_id, guild_id, channel_id, message_id, content, words, created_at.replace(tzinfo=None)),
            )


async def record_ticket_open(
    pool,
    *,
    guild_id: int,
    channel_id: int,
    category: str,
    opener_id: int,
    created_at: datetime,
) -> None:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "INSERT INTO tickets (guild_id, channel_id, category, opener_id, status, created_at) "
                "VALUES (%s,%s,%s,%s,%s,%s) "
                "ON DUPLICATE KEY UPDATE category=VALUES(category), opener_id=VALUES(opener_id), status=VALUES(status)",
                (guild_id, channel_id, category[:64], opener_id, "open", created_at.replace(tzinfo=None)),
            )


async def record_ticket_close(
    pool,
    *,
    channel_id: int,
    closed_by: Optional[int],
    close_reason: str,
    transcript_path: Optional[str],
    closed_at: datetime,
) -> None:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "UPDATE tickets SET status=%s, closed_at=%s, closed_by=%s, close_reason=%s, transcript_path=%s "
                "WHERE channel_id=%s",
                (
                    "closed",
                    closed_at.replace(tzinfo=None),
                    closed_by,
                    (close_reason or "")[:255],
                    transcript_path,
                    channel_id,
                ),
            )


async def record_application_submit(
    pool,
    *,
    guild_id: Optional[int],
    user_id: int,
    role_applied: str,
    answers: Any,
    submitted_at: datetime,
    thread_id: Optional[int],
    transcript_channel_id: Optional[int],
) -> None:
    import json

    answers_json = json.dumps(answers, ensure_ascii=False)
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "INSERT INTO applications (guild_id, user_id, role_applied, answers_json, submitted_at, thread_id, transcript_channel_id) "
                "VALUES (%s,%s,%s,%s,%s,%s,%s)",
                (
                    guild_id,
                    user_id,
                    role_applied[:64],
                    answers_json,
                    submitted_at.replace(tzinfo=None),
                    thread_id,
                    transcript_channel_id,
                ),
            )


async def record_mod_warn(
    pool,
    *,
    guild_id: int,
    user_id: int,
    moderator_id: int,
    reason: str,
    created_at: datetime,
) -> None:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "INSERT INTO mod_warns (guild_id, user_id, moderator_id, reason, created_at) VALUES (%s,%s,%s,%s,%s)",
                (guild_id, user_id, moderator_id, (reason or "No reason provided")[:512], created_at.replace(tzinfo=None)),
            )


async def upsert_count_state(pool, *, guild_id: int, state_json: str, updated_at: datetime) -> None:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "INSERT INTO count_state (guild_id, state_json, updated_at) VALUES (%s,%s,%s) "
                "ON DUPLICATE KEY UPDATE state_json=VALUES(state_json), updated_at=VALUES(updated_at)",
                (guild_id, state_json, updated_at.replace(tzinfo=None)),
            )


async def get_count_state(pool, *, guild_id: int) -> Optional[str]:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("SELECT state_json FROM count_state WHERE guild_id=%s", (guild_id,))
            row = await cur.fetchone()
            return row[0] if row else None


async def record_invite_join(
    pool,
    *,
    guild_id: int,
    member_id: int,
    inviter_id: Optional[int],
    joined_at: datetime,
) -> None:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "INSERT INTO invite_joins (guild_id, member_id, inviter_id, joined_at) VALUES (%s,%s,%s,%s) "
                "ON DUPLICATE KEY UPDATE inviter_id=VALUES(inviter_id), joined_at=VALUES(joined_at)",
                (guild_id, member_id, inviter_id, joined_at.replace(tzinfo=None)),
            )


async def get_invite_count(pool, *, guild_id: int, inviter_id: int) -> int:
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT COUNT(*) FROM invite_joins WHERE guild_id=%s AND inviter_id=%s",
                (guild_id, inviter_id),
            )
            row = await cur.fetchone()
            return int(row[0] if row else 0)


async def get_top_invites(pool, *, guild_id: int, limit: int = 10) -> list[tuple[int, int]]:
    lim = max(1, min(int(limit), 25))
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT inviter_id, COUNT(*) AS c "
                "FROM invite_joins WHERE guild_id=%s AND inviter_id IS NOT NULL "
                "GROUP BY inviter_id ORDER BY c DESC LIMIT %s",
                (guild_id, lim),
            )
            rows = await cur.fetchall() or []
            return [(int(inviter_id), int(count)) for (inviter_id, count) in rows if inviter_id is not None]


async def badword_count_last_7d(pool, *, user_id: int, now: datetime) -> int:
    window_start = now - __import__("datetime").timedelta(days=7)
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT COUNT(*) FROM badword_offenses WHERE user_id=%s AND created_at >= %s",
                (user_id, window_start.replace(tzinfo=None)),
            )
            row = await cur.fetchone()
            return int(row[0] if row else 0)


async def purge_badwords_older_than_7d(pool, *, now: datetime) -> int:
    cutoff = now - __import__("datetime").timedelta(days=7)
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "DELETE FROM badword_offenses WHERE created_at < %s",
                (cutoff.replace(tzinfo=None),),
            )
            return int(cur.rowcount or 0)