import os
import sys

import xml.etree.ElementTree as ET


def output_xml_tree(tree, path):
    with open(path, 'wt') as handle:
        handle.write('<?xml version="1.0" encoding="utf-8"?>\n')
        def output_xml_node(node, handle, depth):
            indent = '  ' * depth
            attributes = ''
            for attribute_name in node.attrib.keys():
                attributes += ' {}="{}"'.format(attribute_name, node.attrib[attribute_name])
            if node.getchildren():
                assert (node.text is None) or (not node.text.strip())
                handle.write('{}<{}{}>\n'.format(indent, node.tag, attributes))
                for child in node.getchildren():
                    output_xml_node(child, handle, depth + 1)
                handle.write('{}</{}>\n'.format(indent, node.tag))
            else:
                if node.text is None:
                    handle.write('{}<{}{}/>\n'.format(indent, node.tag, attributes))
                else:
                    handle.write('{}<{}{}>{}</{}>\n'.format(indent, node.tag, attributes, node.text, node.tag))
        output_xml_node(tree.getroot(), handle, 0)

def remove_terrain_materials(path):
    with open(path, 'rt') as handle:
        xml = handle.read()

    terrain_base = 'terrain_base.xml' in xml
    terrain_norm = 'terrain_norm.xml' in xml
    terrain_spec = 'terrain_spec.xml' in xml
    terrain_triplanar = 'terrain_triplanar.xml' in xml

    should_save = terrain_base or terrain_norm or terrain_spec or terrain_triplanar

    assert xml.count('<material>') == 1

    tree = ET.parse(path)
    root = tree.getroot()
    for element in root.getchildren():
        if len(element.getchildren()) > 0:
            assert element.tag == 'textures'
            for child in element.getchildren():
                assert child.tag == 'texture'
            if terrain_base or terrain_spec or terrain_triplanar:
                texture_element = ET.SubElement(element, 'texture')
                texture_element.set('name', 'normTex')
                texture_element.set('file', 'types/default_norm.png')
            if terrain_base or terrain_norm or terrain_triplanar:
                texture_element = ET.SubElement(element, 'texture')
                texture_element.set('name', 'specTex')
                texture_element.set('file', 'types/blackness.dds')
        elif element.tag == 'material':
            if terrain_base or terrain_norm or terrain_spec:
                element.text = 'terrain_norm_spec.xml'
            elif terrain_triplanar:
                element.text = 'terrain_triplanar_norm_spec.xml'

    if should_save:
        output_xml_tree(tree, path)

