from bs4 import BeautifulSoup
import os
import xml.etree.ElementTree as ET
from tqdm import tqdm
from tqdm import trange

def indent_and_correct(xml_file):
    indented_xml = ""
    with open(file=xml_file,mode="r",encoding="utf-8") as file:
        indented_xml = BeautifulSoup(file, "xml")
    with open(file=xml_file,mode="w",encoding="utf-8") as file:
        file.write(str(indented_xml))

    tree = ET.parse(xml_file)
    root = tree.getroot()

    attack_types = [
        "Crush", "Hack", "Pierce", "Fire", "Forest", "Siege", "Light", "Shadow", "Spirit", "Water", "Poison"
    ]

    for component in root:
        # Fix for attack which now has a damage node
        if(component.tag == "Attack"):
            for damageTypes in component:
                hasName = False
                for componentProp in damageTypes:
                    if(componentProp.tag =="AttackName"):
                        hasName = True
                        break

                if not hasName:
                    attacknameel = ET.Element("AttackName")
                    attacknameel.text = "TBD."
                    damageTypes.append(attacknameel)

                hasdamage = False
                for componentProp in damageTypes:
                    if(componentProp.tag =="Damage"):
                        hasdamage = True
                        break

                if not hasdamage:
                    propDamage = ET.Element("Damage")
                    for componentProp in damageTypes:
                        if componentProp.tag not in attack_types:
                            continue
                        if root.attrib["parent"] is not None and componentProp.text != "0.0" and componentProp.text != "0":
                            propDamage.append(componentProp)

                    for componentProp in reversed(damageTypes):
                        if componentProp.tag not in attack_types:
                            continue

                        damageTypes.remove(componentProp)

                    if len(propDamage) > 0:
                        damageTypes.append(propDamage)

                if damageTypes.tag == "Ranged":
                    projectileEl = damageTypes.find("Projectile")
                    projectileSpeedEl = damageTypes.find("ProjectileSpeed")
                    projectileSpreadEl = damageTypes.find("Spread")
                    if projectileEl is None:
                        projectileEl = ET.Element("Projectile")

                    if projectileSpeedEl is not None:
                        speedEL = ET.Element("Speed")
                        speedEL.text = projectileSpeedEl.text
                        projectileEl.append(speedEL)
                        damageTypes.remove(projectileSpeedEl)

                    if projectileSpreadEl is not None:
                        spreadEL = ET.Element("Spread")
                        spreadEL.text = projectileSpreadEl.text
                        projectileEl.append(spreadEL)
                        damageTypes.remove(projectileSpreadEl)


        # Cost population bonus was moved to a population component
        # In the future the population count might actually be moved there.
        if(component.tag == "Cost"):
            for componentProp in reversed(component):
                 if(componentProp.tag == "PopulationBonus"):
                    bonusEL = ET.Element("Bonus")
                    bonusEL.text = componentProp.text
                    populationEL = ET.Element("Population")
                    populationEL.append(bonusEL)
                    root.append(populationEL)
                    component.remove(componentProp)

        # Fix for UnitMotion, which no longer has a run component but rather a run multiplier
        if(component.tag == "UnitMotion"):
            for componentProp in reversed(component):
                if componentProp.tag == "Run":
                    walkspeedel = component.find("WalkSpeed")
                    runspeedel = componentProp.find("Speed")
                    if(walkspeedel is not None and runspeedel is not None):
                        multiplierel = ET.Element("RunMultiplier")
                        multiplierel.text = str(float(runspeedel.text) / float(walkspeedel.text))[0:6]
                        component.remove(componentProp)
                        component.append(multiplierel)


        # Fix for Armour which now has a damage node
        if(component.tag == "Armour"):
            component.tag = "Resistance"
            propDamage = ET.Element("Damage")
            for componentProp in component:

                if componentProp.tag == "Foundation" or componentProp.tag == "KnockbackResistance":
                    continue

                if root.attrib["parent"] is not None and componentProp.text != "0.0" and componentProp.text != "0":
                    propDamage.append(componentProp)
                if componentProp.text is not None:
                    componentProp.text = str(float(componentProp.text))

            for componentProp in reversed(component):
                if componentProp.tag == "Foundation" or componentProp.tag == "KnockbackResistance":
                    continue

                component.remove(componentProp)


            component.append(propDamage)
    def indent(elem, level=0, more_sibs=False):
        i = "\n"
        if level:
            i += (level-1) * '  '
        num_kids = len(elem)
        if num_kids:
            if not elem.text or not elem.text.strip():
                elem.text = i + "  "
                if level:
                    elem.text += '  '
            count = 0
            for kid in elem:
                indent(kid, level+1, count < num_kids - 1)
                count += 1
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
                if more_sibs:
                    elem.tail += '  '
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = i
                if more_sibs:
                    elem.tail += '  '
    def sort(root):
        # sort the first layer
        root[:] = sorted(root, key=lambda child: (child.tag,child.get('name')))

        # sort the second layer
        for c in root:
            c[:] = sorted(c, key=lambda child: (child.tag,child.get('name')))
            for cp in c:
                cp[:] = sorted(cp, key=lambda child: (child.tag,child.get('name')))
                for scp in cp:
                    scp[:] = sorted(scp, key=lambda child: (child.tag,child.get('name')))
    sort(root)
    indent(root)
    tree.write(xml_file)


listOfFiles = os.walk("templates/")
flatFileList = []
for root, directories, filenames in listOfFiles:
    for filename in filenames:
        if ".xml" in filename:
            flatFileList.append(os.path.join(root,filename))

for filename in tqdm(flatFileList):
    indent_and_correct(filename)