#!/usr/bin/env python3
"""Prepare office.sdf for Gazebo Sim + dartsim.

1) Merge poses from Gazebo Classic's <state> into each <model><pose>. gz-sim
   ignores <state>, so without this step furniture stays at wrong coordinates.
2) Remove the <state> block.
3) Strip <collision> blocks that use triangle <mesh> (dartsim mesh crash).

Usage:
  python3 gen_office_gz_dartsim.py office.sdf office_gz_dartsim.sdf
"""
import sys
import xml.etree.ElementTree as ET


def _geom_has_mesh(geom: ET.Element) -> bool:
    return geom.find('mesh') is not None


def strip_mesh_collisions(root: ET.Element) -> int:
    removed = 0
    for link in root.iter('link'):
        for coll in list(link.findall('collision')):
            geom = coll.find('geometry')
            if geom is not None and _geom_has_mesh(geom):
                link.remove(coll)
                removed += 1
    return removed


def apply_state_poses(world: ET.Element) -> tuple[int, int]:
    """Copy model poses from <state> into world models; drop <state>."""
    state = None
    for child in list(world):
        if child.tag == 'state':
            state = child
            break
    if state is None:
        return 0, 0

    poses: dict[str, str] = {}
    for model in state.findall('model'):
        name = model.get('name')
        pose_el = model.find('pose')
        if name and pose_el is not None and (pose_el.text is not None):
            poses[name] = pose_el.text.strip()

    updated = 0
    for model in world.findall('model'):
        name = model.get('name')
        if not name or name not in poses:
            continue
        new_pose = poses[name]
        pe = model.find('pose')
        if pe is not None:
            pe.text = new_pose
        else:
            el = ET.Element('pose')
            el.text = new_pose
            model.insert(0, el)
        updated += 1

    world.remove(state)
    return updated, len(poses)


def main() -> int:
    if len(sys.argv) != 3:
        print(__doc__.strip(), file=sys.stderr)
        return 2
    inp, outp = sys.argv[1], sys.argv[2]
    tree = ET.parse(inp)
    root = tree.getroot()
    if root.tag != 'sdf':
        print('Expected root <sdf>', file=sys.stderr)
        return 1
    world = root.find('world')
    if world is None:
        print('Expected <world> inside <sdf>', file=sys.stderr)
        return 1
    n_pose, n_state_models = apply_state_poses(world)
    n_mesh = strip_mesh_collisions(root)
    ET.indent(tree, space='  ')
    tree.write(outp, encoding='unicode', xml_declaration=True)
    print(
        f'Wrote {outp}: state poses merged for {n_pose} models '
        f'({n_state_models} in <state>), removed {n_mesh} mesh collision blocks',
        file=sys.stderr,
    )
    return 0


if __name__ == '__main__':
    raise SystemExit(main())
