OriginDex/Utilities/DatabaseBuilder.py

import sqlite3
import csv
import re

def create_connection():
    conn = sqlite3.connect('pokemon_database.db')
    return conn

def create_tables(conn):
    cursor = conn.cursor()
    
    # Create games table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS games (
        id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        generation INTEGER NOT NULL
    )
    ''')
    
    # Create marks table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS marks (
        id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        icon_path TEXT NOT NULL
    )
    ''')
    
    # Create mark_game_associations table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS mark_game_associations (
        mark_id INTEGER,
        game_id INTEGER,
        FOREIGN KEY (mark_id) REFERENCES marks (id),
        FOREIGN KEY (game_id) REFERENCES games (id),
        PRIMARY KEY (mark_id, game_id)
    )
    ''')
    
    # Create pokemon table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS pokemon (
        national_dex_number INTEGER PRIMARY KEY,
        name TEXT NOT NULL
    )
    ''')
    
    # Create pokemon_forms table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS pokemon_forms (
        id INTEGER PRIMARY KEY,
        pokemon_id INTEGER NOT NULL,
        form_name TEXT NOT NULL,
        is_default BOOLEAN NOT NULL,
        image_path TEXT NOT NULL,
        FOREIGN KEY (pokemon_id) REFERENCES pokemon (national_dex_number),
        UNIQUE (pokemon_id, form_name)
    )
    ''')
    
    # Create encounter_methods table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS encounter_methods (
        id INTEGER PRIMARY KEY,
        name TEXT UNIQUE NOT NULL
    )
    ''')
    
    # Create locations table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS locations (
        id INTEGER PRIMARY KEY,
        name TEXT UNIQUE NOT NULL,
        description TEXT
    )
    ''')
    
    # Create form_encounters table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS form_encounters (
        id INTEGER PRIMARY KEY,
        form_id INTEGER NOT NULL,
        game_id INTEGER NOT NULL,
        location_id INTEGER NOT NULL,
        encounter_method_id INTEGER NOT NULL,
        FOREIGN KEY (form_id) REFERENCES pokemon_forms (id),
        FOREIGN KEY (game_id) REFERENCES games (id),
        FOREIGN KEY (location_id) REFERENCES locations (id),
        FOREIGN KEY (encounter_method_id) REFERENCES encounter_methods (id),
        UNIQUE (form_id, game_id, location_id, encounter_method_id)
    )
    ''')
    
    # Create alternate_game_names table
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS alternate_game_names (
        id INTEGER PRIMARY KEY,
        game_id INTEGER NOT NULL,
        alternate_name TEXT NOT NULL,
        FOREIGN KEY (game_id) REFERENCES games (id),
        UNIQUE (alternate_name COLLATE NOCASE)
    )
    ''')
    
    conn.commit()

def load_game_data(conn):
    cursor = conn.cursor()
    
    games = [
        ("Red", 1, ["Red Version"]),
        ("Blue", 1, ["Blue Version"]),
        ("Yellow", 1, ["Yellow Version"]),
        ("Gold", 2, ["Gold Version"]),
        ("Silver", 2, ["Silver Version"]),
        ("Crystal", 2, ["Crystal Version"]),
        ("Ruby", 3, ["Ruby Version"]),
        ("Sapphire", 3, ["Sapphire Version"]),
        ("Emerald", 3, ["Emerald Version"]),
        ("FireRed", 3, ["Fire Red", "Fire-Red"]),
        ("LeafGreen", 3, ["Leaf Green", "Leaf-Green"]),
        ("Diamond", 4, ["Diamond Version"]),
        ("Pearl", 4, ["Pearl Version"]),
        ("Platinum", 4, ["Platinum Version"]),
        ("HeartGold", 4, ["Heart Gold", "Heart-Gold"]),
        ("SoulSilver", 4, ["Soul Silver", "Soul-Silver"]),
        ("Black", 5, ["Black Version"]),
        ("White", 5, ["White Version"]),
        ("Black 2", 5, ["Black Version 2", "Black-2"]),
        ("White 2", 5, ["White Version 2", "White-2"]),
        ("X", 6, ["X Version"]),
        ("Y", 6, ["Y Version"]),
        ("Omega Ruby", 6, ["Omega Ruby Version", "Omega-Ruby"]),
        ("Alpha Sapphire", 6, ["Alpha Sapphire Version", "Alpha-Sapphire"]),
        ("Sun", 7, ["Sun Version"]),
        ("Moon", 7, ["Moon Version"]),
        ("Ultra Sun", 7, ["Ultra Sun Version", "Ultra-Sun"]),
        ("Ultra Moon", 7, ["Ultra Moon Version", "Ultra-Moon"]),
        ("Let's Go Pikachu", 7, ["Let's Go, Pikachu!", "Lets Go Pikachu"]),
        ("Let's Go Eevee", 7, ["Let's Go, Eevee!", "Lets Go Eevee"]),
        ("Sword", 8, ["Sword Version"]),
        ("Shield", 8, ["Shield Version"]),
        ("Brilliant Diamond", 8, ["Brilliant Diamond Version", "Brilliant-Diamond"]),
        ("Shining Pearl", 8, ["Shining Pearl Version", "Shining-Pearl"]),
        ("Legends Arceus", 8, ["Legends: Arceus", "Legends-Arceus"]),
        ("Scarlet", 9, ["Scarlet Version"]),
        ("Violet", 9, ["Violet Version"]),
    ]
    
    for game in games:
        cursor.execute('''
        INSERT OR IGNORE INTO games (name, generation)
        VALUES (?, ?)
        ''', (game[0], game[1]))
        
        game_id = cursor.lastrowid
        
        # Insert alternate names
        for alt_name in game[2]:
            cursor.execute('''
            INSERT OR IGNORE INTO alternate_game_names (game_id, alternate_name)
            VALUES (?, ?)
            ''', (game_id, alt_name))
    
    conn.commit()

def load_mark_data(conn):
    cursor = conn.cursor()
    
    marks = [
        ("Game Boy", "images/marks/GB_icon_HOME.png", ["Red", "Blue", "Yellow", "Gold", "Silver", "Crystal", "Ruby", "Sapphire", "Emerald", "FireRed", "LeafGreen"]),
        ("Kalos", "images/marks/Blue_pentagon_HOME.png", ["X", "Y", "Omega Ruby", "Alpha Sapphire"]),
        ("Alola", "images/marks/Black_clover_HOME.png", ["Sun", "Moon", "Ultra Sun", "Ultra Moon"]),
        ("Let's Go", "images/marks/Let's_Go_icon_HOME.png", ["Let's Go Pikachu", "Let's Go Eevee"]),
        ("Galar", "images/marks/Galar_symbol_HOME.png", ["Sword", "Shield"]),
        ("Sinnoh", "images/marks/BDSP_icon_HOME.png", ["Brilliant Diamond", "Shining Pearl"]),
        ("Hisui", "images/marks/Arceus_mark_HOME.png", ["Legends Arceus"]),
        ("Paldea", "images/marks/Paldea_icon_HOME.png", ["Scarlet", "Violet"]),
    ]
    
    for mark in marks:
        cursor.execute('''
        INSERT OR IGNORE INTO marks (name, icon_path)
        VALUES (?, ?)
        ''', (mark[0], mark[1]))
        
        mark_id = cursor.lastrowid
        
        for game_name in mark[2]:
            cursor.execute('''
            INSERT OR IGNORE INTO mark_game_associations (mark_id, game_id)
            SELECT ?, id FROM games WHERE name = ?
            ''', (mark_id, game_name))
    
    conn.commit()

def load_pokemon_data(conn):
    cursor = conn.cursor()
    
    with open('pokemon_home_list.csv', 'r') as f:
        reader = csv.reader(f)
        next(reader)  # Skip header row if it exists
        for row in reader:
            national_dex_number = int(row[0])
            full_name = row[1]
            
            # Extract the base name and form name
            match = re.match(r'([^(]+)(?:\s*\(([^)]+)\))?', full_name)
            if match:
                base_name = match.group(1).strip()
                form_name = match.group(2).strip() if match.group(2) else "Default"
            else:
                base_name = full_name
                form_name = "Default"
            
            # Insert or update the pokemon entry
            cursor.execute('''
            INSERT OR IGNORE INTO pokemon (national_dex_number, name)
            VALUES (?, ?)
            ''', (national_dex_number, base_name))
            
            # Create the image path
            padded_dex = f"{national_dex_number:04d}"
            if form_name == "Default":
                image_path = f"images/pokemon/{padded_dex}_{base_name}.png"
            else:
                image_path = f"images/pokemon/{padded_dex}_{base_name}_({form_name}).png".replace(" ", "_")
            
            # Insert the form entry
            cursor.execute('''
            INSERT OR IGNORE INTO pokemon_forms (pokemon_id, form_name, is_default, image_path)
            VALUES (?, ?, ?, ?)
            ''', (national_dex_number, form_name, form_name == "Default", image_path))
    
    conn.commit()

def load_encounter_data(conn):
    cursor = conn.cursor()
    
    with open('pokemon_earliest_games.csv', 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            national_dex_number = int(row['number'])
            full_name = row['name']
            earliest_game = row['earliest_game']
            obtain_method = row['obtain_method']
            encounter_locations = row['encounter_locations']
            
            # Extract the base name and form name
            match = re.match(r'([^(]+)(?:\s*\(([^)]+)\))?', full_name)
            if match:
                base_name = match.group(1).strip()
                form_name = match.group(2).strip() if match.group(2) else "Default"
            else:
                base_name = full_name
                form_name = "Default"
            
            # Ensure the Pokémon and form exist in the database
            cursor.execute('''
            INSERT OR IGNORE INTO pokemon (national_dex_number, name)
            VALUES (?, ?)
            ''', (national_dex_number, base_name))
            
            cursor.execute('''
            INSERT OR IGNORE INTO pokemon_forms (pokemon_id, form_name, is_default, image_path)
            VALUES (?, ?, ?, ?)
            ''', (national_dex_number, form_name, form_name == "Default", f"images/pokemon/{national_dex_number:04d}_{base_name}.png"))
            
            # Skip encounter data if it's unknown or N/A
            if earliest_game == "Unknown" or obtain_method == "Unknown":
                continue
            
            # Get the form_id
            cursor.execute('''
            SELECT id FROM pokemon_forms
            WHERE pokemon_id = ? AND form_name = ?
            ''', (national_dex_number, form_name))
            form_id = cursor.fetchone()[0]
            
            # Get the game_id (now case-insensitive and including alternate names)
            cursor.execute('''
            SELECT g.id FROM games g
            LEFT JOIN alternate_game_names agn ON g.id = agn.game_id
            WHERE g.name = ? COLLATE NOCASE OR agn.alternate_name = ? COLLATE NOCASE
            ''', (earliest_game, earliest_game))
            result = cursor.fetchone()
            if result:
                game_id = result[0]
            else:
                print(f"Warning: Game '{earliest_game}' not found for {full_name}")
                continue
            
            # Handle gift Pokémon
            if obtain_method.lower() == "gift" and (encounter_locations == "N/A" or not encounter_locations):
                # Insert or get the "Gift" location
                cursor.execute('''
                INSERT OR IGNORE INTO locations (name, description)
                VALUES (?, ?)
                ''', ("Gift", "Pokémon received as a gift"))
                cursor.execute('SELECT id FROM locations WHERE name = ?', ("Gift",))
                location_id = cursor.fetchone()[0]
                
                # Insert or get the "gift" encounter method
                cursor.execute('INSERT OR IGNORE INTO encounter_methods (name) VALUES (?)', ("gift",))
                cursor.execute('SELECT id FROM encounter_methods WHERE name = ?', ("gift",))
                method_id = cursor.fetchone()[0]
                
                # Insert form_encounter for gift Pokémon
                cursor.execute('''
                INSERT OR IGNORE INTO form_encounters 
                (form_id, game_id, location_id, encounter_method_id)
                VALUES (?, ?, ?, ?)
                ''', (form_id, game_id, location_id, method_id))
            elif encounter_locations != "N/A" and encounter_locations:
                # Process each encounter location
                for location_info in encounter_locations.split('|'):
                    location, method = location_info.strip().rsplit(' ', 1)
                    method = method.strip('()')
                    
                    # Insert or get location_id
                    cursor.execute('''
                    INSERT OR IGNORE INTO locations (name, description)
                    VALUES (?, ?)
                    ''', (location, None))
                    cursor.execute('SELECT id FROM locations WHERE name = ?', (location,))
                    location_id = cursor.fetchone()[0]
                    
                    # Insert or get encounter_method_id
                    cursor.execute('INSERT OR IGNORE INTO encounter_methods (name) VALUES (?)', (method,))
                    cursor.execute('SELECT id FROM encounter_methods WHERE name = ?', (method,))
                    method_id = cursor.fetchone()[0]
                    
                    # Insert form_encounter
                    cursor.execute('''
                    INSERT OR IGNORE INTO form_encounters 
                    (form_id, game_id, location_id, encounter_method_id)
                    VALUES (?, ?, ?, ?)
                    ''', (form_id, game_id, location_id, method_id))
    
    conn.commit()

def main():
    conn = create_connection()
    create_tables(conn)
    load_game_data(conn)
    load_mark_data(conn)
    load_pokemon_data(conn)
    load_encounter_data(conn)
    conn.close()
    print("All data has been successfully added to the database.")

if __name__ == "__main__":
    main()