'''
Author: Dawson Matthews
Date: 01/14/26

Takes as input: 
    1. The directory in which all of the json files are listed in their respective 
        directories seperated by generation from generation1 to generation9
    2. The file containing all desired mega pokemon and their target abilities in 
        the format as shown by: Darkrai-Mega: Dark Aura
    3. The name of the new directory to put the changed files in

Produces as output a new directory with the fixed files in it.
'''

import sys
import os
import json

def get_target_abilities(ability_data_path):
    '''
    Reads in the data from the path and converts date from the showdown format to the
    cobblemon format as a pythyon dictionary i.e. "Darkrai-Mega: Dark Aura" -> {"darkrai": "darkaura"}
    '''

    ability_data = {}

    with open(ability_data_path, "r", encoding="utf-8") as f:
        lines = f.readlines()
        for line in lines:
            line = (line.strip()).split(":") # Removing new line characters, as well as leading and trailing spaces. Then split it by the middle colon

            species_name = line[0].split("-")[0].lower() # Removes the "-Mega" suffix and sends the name to lower case

            ability = (line[1].replace(" ", "")) # Removes the spaces 
            ability = ability.replace("-", "") # Removes the dashes
            ability = ability.lower()
            ability_data[species_name] = ability

    return ability_data

def write_augmented_species_data(species_additions_dir, dst_dir, ability_data):
    '''
    Takes the new species data in dictionary form and creates the new data
    '''

    gen_dirs = os.listdir(species_additions_dir) # Get a list of all directories in the data path, each corresponding to one generation
    for gen_dir in gen_dirs:

        dst_gen_dir_path = os.path.join(dst_dir, gen_dir)
        os.makedirs(dst_gen_dir_path, exist_ok=True) # Make the new output dir for later
        dir_path = os.path.join(species_additions_dir, gen_dir)

        if os.path.isdir(dir_path) == False: # Make sure it's actually a real file path
            return 1

        species_files = os.listdir(dir_path)

        for species_file_name in species_files:

            species_file_path = os.path.join(dir_path, species_file_name)

            species_data = None
            with open(species_file_path, "r", encoding="utf-8") as f:
                species_data = json.load(f)

                species_name = species_data["target"] 
                species_name = species_name.split(":")[1] # The species target name in the file is in the format: "cobblemon:species_name", so we split it by the colon and take the second half

                # Edit the ability data (abities is a list of strings, sometimes having an "h:" in front of the name if it's the hidden ability)
                forms = species_data["forms"]
                for form in forms:

                    if species_name not in ability_data:
                        print(f"WARNING: Species {species_name} not found in ability list. Skipping...")
                        continue

                    new_ability_name = ability_data[species_name]

                    if "abilities" not in form:
                        print(f"WARNING: Species {species_name} has no abilities list. Skipping...")
                        continue

                    abilities = form["abilities"]
                    new_abilities = []
                    for ability in abilities:
                        
                        if ":" in ability:
                            ability_prefix, ability = ability.split(":")    
                            new_abilities.append(ability_prefix + ':' + new_ability_name)                        
                        else:
                            new_abilities.append(new_ability_name)

                    form["abilities"] = new_abilities

            output_file_path = os.path.join(dst_gen_dir_path, species_file_name)
            with open(output_file_path, "w", encoding="utf-8") as f:
                json.dump(species_data, f, indent=4)

    return



def main():
    program_name = sys.argv[0]

    usage_string = f"Usage: {program_name} <source directory> <target directory> <target data>"

    num_args = len(sys.argv)
    if num_args != 4:
        print(usage_string)

    _, src_dir, dst_dir, data_path = sys.argv # Gets the required information from the program arguments

    ability_data = get_target_abilities(data_path)

    write_augmented_species_data(src_dir, dst_dir, ability_data)


if __name__ == "__main__":
    main()

# My args:
# /home/dawson/Desktop/Navas-Investigation/zamega-fabric-1.4.6/data/zamegas/species_additions/ /home/dawson/Desktop/Navas-Investigation/output/ /home/dawson/Desktop/Navas-Investigation/ZAOUMegaList.txt