"""URDF XML helpers shared by sysid rollout and export modules."""

from __future__ import annotations

import xml.etree.ElementTree as ET


def set_mass(inertial: ET.Element, mass: float | None) -> None:
    if mass is None:
        return
    mass_el = inertial.find("mass")
    if mass_el is not None:
        mass_el.set("value", str(mass))


def set_com(
    inertial: ET.Element,
    x: float | None,
    y: float | None,
    z: float | None,
) -> None:
    origin = inertial.find("origin")
    if origin is None:
        return
    xyz = origin.get("xyz", "0 0 0").split()
    if x is not None:
        xyz[0] = str(x)
    if y is not None:
        xyz[1] = str(y)
    if z is not None:
        xyz[2] = str(z)
    origin.set("xyz", " ".join(xyz))


def set_inertia(
    inertial: ET.Element,
    ixx: float | None = None,
    iyy: float | None = None,
    izz: float | None = None,
    ixy: float | None = None,
    iyz: float | None = None,
    ixz: float | None = None,
) -> None:
    ine = inertial.find("inertia")
    if ine is None:
        return
    for attr, val in [
        ("ixx", ixx), ("iyy", iyy), ("izz", izz),
        ("ixy", ixy), ("iyz", iyz), ("ixz", ixz),
    ]:
        if val is not None:
            ine.set(attr, str(val))


def patch_link_inertials(
    root: ET.Element,
    params: dict[str, float],
) -> None:
    """Patch arm and pendulum inertial parameters in a URDF ElementTree root."""
    for link in root.iter("link"):
        link_name = link.get("name", "")
        inertial = link.find("inertial")
        if inertial is None:
            continue

        if link_name == "arm":
            set_mass(inertial, params.get("arm_mass"))
            set_com(inertial, params.get("arm_com_x"),
                     params.get("arm_com_y"), params.get("arm_com_z"))

        elif link_name == "pendulum":
            set_mass(inertial, params.get("pendulum_mass"))
            set_com(inertial, params.get("pendulum_com_x"),
                     params.get("pendulum_com_y"), params.get("pendulum_com_z"))
            set_inertia(inertial,
                         ixx=params.get("pendulum_ixx"),
                         iyy=params.get("pendulum_iyy"),
                         izz=params.get("pendulum_izz"),
                         ixy=params.get("pendulum_ixy"))
