kemoverse/bot/db_utils.py

from random import choices
import sqlite3
import config

DB_PATH = config.DB_PATH
CONNECTION: sqlite3.Connection
CURSOR: sqlite3.Cursor

def connect() -> None:
    '''Creates a connection to the database'''
    print('Connecting to the database...')
    global CONNECTION
    global CURSOR
    CONNECTION = sqlite3.connect(DB_PATH, autocommit=True)
    CONNECTION.row_factory = sqlite3.Row
    CURSOR = CONNECTION.cursor()

def get_random_character():
    ''' Gets a random character from the database'''
    CURSOR.execute('SELECT * FROM characters')
    characters = CURSOR.fetchall()

    if not characters:
        return None, None, None, None

    weights = [c['weight'] for c in characters]
    chosen = choices(characters, weights=weights, k=1)[0]

    return chosen['id'], chosen['name'], chosen['file_id'], chosen['rarity']


# User functions

def get_or_create_user(username):
    '''Retrieves an ID for a given user, if the user does not exist, it will be
    created.'''
    CURSOR.execute('SELECT id FROM users WHERE username = ?', (username,))
    user = CURSOR.fetchone()
    if user:
        return user[0]

    # New user starts with has_rolled = False
    CURSOR.execute(
        'INSERT INTO users (username, has_rolled) VALUES (?, ?)',
        (username, False)
    )
    user_id = CURSOR.lastrowid
    return user_id

def insert_character(name: str, rarity: int, weight: float, file_id: str) -> int:
    '''Inserts a character'''
    CURSOR.execute(
        'INSERT INTO characters (name, rarity, weight, file_id) VALUES (?, ?, ?, ?)',
        (name, rarity, weight, file_id)
    )
    character_id = CURSOR.lastrowid
    return character_id if character_id else 0

def insert_pull(user_id, character_id):
    '''Creates a pull in the database'''
    CURSOR.execute(
        'INSERT INTO pulls (user_id, character_id) VALUES (?, ?)',
        (user_id, character_id)
    )

def get_last_rolled_at(user_id):
    '''Gets the timestamp when the user last rolled'''
    CURSOR.execute("SELECT timestamp FROM pulls WHERE user_id = ? ORDER BY timestamp DESC", \
            (user_id,))
    row = CURSOR.fetchone()
    return row[0] if row else None

# Configuration

def get_config(key):
    '''Reads the value for a specified config key from the db'''
    CURSOR.execute("SELECT value FROM config WHERE key = ?", (key,))
    row = CURSOR.fetchone()
    return row[0] if row else None

def set_config(key, value):
    '''Writes the value for a specified config key to the db'''
    cur.execute("INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)", (key, value))


# Character stat functions

def add_character_stats(character_id, stats):
    '''
    Adds or updates character stats in the character_stats table.
    `stats` should be a dictionary like {'power': 5, 'charm': 3}
    '''
    if not stats:
        return

    conn = get_db_connection()
    cur = conn.cursor()

    columns = ', '.join(stats.keys())
    placeholders = ', '.join(['?'] * len(stats))
    updates = ', '.join([f"{col}=excluded.{col}" for col in stats.keys()])
    
    values = list(stats.values())

    sql = f'''
        INSERT INTO character_stats (character_id, {columns})
        VALUES (?, {placeholders})
        ON CONFLICT(character_id) DO UPDATE SET {updates}
    '''
    cur.execute(sql, [character_id] + values)
    conn.commit()
    conn.close()


def update_character_stat(character_id, stat_name, value):
    '''Updates a single stat field for a character'''
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute(f'''
        UPDATE character_stats SET {stat_name} = ? WHERE character_id = ?
    ''', (value, character_id))
    conn.commit()
    conn.close()

def get_character_stats(character_id):
    '''Retrieves all stats for a single character dynamically'''
    conn = get_db_connection()
    conn.row_factory = sqlite3.Row  # Enables dict-style access to rows
    cur = conn.cursor()
    cur.execute('SELECT * FROM character_stats WHERE character_id = ?', (character_id,))
    row = cur.fetchone()
    conn.close()

    if row:
        return {key: row[key] for key in row.keys() if key != 'character_id'}
    else:
        return {}

def get_character_stat(character_id, stat_name):
    '''Retrieves a single stat value for a character'''
    if stat_name not in ('power', 'charm'):
        raise ValueError("Invalid stat name")
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute(f'SELECT {stat_name} FROM character_stats WHERE character_id = ?', (character_id,))
    row = cur.fetchone()
    conn.close()
    return row[0] if row else 0

def get_stats_for_multiple_characters(character_ids):
    '''
    Retrieves stats for a list of character IDs.
    Returns a dictionary of character_id -> {stat_name: value, ...}
    '''
    if not character_ids:
        return {}

    placeholders = ','.join('?' for _ in character_ids)
    query = f'''
        SELECT *
        FROM character_stats
        WHERE character_id IN ({placeholders})
    '''

    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute(query, character_ids)
    rows = cur.fetchall()
    col_names = [desc[0] for desc in cur.description]
    conn.close()

    result = {}
    for row in rows:
        character_id = row[0]
        stats = dict(zip(col_names[1:], row[1:]))  # Skip character_id
        result[character_id] = stats

    return result