"""
convert_urdf.py — Compile the RX1 humanoid-torso URDF into a clean MuJoCo MJCF.

The shipped URDF uses ROS `package://rx1_description/...` mesh URIs that MuJoCo
cannot resolve.  We rewrite those to plain relative paths, inject a <mujoco>
compiler block, let MuJoCo parse the URDF, then save the resulting MJCF and
hand-augment it (fixed base, floor, head camera, position actuators) in a
separate step (build_scene.py).

Run:
    python convert_urdf.py
Produces:
    rx1_mujoco/rx1_compiled.xml   (raw MuJoCo dump of the URDF)
"""
from __future__ import annotations

import pathlib
import re
import tempfile

import mujoco

HERE = pathlib.Path(__file__).resolve().parent
URDF_SRC = HERE.parent / "humanoid_torso" / "rx1.urdf"
MESH_DIR = HERE / "meshes_mjcf"   # decimated, MuJoCo-safe copies (see prepare_meshes.py)
OUT_XML = HERE / "rx1_compiled.xml"


def main() -> None:
    urdf_text = URDF_SRC.read_text()

    # Strip the ROS package prefix down to the bare path inside the mesh tree,
    # matching the flattened layout of meshes_mjcf/ (base_link.stl, head/..., etc.)
    urdf_text = urdf_text.replace("package://rx1_description/meshes/", "")

    # Inject a <mujoco> compiler block telling MuJoCo where the meshes live and
    # to auto-fix any degenerate inertias coming from the URDF.
    compiler = (
        f'  <mujoco>\n'
        f'    <compiler meshdir="{MESH_DIR}" balanceinertia="true" '
        f'discardvisual="false" fusestatic="false"/>\n'
        f'  </mujoco>\n'
    )
    urdf_text = re.sub(r"(<robot[^>]*>)",
                       lambda m: m.group(0) + "\n" + compiler,
                       urdf_text, count=1)

    # MuJoCo resolves relative mesh paths against the XML file's directory, so we
    # write the patched URDF to a temp file (meshdir is absolute anyway).
    with tempfile.NamedTemporaryFile("w", suffix=".urdf", delete=False) as f:
        f.write(urdf_text)
        tmp_path = f.name

    model = mujoco.MjModel.from_xml_path(tmp_path)
    print(f"[convert] URDF compiled: nq={model.nq} nu={model.nu} "
          f"nbody={model.nbody} njnt={model.njnt} nmesh={model.nmesh}")

    mujoco.mj_saveLastXML(str(OUT_XML), model)
    print(f"[convert] wrote {OUT_XML}")


if __name__ == "__main__":
    main()