OriginDex/DBEditor/DBEditor.py

import sys
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QListWidget, QLineEdit, 
                             QLabel, QCheckBox, QPushButton, QFormLayout, QListWidgetItem, QSplitter, QTreeWidget, 
                             QTreeWidgetItem, QDialog, QDialogButtonBox, QComboBox, QMessageBox, QSpinBox)
from PyQt6.QtCore import Qt, QSize
from PyQt6.QtGui import QPixmap, QFontMetrics, QColor
import sqlite3
import json
import os

class EvolutionEditDialog(QDialog):
    def __init__(self, parent=None, from_pfic=None, to_pfic=None, method=None):
        super().__init__(parent)
        self.setWindowTitle("Edit Evolution")
        self.setModal(True)

        layout = QVBoxLayout(self)

        self.from_combo = QComboBox()
        self.to_combo = QComboBox()
        self.method_edit = QLineEdit(method)

        layout.addWidget(QLabel("From Pokémon:"))
        layout.addWidget(self.from_combo)
        layout.addWidget(QLabel("To Pokémon:"))
        layout.addWidget(self.to_combo)
        layout.addWidget(QLabel("Evolution Method:"))
        layout.addWidget(self.method_edit)

        self.button_box = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        self.button_box.accepted.connect(self.accept)
        self.button_box.rejected.connect(self.reject)
        layout.addWidget(self.button_box)

        self.delete_button = QPushButton("Delete Evolution")
        self.delete_button.clicked.connect(self.delete_evolution)
        layout.addWidget(self.delete_button)

        self.populate_combos(from_pfic, to_pfic)

    def populate_combos(self, from_pfic, to_pfic):
        cursor = self.parent().cursor
        cursor.execute('SELECT PFIC, name, form_name FROM pokemon_forms ORDER BY name, form_name')
        for pfic, name, form_name in cursor.fetchall():
            display_name = f"{name} ({form_name})" if form_name else name
            self.from_combo.addItem(display_name, pfic)
            self.to_combo.addItem(display_name, pfic)
        
        if from_pfic:
            self.from_combo.setCurrentIndex(self.from_combo.findData(from_pfic))
        if to_pfic:
            self.to_combo.setCurrentIndex(self.to_combo.findData(to_pfic))

    def delete_evolution(self):
        self.done(2)  # Use a custom return code for delete action

def parse_pfic(pfic):
    parts = pfic.split('-')
    return tuple(int(part) if part.isdigit() else part for part in parts)

