import bpy
import os

export_path = r"D:\Unity\little_sophia_app\Assets\robots\LittleSophia"
mesh_folder = os.path.join(export_path, "meshes")
urdf_file = os.path.join(export_path, "LittleSophia.urdf")

if not os.path.exists(mesh_folder):
    os.makedirs(mesh_folder)

urdf_lines = []
urdf_lines.append('<robot name="LittleSophia">')

for obj in bpy.data.objects:
    if obj.type == 'MESH':
        # Export mesh with origin at zero (like before)
        original_location = obj.location.copy()
        bpy.ops.object.select_all(action='DESELECT')
        obj.select_set(True)
        bpy.context.view_layer.objects.active = obj
        obj.location = (0,0,0)

        mesh_path = f"meshes/{obj.name}.obj"
        bpy.ops.wm.obj_export(
            filepath=os.path.join(mesh_folder, f"{obj.name}.obj"),
            export_selected_objects=True,
            export_uv=True,
            export_normals=True,
            export_materials=True
        )
        obj.location = original_location

        # Write link entry
        urdf_lines.append(f'  <link name="{obj.name}">')
        urdf_lines.append('    <visual>')
        urdf_lines.append('      <geometry>')
        urdf_lines.append(f'        <mesh filename="{mesh_path}"/>')
        urdf_lines.append('      </geometry>')
        urdf_lines.append('    </visual>')
        urdf_lines.append('  </link>')

        # If object has a parent, write joint
        if obj.parent:
            parent = obj.parent
            rel_loc = obj.location - parent.location
            urdf_lines.append(f'  <joint name="{parent.name}_to_{obj.name}" type="fixed">')
            urdf_lines.append(f'    <parent link="{parent.name}"/>')
            urdf_lines.append(f'    <child link="{obj.name}"/>')
            urdf_lines.append(f'    <origin xyz="{rel_loc.x} {rel_loc.y} {rel_loc.z}" rpy="0 0 0"/>')
            urdf_lines.append('  </joint>')

urdf_lines.append('</robot>')

with open(urdf_file, "w") as f:
    f.write("\n".join(urdf_lines))

print("✅ URDF file generated:", urdf_file)
