import os, shutil, xml.etree.ElementTree as ET, mujoco
D = "/data/cameron/repos/smith300_para_stuff"
Y = "/data/cameron/para/robot/yam/raiden_fork/third_party/exo_redo/robot_models/i2rt_yam_fidex/assets"
OUT = f"{D}/mj_stuff"; os.makedirs(OUT, exist_ok=True)
search = [D, Y, f"{D}/assets"]
def bundle(xin, xout):
    t = ET.parse(xin); r = t.getroot()
    c = r.find("compiler") or ET.SubElement(r, "compiler"); c.set("meshdir", "."); c.set("texturedir", ".")
    n = 0
    for tag in ("mesh", "texture"):
        for e in r.iter(tag):
            f = e.get("file")
            if not f: continue
            b = os.path.basename(f)
            for s in search:
                if os.path.exists(f"{s}/{b}"):
                    if not os.path.exists(f"{OUT}/{b}"): shutil.copy(f"{s}/{b}", OUT)
                    n += 1; break
            e.set("file", b)
    t.write(xout); return n
for x in ["two_servo.xml", "two_servo_fidex.xml", "arm_vs_yam.xml"]:
    k = bundle(f"{D}/{x}", f"{OUT}/{x}")
    m = mujoco.MjModel.from_xml_path(f"{OUT}/{x}")
    print(f"{x}: {k} assets bundled, reloads OK (nbody {m.nbody})")