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, QMenu, QTabWidget,
                             QTextEdit)
from PyQt6.QtCore import Qt, QSize, qInstallMessageHandler, QThread, pyqtSignal
from PyQt6.QtGui import QPixmap, QFontMetrics, QColor, QAction
import sqlite3
import json
import os
import traceback
import time
import pdb
import debugpy

# Add the parent directory to the Python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Now try to import
from db_controller import DBController
from pokemon_db_ui import PokemonUI
from DataGatherers.update_location_information import process_pokemon_for_location_data
from DataGatherers.pokemondb_scraper import retrieve_all_pokemon_forms
from DataGatherers.update_storable_in_home import update_storable_in_home
from DataGatherers.update_location_information import profile_update_location_information
from DataGatherers.Update_evolution_information import update_evolution_chains
from DataGatherers.cache_manager import CacheManager
from event_system import event_system
from db_classes import PokemonForm

import logging

class UILogHandler(logging.Handler):
    def __init__(self):
        super().__init__()

    def emit(self, record):
        log_entry = self.format(record)  # Format the log message
        event_system.emit('update_log_display', log_entry)

class PokemonFormGatherer(QThread):
    progress_signal = pyqtSignal(str)
    finished_signal = pyqtSignal()

    def run(self):
        cache = CacheManager()
        debugpy.debug_this_thread()
        retrieve_all_pokemon_forms(cache, progress_callback=self.progress_signal.emit)
        self.finished_signal.emit()

class PokemonHomeStatusUpdater(QThread):
    progress_signal = pyqtSignal(str)
    finished_signal = pyqtSignal()

    def run(self):
        cache = CacheManager()
        debugpy.debug_this_thread()
        update_storable_in_home(cache, progress_callback=self.progress_signal.emit)
        self.finished_signal.emit()

class PokemonEvolutionUpdater(QThread):
    progress_signal = pyqtSignal(str)
    finished_signal = pyqtSignal()

    def run(self):
        time.sleep(0.1)
        debugpy.debug_this_thread()
        cache = CacheManager()
        update_evolution_chains(cache, progress_callback=self.progress_signal.emit)
        self.finished_signal.emit()