class DBEditor(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Pokémon Database Editor")
        self.setGeometry(100, 100, 1000, 600)

        self.conn = sqlite3.connect(':memory:')  # Use in-memory database for runtime
        self.cursor = self.conn.cursor()
        self.init_database()
        self.patches = self.load_and_apply_patches()

        self.encounter_cache = {}  # Add this line
        self.init_ui()
        #self.load_pokemon_list()  # Make sure this is called after init_ui

    def init_database(self):
        # Copy the original database to the in-memory database
        disk_conn = sqlite3.connect('pokemon_forms.db')
        disk_conn.backup(self.conn)
        disk_conn.close()

    def load_and_apply_patches(self):
        try:
            with open('patches.json', 'r') as f:
                patches = json.load(f)
        except FileNotFoundError:
            patches = {}

        # Apply patches to the in-memory database
        for patch_key, patch in patches.items():
            if patch_key.startswith('evolution_'):
                from_pfic, to_pfic = patch_key.split('_')[1:]
                if patch['action'] == 'delete':
                    self.cursor.execute('''
                        DELETE FROM evolution_chains
                        WHERE from_pfic = ? AND to_pfic = ?
                    ''', (from_pfic, to_pfic))
                elif patch['action'] == 'update':
                    self.cursor.execute('''
                        UPDATE evolution_chains
                        SET from_pfic = ?, to_pfic = ?, method = ?
                        WHERE from_pfic = ? AND to_pfic = ?
                    ''', (patch['new_from_pfic'], patch['new_to_pfic'], patch['new_method'], from_pfic, to_pfic))
                elif patch['action'] == 'add':
                    self.cursor.execute('''
                        INSERT OR REPLACE INTO evolution_chains (from_pfic, to_pfic, method)
                        VALUES (?, ?, ?)
                    ''', (patch['from_pfic'], patch['to_pfic'], patch['method']))
            else:  # pokemon_storage patches
                self.cursor.execute('''
                    UPDATE pokemon_storage
                    SET storable_in_home = ?
                    WHERE PFIC = ?
                ''', (patch['storable_in_home'], patch_key))

        self.conn.commit()
        return patches

    def save_patches(self):
        with open('patches.json', 'w') as f:
            json.dump(self.patches, f)

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)

        # Left side: Search and List
        left_layout = QVBoxLayout()
        search_layout = QHBoxLayout()
        self.search_bar = QLineEdit()
        self.search_bar.setPlaceholderText("Search Pokémon...")
        self.search_bar.textChanged.connect(self.filter_pokemon_list)
        search_layout.addWidget(self.search_bar)

        left_layout.addLayout(search_layout)

        self.pokemon_list = QListWidget()
        self.pokemon_list.currentItemChanged.connect(self.load_pokemon_details)
        left_layout.addWidget(self.pokemon_list)

        # Move the checkbox here, after the pokemon_list
        self.highlight_no_encounters = QCheckBox("Highlight Pokémon without encounters")
        self.highlight_no_encounters.stateChanged.connect(self.toggle_highlight_mode)
        left_layout.addWidget(self.highlight_no_encounters)

            # Add the new checkbox for filtering Home-storable Pokémon
        self.filter_home_storable = QCheckBox("Show only Home-storable Pokémon")
        self.filter_home_storable.stateChanged.connect(self.filter_pokemon_list)
        left_layout.addWidget(self.filter_home_storable)

        # Right side: Edit panel
        right_layout = QHBoxLayout()
        
        # Left side of right panel: Text information
        text_layout = QVBoxLayout()
        self.edit_form = QFormLayout()
        self.name_label = QLabel()
        self.form_name_label = QLabel()
        self.national_dex_label = QLabel()
        self.generation_label = QLabel()
        self.home_checkbox = QCheckBox("Available in Home")
        
        self.edit_form.addRow("Name:", self.name_label)
        self.edit_form.addRow("Form:", self.form_name_label)
        self.edit_form.addRow("National Dex:", self.national_dex_label)
        self.edit_form.addRow("Generation:", self.generation_label)
        self.edit_form.addRow(self.home_checkbox)

        text_layout.addLayout(self.edit_form)

        # Evolution chain tree
        self.evolution_tree = QTreeWidget()
        self.evolution_tree.setHeaderLabels(["Pokémon", "Evolution Method"])
        self.evolution_tree.setColumnWidth(0, 200)
        text_layout.addWidget(self.evolution_tree)

        # Add Locations tree
        self.locations_tree = QTreeWidget()
        self.locations_tree.setHeaderLabels(["Game/Location", "Details"])
        self.locations_tree.setColumnWidth(0, 200)
        self.locations_tree.itemDoubleClicked.connect(self.edit_encounter)
        text_layout.addWidget(QLabel("Locations:"))
        text_layout.addWidget(self.locations_tree)

        # Add New Encounter button
        self.add_encounter_button = QPushButton("Add New Encounter")
        self.add_encounter_button.clicked.connect(self.add_new_encounter)
        text_layout.addWidget(self.add_encounter_button)

        # Move buttons to the bottom
        text_layout.addStretch(1)

        self.save_button = QPushButton("Save Changes")
        self.save_button.clicked.connect(self.save_changes)
        text_layout.addWidget(self.save_button)

        self.export_button = QPushButton("Export Database")
        self.export_button.clicked.connect(self.export_database)
        text_layout.addWidget(self.export_button)

        # Add New Evolution button
        self.add_evolution_button = QPushButton("Add New Evolution")
        self.add_evolution_button.clicked.connect(self.add_new_evolution)
        text_layout.addWidget(self.add_evolution_button)

        # Right side of right panel: Image
        image_layout = QVBoxLayout()
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.image_label.setFixedSize(200, 200)
        image_layout.addWidget(self.image_label)
        image_layout.addStretch(1)

        right_layout.addLayout(text_layout)
        right_layout.addLayout(image_layout)

        main_layout.addLayout(left_layout, 1)
        main_layout.addLayout(right_layout, 1)

        self.load_pokemon_list()

    def adjust_list_width(self):
        max_width = 0
        font_metrics = QFontMetrics(self.pokemon_list.font())
        for i in range(self.pokemon_list.count()):
            item_width = font_metrics.horizontalAdvance(self.pokemon_list.item(i).text())
            max_width = max(max_width, item_width)

        # Add some padding to the width
        list_width = max_width + 50
        self.pokemon_list.setFixedWidth(list_width)
        self.search_bar.setFixedWidth(list_width)

    def load_pokemon_list(self):
        self.pokemon_list.clear()
        self.cursor.execute('''
            SELECT pf.PFIC, pf.name, pf.form_name, pf.national_dex
            FROM pokemon_forms pf
        ''')
        pokemon_data = self.cursor.fetchall()

        # Sort the pokemon_data based on PFIC
        pokemon_data.sort(key=lambda x: parse_pfic(x[0]))

        for pfic, name, form_name, national_dex in pokemon_data:
            display_name = f"{national_dex:04d} - {name}"
            if form_name:
                display_name += f" ({form_name})"
            item = QListWidgetItem(display_name)
            item.setData(Qt.ItemDataRole.UserRole, pfic)
            self.pokemon_list.addItem(item)

        self.update_encounter_cache()
        self.update_pokemon_list_highlights()
        self.adjust_list_width()
        self.filter_pokemon_list()

    def filter_pokemon_list(self):
        search_text = self.search_bar.text().lower()
        show_only_home_storable = self.filter_home_storable.isChecked()

        for i in range(self.pokemon_list.count()):
            item = self.pokemon_list.item(i)
            pfic = item.data(Qt.ItemDataRole.UserRole)
            
            # Check if the item matches the search text
            text_match = search_text in item.text().lower()
            
            # Check if the item is storable in Home (if the filter is active)
            home_storable = True
            if show_only_home_storable:
                self.cursor.execute('SELECT storable_in_home FROM pokemon_storage WHERE PFIC = ?', (pfic,))
                result = self.cursor.fetchone()
                home_storable = result[0] if result else False

            # Show the item only if it matches both filters
            item.setHidden(not (text_match and home_storable))

        self.update_pokemon_list_highlights()

    def load_pokemon_details(self, current, previous):
        if not current:
            return

        pfic = current.data(Qt.ItemDataRole.UserRole)
        self.cursor.execute('''
            SELECT pf.name, pf.form_name, pf.national_dex, pf.generation, ps.storable_in_home
            FROM pokemon_forms pf
            LEFT JOIN pokemon_storage ps ON pf.PFIC = ps.PFIC
            WHERE pf.PFIC = ?
        ''', (pfic,))
        pokemon_data = self.cursor.fetchone()

        if pokemon_data:
            name, form_name, national_dex, generation, storable_in_home = pokemon_data
            self.name_label.setText(name)
            self.form_name_label.setText(form_name if form_name else "")
            self.national_dex_label.setText(str(national_dex))
            self.generation_label.setText(str(generation))
            self.home_checkbox.setChecked(bool(storable_in_home))
            self.home_checkbox.stateChanged.connect(self.update_home_storable)

            # Load and display the image
            image_path = f"images-new/{pfic}.png"
            if os.path.exists(image_path):
                pixmap = QPixmap(image_path)
                self.image_label.setPixmap(pixmap.scaled(200, 200, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation))
            else:
                self.image_label.setText("Image not found")

            # Load and display evolution chain
            self.load_evolution_chain(pfic)

            # Load and display encounter locations
            self.load_encounter_locations(pfic)

        self.current_pfic = pfic
        self.add_evolution_button.setEnabled(True)  # Enable the button when a Pokémon is selected

    def update_home_storable(self):
        if hasattr(self, 'current_pfic'):
            storable_in_home = self.home_checkbox.isChecked()
            self.cursor.execute('UPDATE pokemon_storage SET storable_in_home = ? WHERE PFIC = ?', (storable_in_home, self.current_pfic))
            self.conn.commit()
            self.filter_pokemon_list()  # Reapply the filter

    def edit_evolution(self, item, column):
        parent = item.parent()
        if not parent:
            return  # Don't edit the root item

        from_pfic = parent.data(0, Qt.ItemDataRole.UserRole)
        to_pfic = item.data(0, Qt.ItemDataRole.UserRole)
        method = item.text(1)

        dialog = EvolutionEditDialog(self, from_pfic, to_pfic, method)
        result = dialog.exec()

        if result == QDialog.DialogCode.Accepted:
            new_from_pfic = dialog.from_combo.currentData()
            new_to_pfic = dialog.to_combo.currentData()
            new_method = dialog.method_edit.text()

            # Update the in-memory database
            self.cursor.execute('''
                UPDATE evolution_chains
                SET from_pfic = ?, to_pfic = ?, method = ?
                WHERE from_pfic = ? AND to_pfic = ?
            ''', (new_from_pfic, new_to_pfic, new_method, from_pfic, to_pfic))

            # Create or update the patch
            patch_key = f"evolution_{from_pfic}_{to_pfic}"
            self.patches[patch_key] = {
                'action': 'update',
                'new_from_pfic': new_from_pfic,
                'new_to_pfic': new_to_pfic,
                'new_method': new_method
            }
            self.save_patches()

            # Refresh the evolution chain display
            self.load_evolution_chain(self.current_pfic)

        elif result == 2:  # Delete action
            confirm = QMessageBox.question(self, "Confirm Deletion", "Are you sure you want to delete this evolution?",
                                           QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
            if confirm == QMessageBox.StandardButton.Yes:
                # Delete from the in-memory database
                self.cursor.execute('''
                    DELETE FROM evolution_chains
                    WHERE from_pfic = ? AND to_pfic = ?
                ''', (from_pfic, to_pfic))

                # Create a delete patch
                patch_key = f"evolution_{from_pfic}_{to_pfic}"
                self.patches[patch_key] = {'action': 'delete'}
                self.save_patches()

                # Refresh the evolution chain display
                self.load_evolution_chain(self.current_pfic)

    def get_evolution_chain(self, pfic):
        def follow_chain(start_pfic, visited=None):
            if visited is None:
                visited = set()
            
            chain = []
            stack = [(start_pfic, None)]
            
            while stack:
                current_pfic, method = stack.pop()
                if current_pfic in visited:
                    continue
                
                visited.add(current_pfic)
                self.cursor.execute('SELECT name, form_name FROM pokemon_forms WHERE PFIC = ?', (current_pfic,))
                name, form_name = self.cursor.fetchone()
                
                chain.append((current_pfic, name, form_name, method))
                
                # Get previous evolutions
                self.cursor.execute('SELECT from_pfic, method FROM evolution_chains WHERE to_pfic = ?', (current_pfic,))
                prev_evolutions = self.cursor.fetchall()
                for prev_pfic, prev_method in prev_evolutions:
                    if prev_pfic not in visited:
                        stack.append((prev_pfic, prev_method))
                
                # Get next evolutions
                self.cursor.execute('SELECT to_pfic, method FROM evolution_chains WHERE from_pfic = ?', (current_pfic,))
                next_evolutions = self.cursor.fetchall()
                for next_pfic, next_method in next_evolutions:
                    if next_pfic not in visited:
                        stack.append((next_pfic, next_method))
            
            return chain

        return follow_chain(pfic)

    def load_evolution_chain(self, pfic):
        self.evolution_tree.clear()
        evolution_chain = self.get_evolution_chain(pfic)

        # Create a dictionary to store tree items
        tree_items = {}

        # First pass: create all tree items
        for current_pfic, name, form_name, method in evolution_chain:
            display_name = f"{name} ({form_name})" if form_name else name
            item = QTreeWidgetItem([display_name, method if method else ""])
            item.setData(0, Qt.ItemDataRole.UserRole, current_pfic)
            tree_items[current_pfic] = item

            if current_pfic == pfic:
                item.setBackground(0, QColor(255, 255, 0, 100))  # Highlight selected Pokémon

        # Second pass: build the tree structure
        root = None
        for current_pfic, name, form_name, method in evolution_chain:
            item = tree_items[current_pfic]
            
            # Find the parent of this item
            self.cursor.execute('SELECT from_pfic FROM evolution_chains WHERE to_pfic = ?', (current_pfic,))
            parent_pfic = self.cursor.fetchone()
            
            if parent_pfic:
                parent_item = tree_items.get(parent_pfic[0])
                if parent_item:
                    parent_item.addChild(item)
            elif not root:
                root = item
                self.evolution_tree.addTopLevelItem(root)

        # Expand the entire tree
        self.evolution_tree.expandAll()

        # Scroll to and select the current Pokémon
        current_item = tree_items[pfic]
        self.evolution_tree.scrollToItem(current_item)
        self.evolution_tree.setCurrentItem(current_item)

        # Connect double-click signal
        self.evolution_tree.itemDoubleClicked.connect(self.edit_evolution)

    def save_changes(self):
        if hasattr(self, 'current_pfic'):
            storable_in_home = self.home_checkbox.isChecked()
            self.patches[self.current_pfic] = {'storable_in_home': storable_in_home}
            
            with open('patches.json', 'w') as f:
                json.dump(self.patches, f)

    def export_database(self):
        export_conn = sqlite3.connect('pokemon_forms_production.db')
        self.conn.backup(export_conn)
        export_conn.close()

    def closeEvent(self, event):
        self.conn.close()
        event.accept()

    def add_new_evolution(self):
        if not hasattr(self, 'current_pfic'):
            return

        dialog = EvolutionEditDialog(self, self.current_pfic, None, "")
        result = dialog.exec()

        if result == QDialog.DialogCode.Accepted:
            new_from_pfic = dialog.from_combo.currentData()
            new_to_pfic = dialog.to_combo.currentData()
            new_method = dialog.method_edit.text()

            # Update the in-memory database
            self.cursor.execute('''
                INSERT INTO evolution_chains (from_pfic, to_pfic, method)
                VALUES (?, ?, ?)
            ''', (new_from_pfic, new_to_pfic, new_method))

            # Create a new patch
            patch_key = f"evolution_{new_from_pfic}_{new_to_pfic}"
            self.patches[patch_key] = {
                'action': 'add',
                'from_pfic': new_from_pfic,
                'to_pfic': new_to_pfic,
                'method': new_method
            }
            self.save_patches()

            # Refresh the evolution chain display
            self.load_evolution_chain(self.current_pfic)

    def load_encounter_locations(self, pfic):
        self.locations_tree.clear()
        self.cursor.execute('''
            SELECT game, location, day, time, dual_slot, static_encounter_count, static_encounter, extra_text, stars, rods, fishing
            FROM encounters
            WHERE pfic = ?
            ORDER BY game, location
        ''', (pfic,))
        encounters = self.cursor.fetchall()

        game_items = {}
        for encounter in encounters:
            game, location, day, time, dual_slot, static_encounter_count, static_encounter, extra_text, stars, rods, fishing = encounter
            
            if game not in game_items:
                game_item = QTreeWidgetItem([game])
                game_items[game] = game_item
                # Use generation for sorting, default to 0 if not found
                game_item.setData(0, Qt.ItemDataRole.UserRole, self.game_generations.get(game, 0))

            location_item = QTreeWidgetItem([location])
            details = []
            if day:
                details.append(f"Day: {day}")
            if time:
                details.append(f"Time: {time}")
            if dual_slot:
                details.append(f"Dual Slot: {dual_slot}")
            if static_encounter:
                details.append(f"Static Encounter (Count: {static_encounter_count})")
            if extra_text:
                details.append(f"Extra: {extra_text}")
            if stars:
                details.append(f"Stars: {stars}")
            if fishing:
                details.append(f"Fishing")
                if rods:
                    details.append(f"Rods: {rods}")
            
            location_item.setText(1, ", ".join(details))
            game_items[game].addChild(location_item)

        # Sort game items by generation and add them to the tree
        sorted_game_items = sorted(game_items.values(), key=lambda x: x.data(0, Qt.ItemDataRole.UserRole))
        self.locations_tree.addTopLevelItems(sorted_game_items)
        self.locations_tree.expandAll()

        # Update the cache for this Pokémon
        self.encounter_cache[pfic] = len(encounters) > 0

        # After updating the locations tree
        self.update_pokemon_list_highlights()

    def edit_encounter(self, item, column):
        if item.parent() is None:  # This is a game item, not a location item
            return

        game = item.parent().text(0)
        location = item.text(0)
        
        dialog = QDialog(self)
        dialog.setWindowTitle("Edit Encounter")
        layout = QFormLayout(dialog)

        game_edit = QLineEdit(game)
        location_edit = QLineEdit(location)
        day_edit = QLineEdit()
        time_edit = QLineEdit()
        dual_slot_edit = QLineEdit()
        static_encounter_check = QCheckBox("Static Encounter")
        static_encounter_count_edit = QSpinBox()
        extra_text_edit = QLineEdit()
        stars_edit = QLineEdit()
        fishing_check = QCheckBox("Fishing")
        rods_edit = QLineEdit()

        layout.addRow("Game:", game_edit)
        layout.addRow("Location:", location_edit)
        layout.addRow("Day:", day_edit)
        layout.addRow("Time:", time_edit)
        layout.addRow("Dual Slot:", dual_slot_edit)
        layout.addRow("Static Encounter:", static_encounter_check)
        layout.addRow("Static Encounter Count:", static_encounter_count_edit)
        layout.addRow("Extra Text:", extra_text_edit)
        layout.addRow("Stars:", stars_edit)
        layout.addRow("Fishing:", fishing_check)
        layout.addRow("Rods:", rods_edit)

        # Fetch current values
        self.cursor.execute('''
            SELECT day, time, dual_slot, static_encounter_count, static_encounter, extra_text, stars, fishing, rods
            FROM encounters
            WHERE pfic = ? AND game = ? AND location = ?
        ''', (self.current_pfic, game, location))
        current_values = self.cursor.fetchone()

        if current_values:
            day, time, dual_slot, static_encounter_count, static_encounter, extra_text, stars, fishing, rods = current_values
            day_edit.setText(day or "")
            time_edit.setText(time or "")
            dual_slot_edit.setText(dual_slot or "")
            static_encounter_check.setChecked(bool(static_encounter))
            static_encounter_count_edit.setValue(static_encounter_count or 0)
            extra_text_edit.setText(extra_text or "")
            stars_edit.setText(stars or "")
            fishing_check.setChecked(bool(fishing))
            rods_edit.setText(rods or "")

        buttons = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel,
            Qt.Orientation.Horizontal, dialog)
        buttons.accepted.connect(dialog.accept)
        buttons.rejected.connect(dialog.reject)
        layout.addRow(buttons)

        if dialog.exec() == QDialog.DialogCode.Accepted:
            new_game = game_edit.text()
            new_location = location_edit.text()
            new_day = day_edit.text() or None
            new_time = time_edit.text() or None
            new_dual_slot = dual_slot_edit.text() or None
            new_static_encounter = static_encounter_check.isChecked()
            new_static_encounter_count = static_encounter_count_edit.value()
            new_extra_text = extra_text_edit.text() or None
            new_stars = stars_edit.text() or None
            new_fishing = fishing_check.isChecked()
            new_rods = rods_edit.text() or None
            
            # Update the database
            self.cursor.execute('''
                UPDATE encounters
                SET game = ?, location = ?, day = ?, time = ?, dual_slot = ?, 
                    static_encounter = ?, static_encounter_count = ?, extra_text = ?,
                    stars = ?, fishing = ?, rods = ?
                WHERE pfic = ? AND game = ? AND location = ?
            ''', (new_game, new_location, new_day, new_time, new_dual_slot, 
                  new_static_encounter, new_static_encounter_count, new_extra_text,
                  new_stars, new_fishing, new_rods,
                  self.current_pfic, game, location))
            self.conn.commit()

            # Update the cache if all encounters for this Pokémon were deleted
            if not self.check_pokemon_has_encounters(self.current_pfic):
                self.encounter_cache[self.current_pfic] = False

            # Refresh the locations tree
            self.load_encounter_locations(self.current_pfic)

    def add_new_encounter(self):
        dialog = QDialog(self)
        dialog.setWindowTitle("Add New Encounter")
        layout = QFormLayout(dialog)

        game_edit = QLineEdit()
        location_edit = QLineEdit()
        day_edit = QLineEdit()
        time_edit = QLineEdit()
        dual_slot_edit = QLineEdit()
        static_encounter_check = QCheckBox("Static Encounter")
        static_encounter_count_edit = QSpinBox()
        extra_text_edit = QLineEdit()
        stars_edit = QLineEdit()
        fishing_check = QCheckBox("Fishing")
        rods_edit = QLineEdit()

        layout.addRow("Game:", game_edit)
        layout.addRow("Location:", location_edit)
        layout.addRow("Day:", day_edit)
        layout.addRow("Time:", time_edit)
        layout.addRow("Dual Slot:", dual_slot_edit)
        layout.addRow("Static Encounter:", static_encounter_check)
        layout.addRow("Static Encounter Count:", static_encounter_count_edit)
        layout.addRow("Extra Text:", extra_text_edit)
        layout.addRow("Stars:", stars_edit)
        layout.addRow("Fishing:", fishing_check)
        layout.addRow("Rods:", rods_edit)

        buttons = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel,
            Qt.Orientation.Horizontal, dialog)
        buttons.accepted.connect(dialog.accept)
        buttons.rejected.connect(dialog.reject)
        layout.addRow(buttons)

        if dialog.exec() == QDialog.DialogCode.Accepted:
            game = game_edit.text()
            location = location_edit.text()
            day = day_edit.text() or None
            time = time_edit.text() or None
            dual_slot = dual_slot_edit.text() or None
            static_encounter = static_encounter_check.isChecked()
            static_encounter_count = static_encounter_count_edit.value()
            extra_text = extra_text_edit.text() or None
            stars = stars_edit.text() or None
            fishing = fishing_check.isChecked()
            rods = rods_edit.text() or None
            
            # Insert new encounter into the database
            self.cursor.execute('''
                INSERT INTO encounters 
                (pfic, game, location, day, time, dual_slot, static_encounter, static_encounter_count, extra_text, stars, fishing, rods)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (self.current_pfic, game, location, day, time, dual_slot, static_encounter, static_encounter_count, extra_text, stars, fishing, rods))
            self.conn.commit()

            # Update the cache
            self.encounter_cache[self.current_pfic] = True

            # Refresh the locations tree
            self.load_encounter_locations(self.current_pfic)

    # Add this as a class attribute in the DBEditor class
    game_generations = {
        "Red": 1, "Blue": 1, "Yellow": 1,
        "Gold": 2, "Silver": 2, "Crystal": 2,
        "Ruby": 3, "Sapphire": 3, "Emerald": 3, "FireRed": 3, "LeafGreen": 3,
        "Diamond": 4, "Pearl": 4, "Platinum": 4, "HeartGold": 4, "SoulSilver": 4,
        "Black": 5, "White": 5, "Black 2": 5, "White 2": 5,
        "X": 6, "Y": 6, "Omega Ruby": 6, "Alpha Sapphire": 6,
        "Sun": 7, "Moon": 7, "Ultra Sun": 7, "Ultra Moon": 7,
        "Sword": 8, "Shield": 8, "Brilliant Diamond": 8, "Shining Pearl": 8, "Expansion Pass": 8,
        "Legends: Arceus": 8,
        "Scarlet": 9, "Violet": 9, "The Teal Mask": 9, "The Hidden Treasure of Area Zero": 9, "The Hidden Treasure of Area Zero (Scarlet)": 9, "The Hidden Treasure of Area Zero (Violet)": 9, "The Teal Mask (Scarlet)": 9, "The Teal Mask (Violet)": 9,   
        "Pokémon Go": 0, "Pokémon Home": 0
    }

    def toggle_highlight_mode(self):
        self.update_pokemon_list_highlights()

    def update_pokemon_list_highlights(self):
        highlight_mode = self.highlight_no_encounters.isChecked()
        for i in range(self.pokemon_list.count()):
            item = self.pokemon_list.item(i)
            pfic = item.data(Qt.ItemDataRole.UserRole)
            
            if highlight_mode:
                has_encounters = self.encounter_cache.get(pfic, False)
                if not has_encounters:
                    item.setData(Qt.ItemDataRole.BackgroundRole, QColor(255, 200, 200))  # Light red background
                else:
                    item.setData(Qt.ItemDataRole.BackgroundRole, None)  # White background
            else:
                item.setData(Qt.ItemDataRole.BackgroundRole, None)  # White background

    def update_encounter_cache(self):
        self.cursor.execute('''
            SELECT DISTINCT pfic
            FROM encounters
        ''')
        pokemon_with_encounters = set(row[0] for row in self.cursor.fetchall())
        
        for i in range(self.pokemon_list.count()):
            item = self.pokemon_list.item(i)
            pfic = item.data(Qt.ItemDataRole.UserRole)
            self.encounter_cache[pfic] = pfic in pokemon_with_encounters

    def check_pokemon_has_encounters(self, pfic):
        return self.encounter_cache.get(pfic, False)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    editor = DBEditor()
    editor.show()
    sys.exit(app.exec())