import trimesh
from lxml import etree
import os
import numpy as np

# Define texture path
texture_path = "palette.png"

# Parse the XML file
xml_path = "scene.xml"
tree = etree.parse(xml_path)
root = tree.getroot()

# Create an empty scene
scene = trimesh.Scene()

length = len(root.findall("mesh"))
print(f"Found {length} meshes in the XML file")
iterator = 0

# Load the texture
if not os.path.exists(texture_path):
    raise FileNotFoundError(f"Texture file not found: {texture_path}")

# Iterate over each mesh in the XML
for mesh_data in root.findall("mesh"):
    iterator += 1
    print(f"Processing mesh {iterator} of {length}")

    obj_file = mesh_data.get("file")
    position = [float(x) for x in mesh_data.get("pos").split()]
    rotation = [float(x) for x in mesh_data.get("rot").split()]

    # Load the OBJ file
    if not os.path.exists(obj_file):
        print(f"File not found: {obj_file}")
        continue

    mesh = trimesh.load_mesh(obj_file)

    # Apply translation
    translation_matrix = trimesh.transformations.translation_matrix(position)

    # Apply rotation (convert degrees to radians for trimesh)
    rotation_matrix = trimesh.transformations.euler_matrix(
        np.radians(rotation[0]),  # Rotation around X-axis
        np.radians(rotation[1]),  # Rotation around Y-axis
        np.radians(rotation[2]),  # Rotation around Z-axis
        axes="sxyz"
    )

    # Combine transformations
    transform_matrix = np.dot(translation_matrix, rotation_matrix)
    mesh.apply_transform(transform_matrix)

    # Apply texture to the mesh
    if hasattr(mesh, "visual"):
        mesh.visual = trimesh.visual.TextureVisuals(image=texture_path)

    # Add the mesh to the scene
    scene.add_geometry(mesh)

    if iterator > 100:
        break

# Export the merged scene
output_path = "output.obj"
scene.export(output_path, include_texture=True)

print(f"Scene exported to {output_path}")