class PokemonEncounterUpdater(QThread):
    progress_signal = pyqtSignal(str)
    finished_signal = pyqtSignal()

    def run(self):
        time.sleep(0.1)
        #debugpy.debug_this_thread()
        cache = CacheManager()
        profile_update_location_information(cache, progress_callback=self.progress_signal.emit)
        self.finished_signal.emit()

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

        event_system.add_listener('save_changes', self.save_changes)
        event_system.add_listener('add_new_encounter', self.add_new_encounter)
        event_system.add_listener('refresh_pokemon_encounters', self.refresh_pokemon_encounters)
        event_system.add_listener('gather_pokemon_forms', self.gather_pokemon_forms)
        event_system.add_listener('gather_home_storage_info', self.gather_home_storage_info)
        event_system.add_listener('gather_evolution_info', self.gather_evolution_info)
        event_system.add_listener('gather_encounter_info', self.gather_encounter_info)
        event_system.add_listener('reinitialize_database', self.reinitialize_database)
        event_system.add_listener('gather_marks_info', self.gather_marks_info)

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

        self.init_ui()
        

    def init_database(self):
        self.db_controller = DBController()
        self.db_controller.init_database()

    def setup_logging(self):
        # Create the logger with the name 'ui_feedback'
        self.logger = logging.getLogger('ui_feedback')
        self.logger.setLevel(logging.DEBUG)  # Set the logging level

        # Create handlers
        console_handler = logging.StreamHandler()  # Log to the console
        file_handler = logging.FileHandler('ui_feedback.log')  # Log to a file
        ui_handler = UILogHandler()

        # Create formatters and add them to handlers
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        console_handler.setFormatter(formatter)
        file_handler.setFormatter(formatter)
        ui_handler.setFormatter(formatter)

        # Add the handler to the logger
        self.logger.addHandler(ui_handler)

        # Add handlers to the logger
        self.logger.addHandler(console_handler)
        self.logger.addHandler(file_handler)

    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_changes(self, data):
        # Your existing code to save changes
        pass

    def add_new_encounter(self, data):
        # Your existing code to add a new encounter
        pass

    def init_ui(self):
        # Create a central widget for the main window
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        # Create the PokemonUI as a child widget
        self.pokemon_ui = PokemonUI(self)
        main_layout.addWidget(self.pokemon_ui)

    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 closeEvent(self, event):
        self.conn.close()
        event.accept()    

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

    def manage_exclusive_groups(self):
        dialog = QDialog(self)
        dialog.setWindowTitle("Manage Exclusive Encounter Groups")
        layout = QVBoxLayout(dialog)

        group_list = QListWidget()
        self.cursor.execute('SELECT id, group_name, description FROM exclusive_encounter_groups')
        for group_id, group_name, description in self.cursor.fetchall():
            item = QListWidgetItem(f"{group_name} - {description}")
            item.setData(Qt.ItemDataRole.UserRole, group_id)
            group_list.addItem(item)

        layout.addWidget(group_list)

        add_button = QPushButton("Add Group")
        edit_button = QPushButton("Edit Group")
        delete_button = QPushButton("Delete Group")

        button_layout = QHBoxLayout()
        button_layout.addWidget(add_button)
        button_layout.addWidget(edit_button)
        button_layout.addWidget(delete_button)
        layout.addLayout(button_layout)

        add_button.clicked.connect(lambda: self.add_edit_exclusive_group(dialog, group_list))
        edit_button.clicked.connect(lambda: self.add_edit_exclusive_group(dialog, group_list, group_list.currentItem()))
        delete_button.clicked.connect(lambda: self.delete_exclusive_group(dialog, group_list, group_list.currentItem()))

        dialog.exec()

    def add_edit_exclusive_group(self, parent_dialog, group_list, item=None):
        dialog = QDialog(parent_dialog)
        dialog.setWindowTitle("Add/Edit Exclusive Group")
        layout = QFormLayout(dialog)

        name_edit = QLineEdit()
        description_edit = QLineEdit()
        game_combo = QComboBox()

        layout.addRow("Group Name:", name_edit)
        layout.addRow("Description:", description_edit)
        layout.addRow("Game:", game_combo)

        # Populate game combo box
        self.cursor.execute('SELECT id, name FROM games ORDER BY name')
        for game_id, game_name in self.cursor.fetchall():
            game_combo.addItem(game_name, game_id)

        if item:
            group_id = item.data(Qt.ItemDataRole.UserRole)
            self.cursor.execute('SELECT group_name, description, game_id FROM exclusive_encounter_groups WHERE id = ?', (group_id,))
            group_name, description, game_id = self.cursor.fetchone()
            name_edit.setText(group_name)
            description_edit.setText(description)
            game_combo.setCurrentIndex(game_combo.findData(game_id))

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

        if dialog.exec() == QDialog.DialogCode.Accepted:
            name = name_edit.text()
            description = description_edit.text()
            game_id = game_combo.currentData()
            if item:
                self.cursor.execute('UPDATE exclusive_encounter_groups SET group_name = ?, description = ?, game_id = ? WHERE id = ?', (name, description, game_id, group_id))
                item.setText(f"{name} - {description} ({game_combo.currentText()})")
            else:
                self.cursor.execute('INSERT INTO exclusive_encounter_groups (group_name, description, game_id) VALUES (?, ?, ?)', (name, description, game_id))
                group_id = self.cursor.lastrowid
                new_item = QListWidgetItem(f"{name} - {description} ({game_combo.currentText()})")
                new_item.setData(Qt.ItemDataRole.UserRole, group_id)
                group_list.addItem(new_item)
            self.conn.commit()

    def delete_exclusive_group(self, parent_dialog, group_list, item):
        if item:
            group_id = item.data(Qt.ItemDataRole.UserRole)
            reply = QMessageBox.question(parent_dialog, 'Delete Group', 'Are you sure you want to delete this group?', QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No)
            if reply == QMessageBox.StandardButton.Yes:
                self.cursor.execute('DELETE FROM exclusive_encounter_groups WHERE id = ?', (group_id,))
                self.cursor.execute('UPDATE encounters SET exclusive_group_id = NULL WHERE exclusive_group_id = ?', (group_id,))
                self.conn.commit()
                group_list.takeItem(group_list.row(item))

    def refresh_pokemon_encounters(self, pfic):
        pokemon_data = event_system.call_sync('get_pokemon_data', pfic)

        if pokemon_data:
            name, form_name, national_dex, generation, storable_in_home, is_baby_form = pokemon_data

            # Create a temporary connection for this operation
            temp_conn = sqlite3.connect('pokemon_forms.db')

            try:
                with open('./DataGatherers/DefaultForms.json', 'r') as f:
                    default_forms = json.load(f)
            except FileNotFoundError:
                default_forms = []

            # Create a cache manager instance
            cache = CacheManager()

            event_system.emit_sync('clear_encounters_for_pokemon', pfic)

            # Process the Pokémon data
            process_pokemon_for_location_data(pfic, name, form_name, national_dex, default_forms, cache, temp_conn)

            event_system.emit_sync('refresh_in_memory_db', temp_conn)

            # Close the temporary connection
            temp_conn.close()

            event_system.emit_sync('refresh_pokemon_details_panel', pfic)

            print(f"Refreshed encounters for {name} {form_name if form_name else ''}")

    def gather_pokemon_forms(self, data):
        event_system.emit_sync('clear_log_display')
        self.logger.info("Starting to gather Pokémon forms...")
        
        self.form_gatherer = PokemonFormGatherer()
        self.form_gatherer.progress_signal.connect(self.update_progress)
        self.form_gatherer.finished_signal.connect(self.gathering_finished)
        self.form_gatherer.start()

    def update_progress(self, message):
        self.logger.info(message)

    def gathering_finished(self):
        self.logger.info("Finished gathering Pokémon forms.")

    def gather_home_storage_info(self, data):
        event_system.emit_sync('clear_log_display')
        self.logger.info("Starting to gather Home storage information...")
        self.home_storage_updater = PokemonHomeStatusUpdater()
        self.home_storage_updater.progress_signal.connect(self.update_progress)
        self.home_storage_updater.finished_signal.connect(self.updating_home_storage_finished)
        self.home_storage_updater.start()

    def updating_home_storage_finished(self):
        self.logger.info("Finished updating Home storage information.")

    def gather_evolution_info(self, data):
        event_system.call_sync('clear_log_display')
        self.logger.info("Starting to gather evolution information...")
        self.evolution_updater = PokemonEvolutionUpdater()
        self.evolution_updater.progress_signal.connect(self.update_progress)
        self.evolution_updater.finished_signal.connect(self.updating_evolution_info_finished)
        self.evolution_updater.start()

    def updating_evolution_info_finished(self):
        self.logger.info("Finished updating evolution information.")


    def gather_encounter_info(self, data):
        event_system.call_sync('clear_log_display')
        self.logger.info("Starting to gather encounter information...")
        self.encounter_updater = PokemonEncounterUpdater()
        self.encounter_updater.progress_signal.connect(self.update_progress)
        self.encounter_updater.finished_signal.connect(self.updating_encounter_info_finished)
        self.encounter_updater.start()

    def updating_encounter_info_finished(self):
        self.logger.info("Finished updating encounter information.")

    def reinitialize_database(self, data):
        reply = QMessageBox.question(self, 'Confirm Action', 
                                     'Are you sure you want to clear and reinitialize the database? This action cannot be undone.',
                                     QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No)
        if reply == QMessageBox.StandardButton.Yes:
            self.logger.info("Starting database reinitialization...")
            # Implement the logic to clear and reinitialize the database
            self.conn.close()
            os.remove('pokemon_forms.db')  # Remove the existing database file
            self.conn = sqlite3.connect('pokemon_forms.db')
            self.cursor = self.conn.cursor()
            self.init_database()  # Reinitialize the database structure
            self.logger.info("Database reinitialized successfully.")
            QMessageBox.information(self, 'Database Reinitialized', 'The database has been cleared and reinitialized.')

    def determine_origin_mark(self, pfic, target_generation):
        encounters = event_system.call_sync('get_encounter_locations', pfic)
        if encounters:
            generation_encounters = []
            for encounter in encounters:
                game_generation = event_system.call_sync('get_game_generation', encounter[0])
                game_id = event_system.call_sync('get_game_id_for_name', encounter[0])  
                encounter = encounter + (game_generation, game_id)
                if game_generation == target_generation:
                    generation_encounters.append(encounter)
            if len(generation_encounters) > 0:
                generation_encounters = sorted(generation_encounters, key=lambda x: x[12])
                form_info = event_system.call_sync('get_pokemon_data', pfic)
                mark_id = event_system.call_sync('get_mark_for_game_name', generation_encounters[0][0])
                if mark_id == None:
                    self.logger.info(f"No mark found for {form_info[0]} {form_info[1]}")
                else:
                    mark_details = event_system.call_sync('get_mark_details', mark_id)
                    self.logger.info(f"Mark for {form_info[0]} {form_info[1]} is {mark_details[0]} - {mark_details[1]}")
                    return mark_id
        return None

    def gather_marks_info(self, data):
        event_system.emit_sync('clear_log_display')
        self.logger.info("Starting to gather marks information...")

        pokemon_list = event_system.call_sync('get_pokemon_list')

        for pfic, name, form_name, national_dex in pokemon_list:
            pokemon_data = event_system.call_sync('get_pokemon_form_by_pfic', pfic)

            if pokemon_data.storable_in_home == False:
                continue;
            
            self.logger.info(f"Determining mark for {pokemon_data.name} {pokemon_data.form_name if pokemon_data.form_name else ''}")

            target_generation = pokemon_data.generation

            #Rule 1
            # 1. If a pokemon form has a previous evolution from within the same generation, 
            # use the mark of the previous evolution. This should be recursive within the same generation.
            chain = event_system.call_sync('get_evolution_chain', pfic)
            if chain:
                base_form_in_generation = None
                last_pfic = pfic
                current_pfic = pfic
                while True:
                    current_pfic = event_system.call_sync('get_evolution_parent', last_pfic)
                    if current_pfic == None:
                        base_form_in_generation = last_pfic
                        break
                    chain_pokemon_data = event_system.call_sync('get_pokemon_form_by_pfic', current_pfic[0])
                    if chain_pokemon_data.generation == target_generation:
                        base_form_in_generation = current_pfic
                    else:
                        base_form_in_generation = last_pfic
                        break
                    last_pfic = current_pfic[0]

                if base_form_in_generation and base_form_in_generation != pfic:
                    self.logger.info(f"Base form in generation for {name} {form_name if form_name else ''} is {base_form_in_generation}")
                    mark_id = self.determine_origin_mark(base_form_in_generation, target_generation)
                    if mark_id != None:
                        event_system.emit_sync('assign_mark_to_form', (pfic, mark_id))
                        continue;
                elif base_form_in_generation == pfic:
                    mark_id = self.determine_origin_mark(pfic, target_generation)
                    if mark_id != None:
                        event_system.emit_sync('assign_mark_to_form', (pfic, mark_id))
                        continue;
            #Rule 2a
            # If a pokemon form has no previous evolution from within the same generation, 
            # look at the encounters of the pokemon form from this generation and use the mark of the earliest
            # game you can encounter that form in from that generation
            mark_id = self.determine_origin_mark(pfic, target_generation)
            if mark_id != None:
                event_system.emit_sync('assign_mark_to_form', (pfic, mark_id))
                continue;

            #Rule 3
            # If there are no encounters for the pokemon form from this generation,
            # look to see if a previous evolution has an encounter from this generation, and use the mark of the earliest
            # game from this generation that the previous evolution is encounterable in.
            evolve_encounters = event_system.call_sync('get_evolve_encounters', pfic)
            if evolve_encounters:
                available_encounters = []
                for encounter in evolve_encounters:
                    game_id = encounter.game_id
                    game_info = event_system.call_sync('get_game_by_id', game_id)
                    if game_info[2] == target_generation:
                        available_encounters.append(encounter)

                if len(available_encounters) > 0:
                    available_encounters = sorted(available_encounters, key=lambda x: x.game_id)
                    mark_id = self.determine_origin_mark(available_encounters[0].from_pfic, target_generation)
                    if mark_id != None:
                        event_system.emit_sync('assign_mark_to_form', (pfic, mark_id))
                        continue;
            
            encounters = event_system.call_sync('get_encounter_locations', pfic)
            count = 0
            if encounters:
                for encounter in encounters:
                    game_id = event_system.call_sync('get_game_id_for_name', encounter[0])
                    game_generation = event_system.call_sync('get_game_generation', game_id)
                    if game_generation == target_generation:
                        count += 1
            if count == 0:
                #Rule 2b
                # Check to see if this is a sub-form pokemon, and if so, use the mark of the base form.
                base = pfic[:-6]
                alternate_forms = event_system.call_sync('find_default_form', base)
                if alternate_forms:
                    form_found = False
                    for alternate_form in alternate_forms:
                        mark_id = self.determine_origin_mark(alternate_form.pfic, target_generation)
                        if mark_id != None:
                            event_system.emit_sync('assign_mark_to_form', (pfic, mark_id))
                            form_found = True
                            break;
                    if form_found:
                        continue;

            #Rule 4
            # If there are no encounters for the pokemon form or its evolution line from this generation,
            # use the mark of the earliest game of the generation is marked as being introducted in.
            if encounters:
                earliest_game = 100
                for encounter in encounters:
                    game_id = event_system.call_sync('get_game_id_for_name', encounter[0])
                    if game_id <= earliest_game:
                        earliest_game = game_id
                if earliest_game < 100:
                    form_info = event_system.call_sync('get_pokemon_data', pfic)
                    mark_id = event_system.call_sync('get_mark_for_game', earliest_game)
                    if mark_id == None:
                        self.logger.info(f"No mark found for {form_info[0]} {form_info[1]}")
                    else:
                        mark_details = event_system.call_sync('get_mark_details', mark_id)
                        self.logger.info(f"Mark for {form_info[0]} {form_info[1]} is {mark_details[0]} - {mark_details[1]}")
                        event_system.emit_sync('assign_mark_to_form', (pfic, mark_id))
                    continue;

            

def qt_message_handler(mode, context, message):
    print(f"Qt Message: {mode} {context} {message}")

def exception_hook(exctype, value, tb):
    print(''.join(traceback.format_exception(exctype, value, tb)))
    sys.exit(1)

sys.excepthook = exception_hook

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