# This script will take the ModelRepo.xml file and rename the models based on the trademark

import argparse
import json
import xml.etree.ElementTree as ET

parser = argparse.ArgumentParser(description='Rename the models in an XML file')
parser.add_argument('--original', metavar='xml_file', type=str, help='The XML file to be parsed', required=True)
parser.add_argument('--out', metavar='output_file', type=str, default='ModelRepo_generated.xml', help='The output file', const='ModelRepo_generated.xml', nargs='?')
args = parser.parse_args()

tree = ET.parse(args.original)
root = tree.getroot()

def load_override_json(path):
    with open(path) as f:
        data = json.load(f)
    return data

def get_actual_model(model):
    name = model.get('name')
    trademark = model.get('tm')
    model_id = model.get('id')

    # Check if the model has a manual override
    if model_id in override.keys():
        print('Overriding model: ' + name + ' with ' + override[model_id])
        return override[model_id]

    # Remove "Based on " from the name
    new_name = trademark.replace('Based on ', '')
    # remove "®" from the name
    new_name = new_name.replace('®', '')
    # Change out the word "loaded" for "|"
    new_name = new_name.replace('loaded with', '|')
    # remove the word "with" from "loaded with"
    new_name = new_name.replace('with', '|')
    # remove the word "tubes" from the name
    new_name = new_name.replace('tubes', '')
    new_name = new_name.replace('Tubes', '')

    # Check if the name contains a Mono or Stereo suffix
    if name.endswith('(M)'):
        new_name += ' (M)'
    elif name.endswith('(S)'):
        new_name += ' (S)'

    return new_name


print('Renaming models in file: ' + args.original)
override = load_override_json('manual_override.json')["ids"]
print(override)

# Itterate over the categories inside the XML file
for category in root.findall('Category'):
    # Itterate over the models inside the category
    for model in category.findall('Model'):
        # First check if the model has a trademark
        trademark = model.get('tm')
        if trademark != "" and trademark != None:
            new_name = get_actual_model(model)
            model.set('name', new_name)

print('Saving file to: ' + args.out)
# Save the XML file
tree.write(args.out, encoding='UTF-8', xml_declaration=True)