database.py

import aiosqlite
import os

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DB_PATH = os.path.join(BASE_DIR, "data", "levels.db")

async def init_db():
    os.makedirs(os.path.join(BASE_DIR, "data"), exist_ok=True)
    async with aiosqlite.connect(DB_PATH) as db:
        # Enable Write-Ahead Logging for better concurrency and performance
        await db.execute("PRAGMA journal_mode=WAL;")
        
        await db.execute("""
            CREATE TABLE IF NOT EXISTS users (
                guild_id    TEXT NOT NULL,
                user_id     TEXT NOT NULL,
                xp          INTEGER DEFAULT 0,
                level       INTEGER DEFAULT 0,
                messages    INTEGER DEFAULT 0,
                voice_time  INTEGER DEFAULT 0,
                PRIMARY KEY (guild_id, user_id)
            )
        """)
        # Create an index to massively speed up leaderboard and get_rank queries
        await db.execute("CREATE INDEX IF NOT EXISTS idx_users_guild_xp ON users(guild_id, xp DESC);")
        
        await db.execute("""
            CREATE TABLE IF NOT EXISTS guild_settings (
                guild_id            TEXT PRIMARY KEY,
                xp_per_message      INTEGER DEFAULT 15,
                xp_per_voice_min    INTEGER DEFAULT 10,
                xp_cooldown         INTEGER DEFAULT 60,
                level_up_channel    TEXT DEFAULT NULL,
                level_up_message    TEXT DEFAULT NULL,
                xp_multiplier       REAL DEFAULT 1.0,
                no_xp_roles         TEXT DEFAULT '',
                no_xp_channels      TEXT DEFAULT ''
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS level_roles (
                guild_id    TEXT NOT NULL,
                level       INTEGER NOT NULL,
                role_id     TEXT NOT NULL,
                PRIMARY KEY (guild_id, level)
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS admins (
                user_id     TEXT PRIMARY KEY
            )
        """)
        await db.commit()

async def get_user(guild_id: int, user_id: int) -> dict:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM users WHERE guild_id=? AND user_id=?",
            (str(guild_id), str(user_id))
        ) as cursor:
            row = await cursor.fetchone()
            if row:
                return dict(row)
            # Auto-create
            await db.execute(
                "INSERT OR IGNORE INTO users (guild_id, user_id) VALUES (?,?)",
                (str(guild_id), str(user_id))
            )
            await db.commit()
            return {"guild_id": str(guild_id), "user_id": str(user_id),
                    "xp": 0, "level": 0, "messages": 0, "voice_time": 0}

async def update_user(guild_id: int, user_id: int, **kwargs):
    sets = ", ".join(f"{k}=?" for k in kwargs)
    vals = list(kwargs.values()) + [str(guild_id), str(user_id)]
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            f"UPDATE users SET {sets} WHERE guild_id=? AND user_id=?", vals
        )
        await db.commit()

async def add_xp(guild_id: int, user_id: int, amount: int) -> dict:
    """Add XP and return updated user dict."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR IGNORE INTO users (guild_id, user_id) VALUES (?,?)",
            (str(guild_id), str(user_id))
        )
        await db.execute(
            "UPDATE users SET xp=xp+? WHERE guild_id=? AND user_id=?",
            (amount, str(guild_id), str(user_id))
        )
        await db.commit()
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM users WHERE guild_id=? AND user_id=?",
            (str(guild_id), str(user_id))
        ) as cur:
            return dict(await cur.fetchone())

async def get_leaderboard(guild_id: int, limit: int = 10, offset: int = 0) -> list:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM users WHERE guild_id=? ORDER BY xp DESC LIMIT ? OFFSET ?",
            (str(guild_id), limit, offset)
        ) as cur:
            return [dict(r) for r in await cur.fetchall()]

async def get_user_count(guild_id: int) -> int:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT COUNT(*) FROM users WHERE guild_id=?", (str(guild_id),)
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else 0

async def get_rank(guild_id: int, user_id: int) -> int:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT COUNT(*)+1 FROM users WHERE guild_id=? AND xp > (SELECT xp FROM users WHERE guild_id=? AND user_id=?)",
            (str(guild_id), str(guild_id), str(user_id))
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else 1

async def get_settings(guild_id: int) -> dict:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM guild_settings WHERE guild_id=?", (str(guild_id),)
        ) as cur:
            row = await cur.fetchone()
            if row:
                return dict(row)
            await db.execute(
                "INSERT OR IGNORE INTO guild_settings (guild_id) VALUES (?)", (str(guild_id),)
            )
            await db.commit()
            return {"guild_id": str(guild_id), "xp_per_message": 15,
                    "xp_per_voice_min": 10, "xp_cooldown": 60,
                    "level_up_channel": None, "level_up_message": None,
                    "xp_multiplier": 1.0, "no_xp_roles": "", "no_xp_channels": "1475908072453964010"}

async def set_setting(guild_id: int, key: str, value):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR IGNORE INTO guild_settings (guild_id) VALUES (?)", (str(guild_id),)
        )
        await db.execute(
            f"UPDATE guild_settings SET {key}=? WHERE guild_id=?", (value, str(guild_id))
        )
        await db.commit()

async def get_level_roles(guild_id: int) -> list:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM level_roles WHERE guild_id=? ORDER BY level",
            (str(guild_id),)
        ) as cur:
            return [dict(r) for r in await cur.fetchall()]

async def set_level_role(guild_id: int, level: int, role_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR REPLACE INTO level_roles (guild_id, level, role_id) VALUES (?,?,?)",
            (str(guild_id), level, str(role_id))
        )
        await db.commit()

async def remove_level_role(guild_id: int, level: int):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "DELETE FROM level_roles WHERE guild_id=? AND level=?",
            (str(guild_id), level)
        )
        await db.commit()


async def reset_all_users():
    """Delete all user leveling data, leaving settings and roles intact."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("DELETE FROM users")
        await db.commit()

async def is_db_admin(user_id: int) -> bool:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT 1 FROM admins WHERE user_id=?", (str(user_id),)
        ) as cur:
            row = await cur.fetchone()
            return row is not None

async def add_admin(user_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR IGNORE INTO admins (user_id) VALUES (?)", (str(user_id),)
        )
        await db.commit()