OriginDex-DataManager/ui/workers/calculate_origin_mark_worker.py

from PyQt6.QtCore import QObject, pyqtSignal, QRunnable
import json

from cache import cache
from db import db

from utility.functions import get_display_name, get_shiftable_forms

class CalculateOriginMarkWorkerSignals(QObject):
    finished = pyqtSignal(dict)

class CalculateOriginMarkWorker(QRunnable):
    def __init__(self):
        super().__init__()
        self.signals = CalculateOriginMarkWorkerSignals()
        self.marks = {}

    def run(self):
        try:
            gathered_data = self.calculate_marks()
            self.signals.finished.emit(gathered_data)
        except Exception as e:
            print(f"Error gathering Pokémon home storage status: {e}")

    def calculate_marks(self):
        all_pokemon_forms = db.get_list_of_pokemon_forms()
        for form_entry in all_pokemon_forms:
            if form_entry["storable_in_home"] == False:
                continue

            print(f"Determining mark for {get_display_name(form_entry)}")

            target_generation = form_entry["generation"]
            pfic = form_entry["pfic"]

            #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.
            print("Checking Rule 1")
            chain = db.get_full_evolution_paths(pfic)
            if chain and (len(chain["predecessors"]) > 0 or len(chain["successors"]) > 0):
                base_form_in_generation = None
                last_pfic = pfic
                current_pfic = pfic
                while True:
                    current_pfic, _ = db.get_previous_evolution(current_pfic)
                    if current_pfic == None:
                        base_form_in_generation = last_pfic
                        break
                    chain_pokemon_data = db.get_pokemon_details(current_pfic)
                    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

                if base_form_in_generation and base_form_in_generation != pfic:
                    print(f"Base form in generation for {get_display_name(form_entry)} is {base_form_in_generation}")
                    mark_id = self.determine_origin_mark(base_form_in_generation, target_generation)
                    if mark_id != None:
                        self.marks[pfic] = mark_id
                        continue
                elif base_form_in_generation == pfic:
                    mark_id = self.determine_origin_mark(pfic, target_generation)
                    if mark_id != None:
                        self.marks[pfic] = mark_id
                        continue
            
            #Rule 2
            # 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
            print("Checking Rule 2")
            mark_id = self.determine_origin_mark(pfic, target_generation)
            if mark_id != None:
                self.marks[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.
            print("Checking Rule 3")
            mark_id = self.test_evolve_encounters(pfic, target_generation)
            if mark_id != None:
                self.marks[pfic] = mark_id
                continue

            #Rule 3b
            # Check to see if this is a sub-form pokemon, and if so, use the mark of the base form.
            random_encounters = db.get_encounters(pfic, "random")
            static_encounters = db.get_encounters(pfic, "static")
            encounters = []
            encounters.extend(random_encounters)
            encounters.extend(static_encounters)
            count = 0
            if encounters:
                for encounter in encounters:
                    game_info = db.get_game_by_id(encounter["game_id"])
                    if game_info["generation"] == target_generation:
                        count += 1
            if count == 0:
                
                shiftable_forms = get_shiftable_forms(pfic)
                if len(shiftable_forms) > 0:
                    form_found = False
                    for shiftable_form in shiftable_forms:
                        mark_id = self.determine_origin_mark(shiftable_form["to_pfic"], target_generation)
                        if mark_id != None:
                            self.marks[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_info = db.get_game_by_id(encounter["game_id"])
                    if game_info["id"] <= earliest_game:
                        earliest_game = game_info["id"]
                if earliest_game < 100:
                    form_info = db.get_pokemon_details(pfic)
                    mark_id = game_info["mark"]
                    if mark_id == None:
                        print(f"No mark found for {get_display_name(form_info)}")
                        pass
                    else:
                        print(f"Mark for {get_display_name(form_info)} is {mark_id}")
                        self.marks[pfic] = mark_id
                    continue
    
            event_encounters = db.get_encounters(pfic, "event")
            if event_encounters:
                earliest_game = 100
                for encounter in event_encounters:
                    game_info = game_info = db.get_game_by_id(encounter["game_id"])
                    if game_info["id"] <= earliest_game:
                        earliest_game = game_info["id"]
                if earliest_game < 100:
                    form_info = db.get_pokemon_details(pfic)
                    mark_id = game_info["mark"]
                    if mark_id == None:
                        print(f"No mark found for {get_display_name(form_info)}")
                    else:
                        print(f"Mark for {get_display_name(form_info)} is {mark_id}")
                        self.marks[pfic] = mark_id
                    continue

        return self.marks

    def determine_origin_mark(self, pfic, target_generation):
        shiftable_forms = get_shiftable_forms(pfic)
        if len(shiftable_forms) > 0:
            for shiftable_form in shiftable_forms:
                mark_id = self.determine_origin_mark(shiftable_form["to_pfic"], target_generation)
                return mark_id
        encounters = db.get_encounters(pfic)
        if encounters:
            generation_encounters = []
            for encounter in encounters:
                game_info = db.get_game_by_id(encounter["game_id"])
                encounter["game"] = game_info
                if encounter["game"]["generation"] == target_generation:
                    generation_encounters.append(encounter)
            if len(generation_encounters) > 0:
                generation_encounters = sorted(generation_encounters, key=lambda x: x["game"]["generation"])
                form_info = db.get_pokemon_details(pfic)
                game_info = generation_encounters[0]["game"]
                mark_id = game_info["mark"]
                if mark_id == None:
                    #self.logger.info(f"No mark found for {form_info[0]} {form_info[1]}")
                    print(f"No mark found for {get_display_name(form_info)}")
                else:
                    #self.logger.info(f"Mark for {form_info[0]} {form_info[1]} is {mark_id}")
                    print(f"Mark for {get_display_name(form_info)} is {mark_id}")
                    return mark_id
        return None
    
    def test_evolve_encounters(self, pfic, target_generation):
        evolve_encounters = db.get_encounters(pfic, "evolve")
        if evolve_encounters:
            available_encounters = []
            for encounter in evolve_encounters:
                game_info = db.get_game_by_id(encounter["game_id"])
                if game_info["generation"] == target_generation:
                    available_encounters.append(encounter)

            if len(available_encounters) > 0:
                available_encounters = sorted(available_encounters, key=lambda x: x.game_id)
                data = json.loads(available_encounters[0]["data"])
                mark_id = self.determine_origin_mark(data["from_pfic"], target_generation)
                if mark_id != None:
                    return mark_id
                
                mark_id = self.test_evolve_encounters(data["from_pfic"], target_generation)
                if mark_id != None:
                    return mark_id
                
        return None