bot/database.py

"""
Database layer for the bot using aiosqlite.
Provides async SQLite operations with schema migrations and persistent connection.
"""

import asyncio
import os
import shutil
import time
from pathlib import Path

import aiosqlite

try:
    from huggingface_hub import HfApi, hf_hub_download
except Exception:  # pragma: no cover
    HfApi = None
    hf_hub_download = None


# Minimum seconds between HuggingFace DB sync pushes to avoid rate limits
_HF_SYNC_INTERVAL_SECONDS = 60


class Database:
    def __init__(self, path: str = "database.db") -> None:
        self.path = path
        self._db: aiosqlite.Connection | None = None
        self._lock = asyncio.Lock()
        self._hf_repo_id = (os.getenv("HF_DB_REPO_ID") or "").strip()
        self._hf_token = (
            os.getenv("HF_TOKEN")
            or os.getenv("HUGGINGFACE_TOKEN")
            or ""
        ).strip()
        self._hf_filename = (os.getenv("HF_DB_FILE") or "database.db").strip()
        self._hf_repo_type = "dataset"
        self._hf_sync_enabled = bool(
            (os.getenv("HF_DB_SYNC", "0").strip().lower() in {"1", "true", "yes", "on"})
            and self._hf_repo_id
            and self._hf_token
            and HfApi
            and hf_hub_download
        )
        self._last_hf_push: float = 0.0

    def _prepare_path(self) -> None:
        """Ensure the sqlite parent directory exists when a directory is provided."""
        parent = os.path.dirname(self.path)
        if not parent:
            return
        os.makedirs(parent, exist_ok=True)

    async def _pull_remote_db(self) -> None:
        if not self._hf_sync_enabled:
            return
        try:
            local_file = await asyncio.to_thread(
                hf_hub_download,
                repo_id=self._hf_repo_id,
                filename=self._hf_filename,
                repo_type=self._hf_repo_type,
                token=self._hf_token,
            )
            Path(self.path).parent.mkdir(parents=True, exist_ok=True)
            await asyncio.to_thread(shutil.copy2, local_file, self.path)
        except Exception:
            # Disable further sync attempts to avoid repeated noisy auth/network failures.
            self._hf_sync_enabled = False
            return

    async def _push_remote_db(self) -> None:
        if not self._hf_sync_enabled:
            return
        # Rate limit: avoid pushing too frequently
        now = time.monotonic()
        elapsed = now - self._last_hf_push
        if elapsed < _HF_SYNC_INTERVAL_SECONDS:
            return
        try:
            api = HfApi(token=self._hf_token)
            await asyncio.to_thread(
                api.upload_file,
                path_or_fileobj=self.path,
                path_in_repo=self._hf_filename,
                repo_id=self._hf_repo_id,
                repo_type=self._hf_repo_type,
                commit_message="Bot DB sync",
            )
            self._last_hf_push = time.monotonic()
        except Exception:
            # Disable further sync attempts to avoid repeated noisy auth/network failures.
            self._hf_sync_enabled = False
            return

    async def _ensure_column(
        self,
        db: aiosqlite.Connection,
        table: str,
        column: str,
        definition: str,
    ) -> None:
        """Ensure a column exists in a table, adding it if missing."""
        cursor = await db.execute(f"PRAGMA table_info({table})")
        rows = await cursor.fetchall()
        existing = {row[1] for row in rows}
        if column not in existing:
            await db.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}")

    async def _get_connection(self) -> aiosqlite.Connection:
        """Get or create the persistent database connection."""
        if self._db is None:
            self._db = await aiosqlite.connect(self.path)
            await self._db.execute("PRAGMA journal_mode=WAL")
            await self._db.execute("PRAGMA busy_timeout=5000")
        return self._db

    async def close(self) -> None:
        """Close the persistent database connection."""
        if self._db is not None:
            await self._db.close()
            self._db = None

    async def setup(self) -> None:
        """Set up all database tables and run migrations."""
        self._prepare_path()
        await self._pull_remote_db()
        db = await self._get_connection()
        await db.executescript(
            """
            CREATE TABLE IF NOT EXISTS guild_config (
                guild_id INTEGER PRIMARY KEY,
                log_channel_id INTEGER,
                welcome_channel_id INTEGER,
                verify_channel_id INTEGER,
                verify_role_id INTEGER,
                mod_role_id INTEGER,
                suggestion_channel_id INTEGER,
                leveling_enabled INTEGER DEFAULT 1,
                automod_enabled INTEGER DEFAULT 0,
                ticket_category_id INTEGER,
                daily_channel_id INTEGER,
                daily_message TEXT DEFAULT '✅ رسالة اليوم: لا تنسوا نشر الطاقة الإيجابية في السيرفر!',
                daily_enabled INTEGER DEFAULT 0,
                daily_time TEXT DEFAULT '09:00',
                daily_utc_offset INTEGER DEFAULT 0,
                last_daily_sent_date TEXT,
                daily_title TEXT DEFAULT '🌅 Daily Message',
                daily_image_url TEXT,
                daily_button_label TEXT,
                daily_button_url TEXT,
                ai_model TEXT,
                ai_chat_channel_id INTEGER,
                ai_trigger TEXT DEFAULT 'bana',
                ai_auto_enabled INTEGER DEFAULT 0,
                support_ai_enabled INTEGER DEFAULT 0,
                support_channel_id INTEGER,
                guild_language TEXT DEFAULT 'en',
                dj_role_id INTEGER,
                premium_tier INTEGER DEFAULT 0,
                music_autoplay INTEGER DEFAULT 0,
                poll_channel_id INTEGER,
                free_games_channel_id INTEGER,
                free_games_role_id INTEGER,
                free_games_last_ids TEXT,
                free_games_platforms TEXT DEFAULT 'epic,steam,gog',
                free_games_mention_type TEXT DEFAULT 'role',
                wisdom_channel_id INTEGER,
                wisdom_enabled INTEGER DEFAULT 0,
                wisdom_last_sent_date TEXT,
                game_news_channel_id INTEGER,
                game_news_role_id INTEGER,
                game_news_last_ids TEXT,
                custom_banner_url TEXT
            );

            CREATE TABLE IF NOT EXISTS user_xp (
                guild_id INTEGER,
                user_id INTEGER,
                xp INTEGER DEFAULT 0,
                level INTEGER DEFAULT 1,
                PRIMARY KEY (guild_id, user_id)
            );

            CREATE TABLE IF NOT EXISTS user_balance (
                guild_id INTEGER,
                user_id INTEGER,
                wallet INTEGER DEFAULT 0,
                bank INTEGER DEFAULT 0,
                PRIMARY KEY (guild_id, user_id)
            );

            CREATE TABLE IF NOT EXISTS warns (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                user_id INTEGER,
                moderator_id INTEGER,
                reason TEXT,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS reminders (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                channel_id INTEGER,
                message TEXT,
                due_unix INTEGER
            );

            CREATE TABLE IF NOT EXISTS user_daily_claim (
                guild_id INTEGER,
                user_id INTEGER,
                claimed_date TEXT,
                PRIMARY KEY (guild_id, user_id)
            );

            CREATE TABLE IF NOT EXISTS user_work_cooldown (
                guild_id INTEGER,
                user_id INTEGER,
                last_work TEXT,
                PRIMARY KEY (guild_id, user_id)
            );

            CREATE TABLE IF NOT EXISTS tournaments (
                guild_id INTEGER,
                name TEXT,
                status TEXT DEFAULT 'open',
                created_by INTEGER,
                winner_id INTEGER,
                PRIMARY KEY (guild_id, name)
            );

            CREATE TABLE IF NOT EXISTS tournament_participants (
                guild_id INTEGER,
                tournament_name TEXT,
                user_id INTEGER,
                PRIMARY KEY (guild_id, tournament_name, user_id)
            );

            CREATE TABLE IF NOT EXISTS tournament_games (
                guild_id INTEGER,
                tournament_name TEXT,
                game_name TEXT,
                PRIMARY KEY (guild_id, tournament_name, game_name)
            );

            CREATE TABLE IF NOT EXISTS giveaways (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                channel_id INTEGER,
                message_id INTEGER,
                prize TEXT,
                end_time INTEGER,
                winner_id INTEGER,
                ended INTEGER DEFAULT 0
            );

            CREATE TABLE IF NOT EXISTS giveaway_entries (
                giveaway_id INTEGER,
                user_id INTEGER,
                PRIMARY KEY (giveaway_id, user_id)
            );

            CREATE TABLE IF NOT EXISTS suggestions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                channel_id INTEGER,
                message_id INTEGER,
                user_id INTEGER,
                content TEXT,
                status TEXT DEFAULT 'pending'
            );

            CREATE TABLE IF NOT EXISTS suggestion_votes (
                suggestion_id INTEGER,
                user_id INTEGER,
                vote_type TEXT CHECK(vote_type IN ('up', 'down')),
                PRIMARY KEY (suggestion_id, user_id)
            );

            CREATE TABLE IF NOT EXISTS tickets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                channel_id INTEGER,
                user_id INTEGER,
                status TEXT DEFAULT 'open',
                created_at INTEGER
            );

            CREATE TABLE IF NOT EXISTS voice_sessions (
                guild_id INTEGER PRIMARY KEY,
                channel_id INTEGER,
                created_at INTEGER
            );

            CREATE TABLE IF NOT EXISTS server_backups (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                backup_data TEXT,
                created_by INTEGER,
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS user_warnings (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                user_id INTEGER,
                moderator_id INTEGER,
                reason TEXT,
                timestamp TEXT
            );

            CREATE TABLE IF NOT EXISTS economy_salaries (
                guild_id INTEGER PRIMARY KEY,
                min_salary INTEGER DEFAULT 50,
                max_salary INTEGER DEFAULT 150,
                daily_min INTEGER DEFAULT 100,
                daily_max INTEGER DEFAULT 250
            );

            CREATE TABLE IF NOT EXISTS saved_playlists (
                user_id INTEGER,
                name TEXT,
                tracks_json TEXT,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (user_id, name)
            );

            CREATE TABLE IF NOT EXISTS economy (
                user_id INTEGER PRIMARY KEY,
                coins INTEGER DEFAULT 0,
                xp INTEGER DEFAULT 0
            );

            CREATE TABLE IF NOT EXISTS shield_settings (
                guild_id INTEGER PRIMARY KEY,
                level TEXT DEFAULT 'medium'
            );

            CREATE TABLE IF NOT EXISTS shield_logs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                user_id INTEGER,
                reason TEXT,
                message_content TEXT,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS scam_images (
                guild_id INTEGER,
                image_hash TEXT,
                created_by INTEGER,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (guild_id, image_hash)
            );

            CREATE TABLE IF NOT EXISTS bot_presence_config (
                id INTEGER PRIMARY KEY CHECK (id = 1),
                status TEXT DEFAULT 'online',
                activity_type TEXT DEFAULT 'playing',
                activity_text TEXT DEFAULT 'CYBER // GRID',
                updated_by INTEGER,
                updated_at TEXT DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS ai_scheduled_tasks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                guild_id INTEGER,
                run_at TEXT,
                action_json TEXT,
                reason TEXT,
                created_by INTEGER,
                executed INTEGER DEFAULT 0,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            );
            """
        )

        db = await self._get_connection()
        await self._ensure_column(db, "guild_config", "daily_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "daily_message", "TEXT")
        await self._ensure_column(db, "guild_config", "daily_enabled", "INTEGER DEFAULT 0")
        await self._ensure_column(db, "guild_config", "daily_time", "TEXT DEFAULT '09:00'")
        await self._ensure_column(db, "guild_config", "daily_utc_offset", "INTEGER DEFAULT 0")
        await self._ensure_column(db, "guild_config", "last_daily_sent_date", "TEXT")
        await self._ensure_column(db, "guild_config", "free_games_platforms", "TEXT DEFAULT 'epic,steam,gog'")
        await self._ensure_column(db, "guild_config", "free_games_mention_type", "TEXT DEFAULT 'role'")
        await self._ensure_column(db, "guild_config", "verify_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "verify_role_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "daily_title", "TEXT DEFAULT '🌅 Daily Message'")
        await self._ensure_column(db, "guild_config", "daily_image_url", "TEXT")
        await self._ensure_column(db, "guild_config", "daily_button_label", "TEXT")
        await self._ensure_column(db, "guild_config", "daily_button_url", "TEXT")
        await self._ensure_column(db, "guild_config", "ai_model", "TEXT")
        await self._ensure_column(db, "guild_config", "ai_chat_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "ai_trigger", "TEXT DEFAULT 'bana'")
        await self._ensure_column(db, "guild_config", "ai_auto_enabled", "INTEGER DEFAULT 0")
        await self._ensure_column(db, "guild_config", "support_ai_enabled", "INTEGER DEFAULT 0")
        await self._ensure_column(db, "guild_config", "support_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "guild_language", "TEXT DEFAULT 'en'")
        await self._ensure_column(db, "guild_config", "dj_role_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "premium_tier", "INTEGER DEFAULT 0")
        await self._ensure_column(db, "guild_config", "music_autoplay", "INTEGER DEFAULT 0")
        await self._ensure_column(db, "guild_config", "poll_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "free_games_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "free_games_role_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "free_games_last_ids", "TEXT")
        await self._ensure_column(db, "guild_config", "wisdom_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "wisdom_enabled", "INTEGER DEFAULT 0")
        await self._ensure_column(db, "guild_config", "wisdom_last_sent_date", "TEXT")
        await self._ensure_column(db, "guild_config", "game_news_channel_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "game_news_role_id", "INTEGER")
        await self._ensure_column(db, "guild_config", "game_news_last_ids", "TEXT")
        await self._ensure_column(db, "guild_config", "custom_banner_url", "TEXT")

        await db.commit()
        await self._push_remote_db()

    async def execute(self, query: str, *params: object) -> None:
        """Execute a query with parameters using persistent connection."""
        async with self._lock:
            db = await self._get_connection()
            await db.execute(query, params)
            await db.commit()
        await self._push_remote_db()

    async def fetchone(self, query: str, *params: object):
        """Fetch a single row from a query using persistent connection."""
        async with self._lock:
            db = await self._get_connection()
            cursor = await db.execute(query, params)
            return await cursor.fetchone()

    async def fetchall(self, query: str, *params: object):
        """Fetch all rows from a query using persistent connection."""
        async with self._lock:
            db = await self._get_connection()
            cursor = await db.execute(query, params)
            return await cursor.fetchall()

    async def cleanup_guild(self, guild_id: int) -> None:
        """Clean up all data for a guild (for when bot leaves a server)."""
        async with self._lock:
            db = await self._get_connection()
            params = (guild_id,)
            await db.execute("DELETE FROM guild_config WHERE guild_id = ?", params)
            await db.execute("DELETE FROM user_xp WHERE guild_id = ?", params)
            await db.execute("DELETE FROM user_balance WHERE guild_id = ?", params)
            await db.execute("DELETE FROM user_daily_claim WHERE guild_id = ?", params)
            await db.execute("DELETE FROM user_work_cooldown WHERE guild_id = ?", params)
            await db.execute("DELETE FROM warns WHERE guild_id = ?", params)
            await db.execute("DELETE FROM tournaments WHERE guild_id = ?", params)
            await db.execute("DELETE FROM tournament_participants WHERE guild_id = ?", params)
            await db.execute("DELETE FROM tournament_games WHERE guild_id = ?", params)
            await db.execute("DELETE FROM voice_sessions WHERE guild_id = ?", params)
            await db.commit()
        await self._push_remote_db()