Discord-Raid-bot/utils/DatabaseManager_class.py

# -*- coding: utf-8 -*-
from config import DATABASE_PATH, CLAN_MIGRATION
import sqlite3, os

class DatabaseManager:
    def __init__(self, db_path=DATABASE_PATH):
        self.db_path = db_path
        self.init_database()

    def init_database(self):
        """Initializes the database with all columns for difficulties"""
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Main table with all difficulties
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            discord_id TEXT UNIQUE,
            discord_username TEXT,

            -- Hydra difficulties
            pb_hydra_normal INTEGER DEFAULT 0,
            pb_hydra_normal_screenshot TEXT,
            pb_hydra_normal_date TIMESTAMP,
            pb_hydra_hard INTEGER DEFAULT 0,
            pb_hydra_hard_screenshot TEXT,
            pb_hydra_hard_date TIMESTAMP,
            pb_hydra_brutal INTEGER DEFAULT 0,
            pb_hydra_brutal_screenshot TEXT,
            pb_hydra_brutal_date TIMESTAMP,
            pb_hydra_nightmare INTEGER DEFAULT 0,
            pb_hydra_nightmare_screenshot TEXT,
            pb_hydra_nightmare_date TIMESTAMP,

            -- Chimera difficulties
            pb_chimera_easy INTEGER DEFAULT 0,
            pb_chimera_easy_screenshot TEXT,
            pb_chimera_easy_date TIMESTAMP,
            pb_chimera_normal INTEGER DEFAULT 0,
            pb_chimera_normal_screenshot TEXT,
            pb_chimera_normal_date TIMESTAMP,
            pb_chimera_hard INTEGER DEFAULT 0,
            pb_chimera_hard_screenshot TEXT,
            pb_chimera_hard_date TIMESTAMP,
            pb_chimera_brutal INTEGER DEFAULT 0,
            pb_chimera_brutal_screenshot TEXT,
            pb_chimera_brutal_date TIMESTAMP,
            pb_chimera_nightmare INTEGER DEFAULT 0,
            pb_chimera_nightmare_screenshot TEXT,
            pb_chimera_nightmare_date TIMESTAMP,
            pb_chimera_ultra INTEGER DEFAULT 0,
            pb_chimera_ultra_screenshot TEXT,
            pb_chimera_ultra_date TIMESTAMP,

            -- CvC (unchanged)
            pb_cvc INTEGER DEFAULT 0,
            pb_cvc_screenshot TEXT,
            pb_cvc_date TIMESTAMP,

            total_attempts INTEGER DEFAULT 0,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')

        # Existing data migration (if needed)
        cursor.execute("PRAGMA table_info(users)")
        columns = [row[1] for row in cursor.fetchall()]

        if 'discord_id' not in columns:
            cursor.execute('ALTER TABLE users ADD COLUMN discord_id TEXT')

        if 'clan' not in columns:
            cursor.execute('ALTER TABLE users ADD COLUMN clan TEXT')
            # Auto-migration: derive clan from old username prefix
            for old_tag, new_clan in CLAN_MIGRATION.items():
                cursor.execute(
                    "UPDATE users SET clan = ? WHERE clan IS NULL AND ("
                    "discord_username LIKE ? OR discord_username LIKE ?)",
                    (new_clan, f'[{old_tag}] %', f'[{old_tag}]%')
                )

        # Global history table
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS pb_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            discord_id TEXT,
            username TEXT,
            boss_type TEXT,
            difficulty TEXT,
            damage INTEGER,
            screenshot_filename TEXT,
            date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')

        conn.commit()
        conn.close()

    def get_user_pb(self, user_id, boss_type, difficulty=None):
        """Returns PB for a user on a specific boss and difficulty"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        if difficulty:
            column_prefix = f"pb_{boss_type}_{difficulty}"
        else:
            column_prefix = f"pb_{boss_type}"

        cursor.execute(
            f"SELECT {column_prefix}, {column_prefix}_screenshot, {column_prefix}_date FROM users WHERE discord_id = ?",
            (str(user_id),)
        )
        result = cursor.fetchone()
        conn.close()

        return result if result else (0, None, None)

    def update_user_pb(self, user_id, username, boss_type, damage, screenshot_filename, difficulty=None, clan=None):
        """Updates a user's PB and deletes the previous screenshot"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Get old screenshot for deletion
        old_data = self.get_user_pb(user_id, boss_type, difficulty)
        old_screenshot = old_data[1] if old_data else None

        if difficulty:
            column_prefix = f"pb_{boss_type}_{difficulty}"
        else:
            column_prefix = f"pb_{boss_type}"

        # COALESCE(?, clan): keeps existing clan if detection returns None
        cursor.execute(f'''
        INSERT INTO users (discord_id, discord_username, clan, {column_prefix}, {column_prefix}_screenshot, {column_prefix}_date, total_attempts)
        VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP, 1)
        ON CONFLICT(discord_id)
        DO UPDATE SET
            discord_username = ?,
            clan = COALESCE(?, clan),
            {column_prefix} = ?,
            {column_prefix}_screenshot = ?,
            {column_prefix}_date = CURRENT_TIMESTAMP,
            total_attempts = total_attempts + 1
        ''', (str(user_id), username, clan, damage, screenshot_filename, username, clan, damage, screenshot_filename))

        # Add to history
        cursor.execute('''
        INSERT INTO pb_history (discord_id, username, boss_type, difficulty, damage, screenshot_filename)
        VALUES (?, ?, ?, ?, ?, ?)
        ''', (str(user_id), username, boss_type, difficulty or 'none', damage, screenshot_filename))

        conn.commit()
        conn.close()

        return old_screenshot

    def get_leaderboard(self, boss_type, difficulty=None, limit=10, clan=None):
        """Returns the leaderboard for a specific boss and difficulty"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        if difficulty:
            column_prefix = f"pb_{boss_type}_{difficulty}"
        else:
            column_prefix = f"pb_{boss_type}"

        base_query = f'''
        SELECT discord_username, {column_prefix}, {column_prefix}_date, clan
        FROM users
        WHERE {column_prefix} > 0
        '''

        params = []
        if clan:
            base_query += ' AND clan = ?'
            params.append(clan)

        base_query += f' ORDER BY {column_prefix} DESC LIMIT ?'
        params.append(limit)

        cursor.execute(base_query, params)
        results = cursor.fetchall()
        conn.close()
        return results

    def get_user_all_pbs(self, user_id):
        """Returns all PBs for a user"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Retrieve all PB columns
        cursor.execute('SELECT * FROM users WHERE discord_id = ?', (str(user_id),))
        result = cursor.fetchone()
        columns = [desc[0] for desc in cursor.description]
        conn.close()

        if not result:
            return None

        return dict(zip(columns, result))

    def find_user_by_name(self, username):
        """Finds a user by name (for backwards compatibility)"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        cursor.execute('SELECT discord_id, discord_username FROM users WHERE discord_username LIKE ?', (f'%{username}%',))
        results = cursor.fetchall()
        conn.close()

        return results