#!/usr/bin/env python3
import fileinput
import glob
import os
import xml.etree.ElementTree as ET

class SedLike:
  def sed(path, changes):
    for line in fileinput.input(path, inplace=True):
      for change in changes:
        line = line.replace(change[0], change[1])
      print(line, end="")

class TemplateFixer:
  def __init__(self, vfs_root):
    self.template_folder = os.path.join(vfs_root, 'simulation', 'templates')

  def fix_template(self, template_path):
    tree = ET.parse(template_path)
    root = tree.getroot()
    production_queue = root.find('ProductionQueue')
    if production_queue == None:
      return False

    technologies = production_queue.find('Technologies')
    tech_cost_multiplier = production_queue.find('TechCostMultiplier')
    if technologies != None or tech_cost_multiplier != None:
      researcher = ET.Element('Researcher')
      if (technologies != None):
        researcher.append(technologies)
        production_queue.remove(technologies)
      if (tech_cost_multiplier != None):
        researcher.append(tech_cost_multiplier)
        production_queue.remove(tech_cost_multiplier)
      researcher[:] = sorted(researcher, key=lambda x: x.tag)
      root.append(researcher)

    entities = production_queue.find('Entities')
    batch_time_modifier = production_queue.find('BatchTimeModifier')
    if entities != None or batch_time_modifier != None:
      trainer = ET.Element('Trainer')
      if (entities != None):
        trainer.append(entities)
        production_queue.remove(entities)
      if (batch_time_modifier != None):
        trainer.append(batch_time_modifier)
        production_queue.remove(batch_time_modifier)
      trainer[:] = sorted(trainer, key=lambda x: x.tag)
      root.append(trainer)

    if production_queue.get('disable') != None:
      for element in ["Researcher", "Trainer"]:
        functionality = ET.Element(element)
        functionality.set('disable', "")
        root.append(functionality)
    else:
      root.remove(production_queue)

    root[:] = sorted(root, key=lambda x: x.tag)
    ET.indent(tree)

    tree.write(template_path, xml_declaration=True, encoding='utf-8')
    return True

  def fix_style(self, template_path):
    self.changes = [
      [' />', '/>'],
      ["version='1.0'", 'version="1.0"'],
      ["'utf-8'", '"utf-8"']
    ]
    SedLike.sed(template_path, self.changes)
    with open(template_path, 'a') as file:
      file.write('\n')

  def run(self):
    for template in glob.iglob(self.template_folder + '/**/*.xml', recursive=True):
      if self.fix_template(template):
        self.fix_style(template)


class ModifiersFixer:
  def __init__(self, vfs_root):
    self.data_folder = os.path.join(vfs_root, 'simulation', 'data')
    self.changes = [
      ["ProductionQueue/Batch", "Trainer/Batch"],
      ["ProductionQueue/Ent", "Trainer/Ent"],
      ["ProductionQueue/Train", "Trainer/Train"],
      ["ProductionQueue/Tech", "Researcher/Tech"],
    ]

  def run(self):
    for modification in glob.iglob(self.data_folder + '/**/*.json', recursive=True):
      SedLike.sed(modification, self.changes)


if __name__ == '__main__':
  script_dir = os.path.dirname(os.path.realpath(__file__))
  template_fixer = TemplateFixer(script_dir)
  template_fixer.run()

  modifier_fixer = ModifiersFixer(script_dir)
  modifier_fixer.run()