"""URDF 处理工具模块

提供 URDF 文件处理功能。
"""

import xml.etree.ElementTree as ET
from pathlib import Path

from utils import make_mesh_path_absolute


def prepare_isaac_urdf(source_urdf_path, output_urdf_path):
    """准备 Isaac Sim 可用的 URDF 文件

    处理网格路径和关节原点设置。
    """
    tree = ET.parse(source_urdf_path)
    root = tree.getroot()
    source_dir = source_urdf_path.parent

    local_mesh_links = {
        "left_steering_hinge",
        "right_steering_hinge",
        "left_wheel",
        "right_wheel",
        "left_rear_wheel",
        "right_rear_wheel",
        "camera",
        "laser",
    }

    for link in root.findall("link"):
        link_name = link.get("name", "")
        for section_name in ("visual", "collision"):
            section = link.find(section_name)
            if section is None:
                continue
            mesh = section.find("geometry/mesh")
            if mesh is not None and mesh.get("filename"):
                mesh.set("filename", make_mesh_path_absolute(mesh.get("filename"), source_dir))
            if link_name in local_mesh_links:
                origin = section.find("origin")
                if origin is None:
                    origin = ET.SubElement(section, "origin")
                origin.set("xyz", "0 0 0")
                origin.set("rpy", "0 0 0")

    output_urdf_path.parent.mkdir(parents=True, exist_ok=True)
    tree.write(output_urdf_path, encoding="utf-8", xml_declaration=True)
    return output_urdf_path