def get_materials():
    return {
        'aura.xml': {'norm': False, 'spec': True, 'replacement': 'aura_norm_spec.xml'},
        'aura_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'aura_norm_spec.xml'},
        'basic_glow.xml': {'norm': False, 'spec': True, 'replacement': 'basic_glow_norm_spec.xml'},
        'basic_glow_norm.xml': {'norm': True, 'spec': True, 'replacement': 'basic_glow_norm_spec.xml'},
        'basic_glow_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_glow_norm_spec.xml'},
        'basic_glow_wind.xml': {'norm': False, 'spec': True, 'replacement': 'basic_glow_wind_norm_spec.xml'},
        'basic_glow_wind_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_glow_wind_norm_spec.xml'},
        'basic_specmap.xml': {'norm': False, 'spec': True, 'replacement': 'no_trans_norm_spec.xml'},
        'basic_trans.xml': {'norm': False, 'spec': False, 'replacement': 'basic_trans_norm_spec.xml'},
        'basic_trans_ao.xml': {'norm': False, 'spec': False, 'replacement': 'basic_trans_ao_norm_spec.xml'},
        'basic_trans_ao_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_trans_ao_norm_spec.xml'},
        'basic_trans_ao_parallax_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_trans_ao_parallax_spec.xml'},
        'basic_trans_ao_spec.xml': {'norm': False, 'spec': True, 'replacement': 'basic_trans_ao_norm_spec.xml'},
        'basic_trans_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_trans_norm_spec.xml'},
        'basic_trans_parallax_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_trans_parallax_spec.xml'},
        'basic_trans_spec.xml': {'norm': False, 'spec': True, 'replacement': 'basic_trans_norm_spec.xml'},
        'basic_trans_wind.xml': {'norm': False, 'spec': False, 'replacement': 'basic_trans_wind_norm_spec.xml'},
        'basic_trans_wind_grain.xml': {'norm': False, 'spec': False, 'replacement': 'basic_trans_wind_grain_norm_spec.xml'},
        'basic_trans_wind_grain_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_trans_wind_grain_norm_spec.xml'},
        'basic_trans_wind_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'basic_trans_wind_norm_spec.xml'},
        'default.xml': {'norm': False, 'spec': False, 'replacement': 'no_trans_norm_spec.xml'},
        'no_trans_ao.xml': {'norm': False, 'spec': False, 'replacement': 'no_trans_ao_norm_spec.xml'},
        'no_trans_ao_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'no_trans_ao_norm_spec.xml'},
        'no_trans_ao_parallax_spec.xml': {'norm': True, 'spec': True, 'replacement': 'no_trans_ao_parallax_spec.xml'},
        'no_trans_ao_spec.xml': {'norm': False, 'spec': True, 'replacement': 'no_trans_ao_norm_spec.xml'},
        'no_trans_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'no_trans_norm_spec.xml'},
        'no_trans_parallax_ao.xml': {'norm': True, 'spec': False, 'replacement': 'no_trans_ao_parallax_spec.xml'},
        'no_trans_parallax_spec.xml': {'norm': True, 'spec': True, 'replacement': 'no_trans_parallax_spec.xml'},
        'no_trans_spec.xml': {'norm': False, 'spec': True, 'replacement': 'no_trans_norm_spec.xml'},
        'objectcolor.xml': {'norm': False, 'spec': False, 'replacement': 'objectcolor_norm_spec.xml'},
        'objectcolor_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'objectcolor_norm_spec.xml'},
        'objectcolor_specmap.xml': {'norm': False, 'spec': True, 'replacement': 'objectcolor_norm_spec.xml'},
        'player_trans.xml': {'norm': False, 'spec': False, 'replacement': 'player_trans_norm_spec.xml'},
        'player_trans_ao.xml': {'norm': False, 'spec': False, 'replacement': 'player_trans_ao_norm_spec.xml'},
        'player_trans_ao_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'player_trans_ao_norm_spec.xml'},
        'player_trans_ao_parallax.xml': {'norm': True, 'spec': False, 'replacement': 'player_trans_ao_parallax_spec.xml'},
        'player_trans_ao_parallax_spec.xml': {'norm': True, 'spec': True, 'replacement': 'player_trans_ao_parallax_spec.xml'},
        'player_trans_ao_spec.xml': {'norm': True, 'spec': True, 'replacement': 'player_trans_ao_spec.xml'},
        'player_trans_norm.xml': {'norm': True, 'spec': False, 'replacement': 'player_trans_norm_spec.xml'},
        'player_trans_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'player_trans_norm_spec.xml'},
        'player_trans_norm_spec_helmet.xml': {'norm': True, 'spec': True, 'replacement': 'player_trans_norm_spec_helmet.xml'},
        'player_trans_parallax.xml': {'norm': True, 'spec': False, 'replacement': 'player_trans_parallax_spec.xml'},
        'player_trans_parallax_spec.xml': {'norm': True, 'spec': True, 'replacement': 'player_trans_parallax_spec.xml'},
        'player_trans_parallax_spec_helmet.xml': {'norm': True, 'spec': True, 'replacement': 'player_trans_parallax_spec_helmet.xml'},
        'player_trans_spec.xml': {'norm': False, 'spec': True, 'replacement': 'player_trans_norm_spec.xml'},
        'player_trans_spec_helmet.xml': {'norm': False, 'spec': True, 'replacement': 'player_trans_norm_spec_helmet.xml'},
        'player_water.xml': {'norm': False, 'spec': False, 'replacement': 'player_water.xml'},
        'rock_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'rock_norm_spec.xml'},
        'rock_norm_spec_ao.xml': {'norm': True, 'spec': True, 'replacement': 'rock_norm_spec_ao.xml'},
        'rock_normstrong_spec.xml': {'norm': True, 'spec': True, 'replacement': 'rock_normstrong_spec.xml'},
        'rock_normstrong_spec_ao.xml': {'norm': True, 'spec': True, 'replacement': 'rock_normstrong_spec_ao.xml'},
        'terrain_base.xml': {'norm': False, 'spec': False, 'replacement': 'terrain_norm_spec.xml'},
        'terrain_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'terrain_norm_spec.xml'},
        'terrain_normstrong_spec.xml': {'norm': True, 'spec': True, 'replacement': 'terrain_normstrong_spec.xml'},
        'terrain_normweak_spec.xml': {'norm': True, 'spec': True, 'replacement': 'terrain_normweak_spec.xml'},
        'terrain_triplanar.xml': {'norm': False, 'spec': False, 'replacement': 'terrain_triplanar_norm_spec.xml'},
        'terrain_triplanar_norm_spec.xml': {'norm': True, 'spec': True, 'replacement': 'terrain_triplanar_norm_spec.xml'},
        'trans_wind.xml': {'norm': False, 'spec': False, 'replacement': 'basic_trans_wind_norm_spec.xml'},
        'waterfall.xml': {'norm': False, 'spec': False, 'replacement': 'waterfall.xml'},
    }

def get_material(material):
    return get_materials()[material]

