import json
import os

from scipy.spatial.transform import Rotation
import xml.etree.ElementTree as ET
import numpy as np

def parse_xml_to_json(xml_file_path, json_file_path, include_unit_id=False):
    # Parse the XML file
    tree = ET.parse(xml_file_path)
    root = tree.getroot()

    # Create a list to store the parsed data
    world_data = []

    # Keep track of unique instances of the "path" attribute
    unique_paths = set()

    # Iterate through the "unit_data" elements and extract the desired attributes
    for unit_data_elem in root.findall('.//unit_data'):
        position_str = unit_data_elem.get('position')
        rotation_quaternion_str = unit_data_elem.get('rotation')
        full_name = unit_data_elem.get('name')
        unit_id = unit_data_elem.get('unit_id') if include_unit_id else None

        # Replace commas with dots in the position and quaternion strings
        position_str = position_str.replace(',', '.')
        rotation_quaternion_str = rotation_quaternion_str.replace(',', '.')

        # Extract name and path
        unit_description = full_name.split('/')[-1]
        path = full_name

        # Track unique instances of the "path" attribute
        unique_paths.add(path)

        # Convert position and quaternion to vectors
        try:
            position_vector = list(map(float, position_str.split()))

            # Convert right-handed to left-handed coordinates
            position_vector[0] = -position_vector[0]

            quaternion = np.array(list(map(float, rotation_quaternion_str.split())))
            
            rotation = Rotation.from_quat(quaternion)
            euler_angles = rotation.as_euler('zyx')
            
            # Convert radians to degrees and clamp to 3 decimal places
            rotation_vector = np.degrees(euler_angles).round(3).tolist()

            rotation_vector[0] = - rotation_vector[0]

            # Convert Euler angles to Rotator format
            rotator = {
                'yaw': rotation_vector[0],
                'pitch': rotation_vector[1],
                'roll': rotation_vector[2]
            }

            # Generate Unreal object reference
            object_reference = f"/Script/Engine.StaticMesh'/Game/{path}/{unit_description}.{unit_description}'"

            # Use 'name' instead of 'unit_id'
            unit_data = {
                'position': {
                    'x': position_vector[0],
                    'y': position_vector[1],
                    'z': position_vector[2]
                },
                'rotation_quaternion': rotation_quaternion_str,
                'rotation_vector': rotator,
                'unit_description': unit_description,
                'path': object_reference,
            }

            if include_unit_id:
                unit_data['name'] = unit_id
            else:
                unit_data['name'] = unit_description

            world_data.append(unit_data)

        except ValueError:
            print(f"Skipping invalid data: {position_str}, {rotation_quaternion_str}")

    # Print the unique instances of the "path" attribute
    unique_dlc_units = set()
    for unique_path in unique_paths:
        dlc = unique_path.split('/')[1]
        unique_dlc_units.add(dlc)

    print(f"Units DLCs in this map: {', '.join(unique_dlc_units)}. Make sure to get them from the game")

    # Convert the list to JSON and write it to the output file
    with open(json_file_path, 'w') as json_file:
        json.dump(world_data, json_file, indent=2)

    print(f'Conversion complete. JSON data written to {json_file_path}')



def process_folder(folder_path):
    # Ensure the folder exists, create it if it doesn't
    full_folder_path = os.path.abspath(folder_path)
    if not os.path.exists(full_folder_path):
        os.makedirs(full_folder_path)
        print(f'Created folder: {full_folder_path}')

    # List all files and folders in the current directory
    files_and_folders = os.listdir(full_folder_path)

    for item in files_and_folders:
        item_path = os.path.join(full_folder_path, item)

        # If it's a directory, recursively process it
        if os.path.isdir(item_path):
            print(f'Processing folder: {item_path}')
            process_folder(item_path)

        # If it's a file with a .continent extension, process it
        elif os.path.isfile(item_path) and item.lower().endswith('.continent'):
            print(f'Processing file: {item_path}')

            # Get the directory of the script using __file__
            script_directory = os.path.abspath(os.path.dirname(__file__))

            json_file_path = os.path.abspath(os.path.join(script_directory, 'WorldFiles', 'UnrealJSON', f'{os.path.splitext(item)[0]}.json'))
            parse_xml_to_json(item_path, json_file_path, include_unit_id=True)
            print(f'Created JSON file: {json_file_path}')


# Get The working Dir
root_folder = r'\WorldFiles\DieselWorldCONTINENT'
dir_path = os.path.dirname(os.path.realpath(__file__))
folder_path = dir_path + root_folder
# Start processing
print(f'Starting processing from: {folder_path}')
process_folder(folder_path)