OriginDex/DBEditor/pokemon_db_ui.py

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

from event_system import event_system

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

class PokemonUI(QWidget):  # Change from QMainWindow to QWidget
    def __init__(self, parent=None):
        super().__init__(parent)
        self.init_ui()
        self.request_pokemon_list()

    def init_ui(self):
        main_layout = QVBoxLayout(self)  # Use self directly instead of a central widget

        # Create a tab widget
        self.tab_widget = QTabWidget()
        main_layout.addWidget(self.tab_widget)

        # Main tab (existing content)
        main_tab = QWidget()
        main_tab_layout = QHBoxLayout(main_tab)
        self.tab_widget.addTab(main_tab, "Main")

        # 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.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        #self.pokemon_list.customContextMenuRequested.connect(self.show_pokemon_context_menu)
        self.pokemon_list.currentItemChanged.connect(self.on_pokemon_selected)
        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)

        main_tab_layout.addLayout(left_layout, 1)

        # 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.is_baby_form_checkbox = QCheckBox("Is Baby Form")
        
        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)
        self.edit_form.addRow(self.is_baby_form_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_tab_layout.addLayout(right_layout, 1)

        # Database Operations tab
        db_tab = QWidget()
        db_tab_layout = QVBoxLayout(db_tab)
        self.tab_widget.addTab(db_tab, "Database Operations")

        # Add buttons to the Database Operations tab
        gather_forms_btn = QPushButton("Gather Pokémon Forms")
        #gather_forms_btn.clicked.connect(self.gather_pokemon_forms)
        db_tab_layout.addWidget(gather_forms_btn)

        gather_home_btn = QPushButton("Gather Home Storage Info")
        #gather_home_btn.clicked.connect(self.gather_home_storage_info)
        db_tab_layout.addWidget(gather_home_btn)

        gather_evolutions_btn = QPushButton("Gather Evolution Information")
        #gather_evolutions_btn.clicked.connect(self.gather_evolution_info)
        db_tab_layout.addWidget(gather_evolutions_btn)

        gather_encounters_btn = QPushButton("Gather Encounter Information")
        #gather_encounters_btn.clicked.connect(self.gather_encounter_info)
        db_tab_layout.addWidget(gather_encounters_btn)

        # Add QTextEdit for progress reporting
        self.progress_text = QTextEdit()
        self.progress_text.setReadOnly(True)
        self.progress_text.setMinimumHeight(200)  # Set a minimum height
        db_tab_layout.addWidget(self.progress_text)

        db_tab_layout.addStretch(1)  # Add stretch to separate the last button

        reinit_db_btn = QPushButton("Clear and Reinitialize Database")
        #reinit_db_btn.clicked.connect(self.reinitialize_database)
        db_tab_layout.addWidget(reinit_db_btn)

    def on_pokemon_selected(self, item):
        pfic = item.data(Qt.ItemDataRole.UserRole)
        pokemon_data = event_system.emit_sync('get_pokemon_data', data=pfic)

        if pokemon_data:
            name, form_name, national_dex, generation, storable_in_home = pokemon_data[0]
            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)
            self.is_baby_form_checkbox.setChecked(bool(False))

            # 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 load_evolution_chain(self, pfic):
        self.evolution_tree.clear()

        evolution_chain = event_system.emit_sync('get_evolution_chain', data=pfic)[0]

        # 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
            parent_pfic = event_system.emit_sync('get_evolution_parent', data=current_pfic)[0]
            
            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 on_save_button_clicked(self):
        # Gather data from UI elements
        data = {
            'name': self.name_input.text(),
            'form': self.form_input.text(),
            # ... other fields ...
        }
        event_system.emit('save_changes', data)

    def request_pokemon_list(self):
        event_system.emit('get_pokemon_list', self.update_pokemon_list)

    def update_pokemon_list(self, pokemon_data):
        self.pokemon_list.clear()

        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 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 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:
                home_storable = event_system.call_sync('get_home_storable', pfic)
            # Show the item only if it matches both filters
            item.setHidden(not (text_match and home_storable))

        self.update_pokemon_list_highlights()

    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 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)

        exclusive_group_combo = QComboBox()
        exclusive_group_combo.addItem("None", None)
        self.cursor.execute('SELECT id, group_name FROM exclusive_encounter_groups')
        for group_id, group_name in self.cursor.fetchall():
            exclusive_group_combo.addItem(group_name, group_id)

        layout.addRow("Exclusive Group:", exclusive_group_combo)

        # Fetch current values
        self.cursor.execute('''
            SELECT day, time, dual_slot, static_encounter_count, static_encounter, extra_text, stars, fishing, rods, exclusive_group_id
            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, exclusive_group_id = 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 "")
            index = exclusive_group_combo.findData(exclusive_group_id)
            if index >= 0:
                exclusive_group_combo.setCurrentIndex(index)

        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
            new_exclusive_group_id = exclusive_group_combo.currentData()
            
            # 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 = ?, exclusive_group_id = ?
                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, new_exclusive_group_id,
                  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)

        exclusive_group_combo = QComboBox()
        exclusive_group_combo.addItem("None", None)
        self.cursor.execute('SELECT id, group_name FROM exclusive_encounter_groups')
        for group_id, group_name in self.cursor.fetchall():
            exclusive_group_combo.addItem(group_name, group_id)

        layout.addRow("Exclusive Group:", exclusive_group_combo)

        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
            exclusive_group_id = exclusive_group_combo.currentData()
            
            # 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, exclusive_group_id)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (self.current_pfic, game, location, day, time, dual_slot, static_encounter, static_encounter_count, extra_text, stars, fishing, rods, exclusive_group_id))
            self.conn.commit()

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

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

    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 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)