def remove_actor_materials(path):
    with open(path, 'rt') as handle:
        xml = handle.read()

    should_save = False

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

    material_name = None

    if '<material>' in xml:
        assert xml.count('<material>') == 1

        if root.tag == 'actor':
            for child in root.getchildren():
                if child.tag == 'material':
                    material_name = child.text.strip()
                    child.text = get_material(material_name)['replacement']
        else:
            assert root.tag == 'qualitylevels'
            for root_child in root.getchildren():
                assert root_child.tag == 'inline' or root_child.tag == 'actor'
                for child in root_child.getchildren():
                    if child.tag == 'material':
                        material_name = child.text.strip()
                        child.text = get_material(material_name)['replacement']
    else:
        assert root.tag == 'actor'
        material_name = 'default.xml'
        material_element = ET.SubElement(root, 'material')
        material_element.text = get_material(material_name)['replacement']

    assert material_name is not None

    material = get_material(material_name)
    if material_name == material['replacement']:
        return
    replacement_material = get_material(material['replacement'])
    assert (not material['norm']) or (not material['spec'])
    assert replacement_material['norm'] and replacement_material['spec']

    if material['norm'] != ('normTex' in xml) or material['spec'] != ('specTex' in xml) or '<!--' in xml:
        sys.stdout.write(
            'Actor "{}" should be handled manually. It should use "{}" instead of "{}".\n'.format(
                path, material['replacement'], material_name))
        return
    should_save = True

    # If there are no textures it means the material isn't really useful. We'll with them later.
    if xml.count('<textures>') > 0:
        def add_default_norm_element(textures_element, index=None):
            if not index:
                norm_element = ET.SubElement(textures_element, 'texture')
            else:
                norm_element = ET.Element('texture')
            norm_element.set('file', 'default_norm.png')
            norm_element.set('name', 'normTex')
            if index:
                textures_element.insert(index, norm_element)
        def add_default_spec_element(textures_element, index=None):
            if not index:
                spec_element = ET.SubElement(textures_element, 'texture')
            else:
                spec_element = ET.Element('texture')
            spec_element.set('file', 'null_black.dds')
            spec_element.set('name', 'specTex')
            if index:
                textures_element.insert(index, spec_element)
        def add_missing_textures(element, needs_norm, needs_spec):
            for child in element.getchildren():
                if child.tag == 'textures':
                    expected_len = 1
                    if not needs_norm:
                        expected_len += 1
                    if not needs_spec:
                        expected_len += 1
                    if expected_len > 1:
                        only_base = True
                        for texture_child in child.getchildren():
                            if texture_child.get('name') == 'normTex' or texture_child.get('name') == 'specTex':
                                only_base = False
                        if only_base:
                            continue
                    if needs_norm and needs_spec:
                        add_default_norm_element(child)
                        add_default_spec_element(child)
                    elif needs_norm:
                        add_default_norm_element(child)
                    elif needs_spec:
                        add_default_spec_element(child)
                else:
                    add_missing_textures(child, needs_norm, needs_spec)
        if not material['norm'] and not material['spec']:
            if xml.count('<textures>') > 1:
                group_element = ET.Element('group')
                variant_element = ET.SubElement(group_element, 'variant')
                textures_element = ET.SubElement(variant_element, 'textures')
                add_default_norm_element(textures_element)
                add_default_spec_element(textures_element)
                if root.tag == 'actor':
                    last_index = len(root.getchildren()) - 1
                    assert root.getchildren()[last_index].tag == 'material'
                    root.insert(last_index, group_element)
                else:
                    assert root.tag == 'qualitylevels'
                    for root_child in root.getchildren():
                        assert root_child.tag == 'inline' or root_child.tag == 'actor'
                        last_index = len(root_child.getchildren()) - 1
                        assert root_child.getchildren()[last_index].tag == 'material'
                        root_child.insert(last_index, group_element)
            else:
                add_missing_textures(root, True, True)
        else:
            add_missing_textures(root, not material['norm'], not material['spec'])

    if should_save:
        sys.stdout.write('Actor "{}" is fixed.\n'.format(path))
        output_xml_tree(tree, path)

def remove_terrain_materials_recursively(path):
    for file_name in os.listdir(path):
        file_path = os.path.join(path, file_name)
        if os.path.isfile(file_path):
            name, ext = os.path.splitext(file_name)
            if ext.lower() == '.xml' and name.lower() != 'terrains':
                remove_terrain_materials(file_path)
        elif os.path.isdir(file_path):
            remove_terrain_materials_recursively(file_path)

def remove_actor_materials_recursively(path):
    for file_name in os.listdir(path):
        file_path = os.path.join(path, file_name)
        if os.path.isfile(file_path):
            name, ext = os.path.splitext(file_name)
            if ext.lower() == '.xml':
                remove_actor_materials(file_path)
        elif os.path.isdir(file_path):
            remove_actor_materials_recursively(file_path)

if __name__ == '__main__':
    # paths should be a list of paths to mods folders
    paths = sys.argv[1:]
    for path in paths:
        remove_terrain_materials_recursively(os.path.join(path, 'art', 'terrains'))
        remove_actor_materials_recursively(os.path.join(path, 'art', 'actors'))