"""Scene mutation via the MuJoCo ``MjSpec`` AST.

This module used to contain ~980 lines of XML-round-trip machinery (tmpdir +
``mj_saveLastXML`` + ElementTree parse + name-mangling + regex path patching).
All of that is replaced by ``spec.recompile(model, data)`` which:

* preserves joint state on unchanged joints automatically,
* initializes new joints to body ``pos``/``quat`` (removing the need to
  delete keyframes on freejoint insertion),
* namespaces robot bodies/joints/geoms/actuators/sensors via ``spec.attach()``
  without us walking the tree manually.

Public API:

* :func:`inject_robot_into_scene` - ``spec.attach(robot_spec, prefix=...)``.
* :func:`inject_object_into_scene` - ``SpecBuilder.add_object(spec, obj)`` + recompile.
* :func:`inject_camera_into_scene` - ``SpecBuilder.add_camera(spec, cam)`` + recompile.
* :func:`install_compiled_model` - install a compiled model/data pair as the
  live scene state; the single point every model swap goes through.
* :func:`eject_body_from_scene` - ``SpecBuilder.remove_body(spec, name)`` + recompile.
* :func:`reposition_body_in_scene` - edit a body's spec ``pos``/``quat`` + recompile.
* :func:`eject_robot_from_scene` - walk the spec, delete everything namespaced
  under ``{robot_name}/``, then recompile.
* :func:`refresh_body_inertial_from_geometry` - re-derive a body's mass /
  center of mass / inertia after one of its geoms was resized at runtime.

Every function takes a ``SimWorld`` whose ``_backend_state["spec"]`` holds the
live ``MjSpec``. They return ``True`` on success, ``False`` on failure (matching
the legacy API) so call sites in :mod:`simulation` don't need to change.
"""

from __future__ import annotations

import difflib
import logging
from collections.abc import Callable, Mapping
from typing import Any

from strands_robots.simulation.models import SimCamera, SimObject, SimRobot, SimWorld
from strands_robots.simulation.mujoco.backend import _ensure_mujoco, filter_mujoco_attach_noise, mj_name_to_id
from strands_robots.simulation.mujoco.spec_builder import SpecBuilder
from strands_robots.utils import coerce_rgba, finite_vector_error, pose_vector_error

logger = logging.getLogger(__name__)


def _get_spec(world: SimWorld) -> Any | None:
    """Fetch the live MjSpec from ``world._backend_state``.

    Callers MUST have run ``_compile_world`` at least once before any scene
    mutation - without a spec we can't recompile. Returns ``None`` if missing
    so callers can return a clean error dict rather than crashing mid-op.
    """
    return world._backend_state.get("spec")


def _sync_cached_xml(world: SimWorld, spec: Any) -> None:
    """Refresh the legacy XML cache in ``world._backend_state["xml"]`` from ``spec``.

    Some readers (the ``load_scene`` + ``add_robot`` round-trip) consume the
    cached XML string rather than the live ``MjSpec``, so it must be refreshed
    after every recompile. ``spec.to_xml()`` can fail on specs MuJoCo cannot
    serialise; that is never fatal to the mutation (the live model/spec are
    already updated), so we leave the previous cache in place - but we always
    log the reason at debug so a silently stale cache stays diagnosable rather
    than being swallowed.
    """
    try:
        with filter_mujoco_attach_noise():
            world._backend_state["xml"] = spec.to_xml()
    except Exception as xml_err:
        logger.debug("spec.to_xml() failed; cached XML left stale: %s", xml_err)


def _snapshot_spec(spec: Any, *, context: str) -> Any | None:
    """Deep-copy ``spec`` so a refused mutation can be rolled back losslessly.

    A spec mutation that is only validated by the recompile it precedes needs a
    way back. Rebuilding the spec from the registered objects / cameras / robots
    is not it: the live spec can carry mutations that exist ONLY on the spec and
    are absent from that registry - weld equalities from ``attach_bodies``,
    actuators from :func:`actuate_robot_in_scene`, bodies from
    :func:`patch_scene_mjcf`, whole scenes from :func:`replace_scene_mjcf`. A
    rebuild silently drops every one of them, which is why
    :meth:`Simulation.remove_robot` refuses outright while an attachment is
    live rather than rebuilding over it.

    ``MjSpec.copy`` is used rather than a ``spec.to_xml()`` round-trip because
    the round-trip is not faithful for a scene holding attached robots: the
    emitted MJCF loses the asset search paths its mesh references were resolved
    against and re-declares the attached model's keyframes, so the restored spec
    no longer compiles ("Error opening file 'link2.stl'", "repeated name
    'panda/home' in key") - a rollback that leaves the scene as broken as the
    orphan it removed.

    Args:
        spec: live scene spec to snapshot.
        context: short description of the caller, used in log messages.

    Returns:
        An independent copy of ``spec``, or ``None`` when the copy failed (the
        reason is logged). A caller that cannot snapshot must refuse its
        mutation rather than proceed with no way back.
    """
    try:
        return spec.copy()
    except (ValueError, RuntimeError) as e:
        logger.error("%s: cannot snapshot the scene spec: %s", context, e)
        return None


def actuator_joint_id(model: Any, act_id: int, mj: Any) -> int:
    """Return the joint id actuator ``act_id`` drives, or ``-1`` if it drives none.

    ``actuator_trnid[act_id, 0]`` holds a JOINT id only when the actuator's
    transmission is ``mjTRN_JOINT`` or ``mjTRN_JOINTINPARENT``. For the other
    transmissions it holds a tendon, site, slider-crank or body id instead --
    separate id spaces that each start at 0, so a raw comparison against joint
    ids matches an unrelated entity that merely shares a number. A fixed tendon
    coupling a gripper's fingers is the standard MJCF idiom for that (the
    Menagerie Panda hand and the Robotiq 2F-85 both use one), so the collision
    is reachable with stock assets rather than only with exotic models.

    Reading the transmission through this function is what keeps "which joint
    does this actuator drive" a single rule, rather than one that each call site
    re-derives and can omit.

    Args:
        model: The compiled ``MjModel``.
        act_id: Actuator index in ``range(model.nu)``.
        mj: The ``mujoco`` module.

    Returns:
        The driven joint id, or ``-1`` when the transmission is not a joint.
    """
    joint_trn = {int(mj.mjtTrn.mjTRN_JOINT)}
    if hasattr(mj.mjtTrn, "mjTRN_JOINTINPARENT"):
        joint_trn.add(int(mj.mjtTrn.mjTRN_JOINTINPARENT))
    if int(model.actuator_trntype[act_id]) not in joint_trn:
        return -1
    return int(model.actuator_trnid[act_id, 0])


def robot_owned_actuator_ids(model: Any, robot: SimRobot, mj: Any) -> list[int]:
    """Return the actuator ids ``robot`` owns, in model order.

    Ownership is the union of two rules, because neither alone covers every
    actuator a robot can carry:

    * **Namespace.** ``spec.attach(other, prefix=...)`` prefixes every actuator
      name it merges in, so a prefixed name settles ownership whatever the
      transmission drives -- the only rule that reaches a tendon, site or body
      transmission.
    * **Driven joint.** :func:`actuate_robot_in_scene` names the position
      actuators it injects ``"<robot>_act_<joint>"``, which carries no namespace
      prefix, so those are recognized by the joint they drive.

    The second rule resolves the joint through :func:`actuator_joint_id`, so an
    actuator whose transmission is not a joint is never assigned by an id-space
    collision. Ungated, a tendon-driven gripper is claimed by whichever robot
    owns the joint whose id equals the tendon's and is missing from the robot
    that actually carries it -- the gripper then has no owner to operate it, and
    the other robot advertises an actuator that moves a different machine.

    Args:
        model: The compiled ``MjModel``.
        robot: The robot to scope. Reads its ``namespace`` and its already
            resolved ``joint_ids``.
        mj: The ``mujoco`` module.

    Returns:
        Owned actuator ids, ascending. Ascending model order is the contract:
        :meth:`~strands_robots.simulation.mujoco.simulation.MuJoCoSimEngine.robot_action_keys`
        and the recorded dataset action columns are ordered by actuator index,
        so the sequence is stable and only membership is corrected here.
    """
    pfx = robot.namespace or ""
    joint_ids = set(robot.joint_ids)
    owned: list[int] = []
    for act_id in range(int(model.nu)):
        name = mj.mj_id2name(model, mj.mjtObj.mjOBJ_ACTUATOR, act_id) or ""
        if pfx and name.startswith(pfx):
            owned.append(act_id)
        elif actuator_joint_id(model, act_id, mj) in joint_ids:
            owned.append(act_id)
    return owned


def install_compiled_model(world: SimWorld, model: Any, data: Any) -> None:
    """Install a compiled ``model``/``data`` pair as ``world``'s live scene state.

    Every path that swaps the live model goes through here, because the swap has
    a consequence that is easy to leave behind: ``world._recompile_generation``
    is the only part of the ``save_state`` / ``load_state`` fingerprint that
    changes when the new model happens to keep the previous ``nq``/``nv``/``na``/
    ``nu``. A site that rebinds ``world._model`` without bumping it lets a
    checkpoint taken against the previous model be applied to this one - the
    state vector is the right length, so it is written index by index into
    whatever those indices now mean.

    Concentrating the rebind here makes that impossible to forget: a new swap
    site cannot install a model without also invalidating the checkpoints taken
    against the old one.

    Args:
        world: The scene whose live model/data are replaced.
        model: The freshly compiled ``MjModel``.
        data: An ``MjData`` matching ``model``.
    """
    world._model = model
    world._data = data
    # Any swap invalidates every outstanding checkpoint, including one whose
    # state vector still has the right length: only this counter distinguishes
    # a same-shape model from the one the checkpoint was taken against.
    world._recompile_generation += 1


def _recompile_preserving_state(world: SimWorld, spec: Any, *, raise_on_refusal: bool = False) -> bool:
    """Recompile ``spec`` in place, replacing ``world._model`` and ``_data``.

    Uses ``spec.recompile(model, data)`` which auto-preserves qpos/qvel for
    existing joints and initializes new joints to their body's pos/quat. No
    manual state-copy loop is required.

    That transfer is POSITIONAL, not by name: the new buffers receive the old
    values for the indices both models share, and the entries past the old size
    are whatever the fresh allocation happened to contain. For ``qpos`` (filled
    from ``qpos0``), ``qvel`` and ``act`` the compiler defines them; for
    ``ctrl`` it does not, so a recompile that grows ``nu`` -- adding a robot,
    adding actuators to one -- leaves the new setpoints undefined. That is not
    a harmless nonsense number: MuJoCo's ``mj_checkCtrl`` disables actuation for
    the WHOLE model on any step where a single ``ctrl`` entry is non-finite, so
    one uninitialized entry can silently release every held pose in the scene,
    only on the runs where the leftover memory happens to be NaN. This function
    therefore defines the new tail as zero -- the value a reset writes, and the
    only setpoint a caller who has not commanded the new actuators can mean --
    before the forward pass below reads it.

    Also re-discovers per-robot joint and actuator IDs (they may have shifted
    as new bodies were inserted earlier in the body tree). Returns True on
    success, False on compile failure (logged).

    Args:
        world: The scene whose model/data are replaced on success.
        spec: The ``MjSpec`` to compile.
        raise_on_refusal: Re-raise the compiler's exception instead of folding
            it into a ``False`` return. A refusal message names what the
            compiler could not honor, and a ``bool`` cannot carry it: a caller
            that reports the reason to a user needs it, whereas a caller that
            only rolls back does not. Off by default so existing callers keep
            the ``bool`` contract.
    """
    mj = _ensure_mujoco()
    # Sizes of the buffers being handed to the compiler: everything at or past
    # these indices in the new data is an entry the old model had no value for.
    old_nu = int(world._model.nu) if world._model is not None else 0
    old_na = int(world._model.na) if world._model is not None else 0
    try:
        with filter_mujoco_attach_noise():
            new_model, new_data = spec.recompile(world._model, world._data)
    except (ValueError, RuntimeError) as e:
        logger.error("spec.recompile failed: %s", e)
        if raise_on_refusal:
            raise
        return False

    install_compiled_model(world, new_model, new_data)
    # Define the control entries the positional transfer left untouched (see the
    # docstring). ``act`` is measurably zero-filled today, but it comes out of
    # the same allocation as ``ctrl`` and carries an actuator's activation -- its
    # effective command -- so it is defined here too rather than by luck.
    if new_model.nu > old_nu:
        new_data.ctrl[old_nu:] = 0.0
    if new_model.na > old_na:
        new_data.act[old_na:] = 0.0
    # Forward pass so newly-injected bodies have valid xpos/xquat and any
    # camera xforms are populated. Without this, the next render() call
    # after add_object / add_robot / add_camera returns a 100% black frame
    # because the MjData arrays still hold their initialization zeros.
    mj.mj_forward(new_model, new_data)

    # Keep the cached XML in sync with the spec for legacy readers (e.g.
    # load_scene + add_robot round-trip).
    _sync_cached_xml(world, spec)

    # Re-discover per-robot IDs. Names inside MuJoCo are namespaced under
    # robot.namespace (e.g. "arm1/shoulder_pan") when robots were attached
    # via SpecBuilder.attach_robot; fall back to the raw name otherwise.
    for robot in world.robots.values():
        pfx = robot.namespace or ""
        robot.joint_ids = []
        for jnt_name in robot.joint_names:
            jid = -1
            if pfx:
                jid = mj_name_to_id(new_model, mj.mjtObj.mjOBJ_JOINT, pfx + jnt_name)
            if jid < 0:
                jid = mj_name_to_id(new_model, mj.mjtObj.mjOBJ_JOINT, jnt_name)
            if jid >= 0:
                robot.joint_ids.append(jid)
        robot.actuator_ids = robot_owned_actuator_ids(new_model, robot, mj)
        # Single-robot fallback. Ownership above is settled by namespace or by
        # driven joint; a lone robot whose actuators are neither prefixed nor
        # joint-driven (a tendon or site transmission in a scene loaded whole,
        # so no attach prefix) matches on neither, and in a one-robot scene
        # every actuator is unambiguously that robot's.
        if not robot.actuator_ids and len(world.robots) == 1:
            robot.actuator_ids = list(range(new_model.nu))

    return True


# Persist


_NO_SPEC_REASON = "the scene has no live spec, so the change cannot be recorded durably"


def _spec_element_by_id(elements: Any, entity_id: int, kind: str) -> tuple[Any | None, str | None]:
    """Return the spec element a compiled entity id was built from.

    The compiler emits one model entity per spec element in declaration order and
    records the resulting index on the element, so a compiled id indexes the
    spec's element list directly. That mapping is what lets an UNNAMED element be
    addressed at all: most geoms in a robot scene carry no name, so resolving by
    name would silently cover almost none of them.

    The recorded ``id`` is verified rather than assumed. Writing a property onto
    the wrong entity is a worse outcome than not recording it, so a spec that no
    longer agrees with the compiled model is reported instead of written to.

    Args:
        elements: The spec's element list (``spec.geoms`` / ``spec.bodies``).
        entity_id: Compiled index of the entity, as resolved by the caller.
        kind: Entity name used in the reason text (``"geom"`` / ``"body"``).

    Returns:
        ``(element, None)`` when the element was located, otherwise
        ``(None, reason)`` naming why it was not.
    """
    count = len(elements)
    if entity_id < 0 or entity_id >= count:
        return None, f"{kind} id {entity_id} is outside the scene spec's {count} {kind}(s)"
    element = elements[entity_id]
    if element.id != entity_id:
        return None, (
            f"the scene spec no longer agrees with the compiled model: the {kind} at index "
            f"{entity_id} reports id {element.id}"
        )
    return element, None


def persist_geom_properties(
    world: SimWorld,
    geom_id: int,
    *,
    color: list[float] | None = None,
    friction: list[float] | None = None,
    size: list[float] | None = None,
) -> str | None:
    """Record a runtime geom property write in the spec the model is compiled from.

    ``world._model`` is DERIVED state: every scene mutation recompiles the spec
    over it (see :func:`_recompile_preserving_state`), so a value written only
    into the model is discarded by the next ``add_object`` / ``add_camera`` /
    ``add_robot`` call and the geom reverts to whatever it was compiled with -
    after the setter already reported the new value. Writing the same value into
    the spec is what makes that reported result durable.

    Nothing is written when a reason is returned, so a caller can refuse before it
    touches the model and keep the two representations in step.

    Args:
        world: The scene holding the live spec.
        geom_id: Compiled geom index, already resolved by the caller.
        color: RGBA components, already validated.
        friction: The three friction coefficients, already validated.
        size: Half-extents, already validated. Only the components the caller
            supplied are written, matching the model write: the unused tail of
            the spec's 3-wide row keeps its declared value.

    Returns:
        ``None`` once the value is recorded, otherwise the reason it could not be.
    """
    spec = _get_spec(world)
    if spec is None:
        return _NO_SPEC_REASON
    spec_geom, reason = _spec_element_by_id(spec.geoms, geom_id, "geom")
    if spec_geom is None:
        return reason

    if color is not None:
        spec_geom.rgba = list(color)
    if friction is not None:
        spec_geom.friction[:] = friction
    if size is not None:
        spec_geom.size[: len(size)] = size
    return None


def refresh_body_inertial_from_geometry(world: SimWorld, geom_id: int) -> str | None:
    """Re-derive the inertial row of the body owning ``geom_id`` from its geometry.

    A body that does not declare an ``<inertial>`` takes its mass, center of mass
    and inertia tensor from the shapes it owns: the compiler integrates them over
    the geoms once, at compile time. ``model.body_mass`` / ``body_ipos`` /
    ``body_iquat`` / ``body_inertia`` are therefore DERIVED from ``geom_size``, and
    no ``mj_forward`` / ``mj_step`` recomputes them. Resizing a geom in the model
    alone leaves the body describing the shape it used to have, so it resists
    rotation as the old geometry did while colliding as the new one - and for a
    geom whose mass comes from a density, it also keeps the old mass and the old
    balance point.

    The values are read from a compile of a *copy* of the live spec, which already
    carries the new geometry (see :func:`persist_geom_properties`). That makes the
    refreshed row equal, by construction, to the one the next scene recompile will
    produce - so the same resize no longer means two different things depending on
    whether an unrelated ``add_object`` happens afterwards. Deriving it from
    MuJoCo's own integrator rather than from per-primitive formulas is also what
    makes a multi-geom body and a shifted center of mass come out right, with no
    shape-by-shape special cases to keep correct.

    The live model is not swapped and the scene's own ``mjData`` is never passed to
    MuJoCo, so entity ids, joint state and the recompile generation are untouched -
    the refresh reads geometry and writes constants, and a resize therefore leaves
    the scene exactly where it was. The cost is one spec compile plus one scratch
    ``mjData``, paid only for a resize of a geom whose body derives its inertia
    from geometry.

    Args:
        world: The scene holding the live spec and compiled model.
        geom_id: Compiled geom index whose size has already been recorded in the
            spec, as resolved by the caller.

    Returns:
        ``None`` once the row is refreshed - including when the owning body
        declares its own ``<inertial>``, which takes nothing from geometry and so
        needs no refresh - otherwise the reason it could not be, leaving both the
        model and the spec untouched so the caller can restore them.
    """
    mj = _ensure_mujoco()
    model = world._model
    if model is None:
        return "the scene has no compiled model whose inertia could be re-derived"

    spec = _get_spec(world)
    if spec is None:
        return _NO_SPEC_REASON

    body_id = int(model.geom_bodyid[geom_id])
    spec_body, reason = _spec_element_by_id(spec.bodies, body_id, "body")
    if spec_body is None:
        return reason
    if spec_body.explicitinertial:
        # The compiler ignores a body's geoms when the body declares its own
        # <inertial>, so that row does not describe the geometry and a resize
        # cannot make it stale. Returning early also keeps resizing a robot
        # link's collision geom free, since those links declare their inertials.
        return None

    try:
        with filter_mujoco_attach_noise():
            candidate = spec.copy().compile()
    except (ValueError, RuntimeError) as e:
        return f"the resized geometry does not compile, so the body's inertia cannot be re-derived: {e}"

    # Writing a row read from a model with a different body ordering would move
    # one body's inertia onto another, which is a worse outcome than leaving the
    # row stale and saying so.
    if candidate.nbody != model.nbody:
        return (
            "the scene spec no longer agrees with the compiled model: it compiles to "
            f"{candidate.nbody} bodies against the model's {model.nbody}"
        )
    if mj.mj_id2name(candidate, mj.mjtObj.mjOBJ_BODY, body_id) != mj.mj_id2name(model, mj.mjtObj.mjOBJ_BODY, body_id):
        return (
            "the scene spec no longer agrees with the compiled model: body "
            f"{body_id} is named differently in a fresh compile"
        )

    model.body_mass[body_id] = candidate.body_mass[body_id]
    model.body_ipos[body_id] = candidate.body_ipos[body_id]
    model.body_iquat[body_id] = candidate.body_iquat[body_id]
    model.body_inertia[body_id] = candidate.body_inertia[body_id]
    # body_subtreemass and the invweight / M0 reference constants the constraint
    # solver scales with are themselves derived from the inertial rows and are not
    # refreshed by a step. mj_setConst recomputes them, which is what makes the
    # whole inertial state of the live model equal to that of a fresh compile.
    #
    # It is handed a scratch mjData, never the scene's own. Those constants are by
    # definition evaluated at the model's reference configuration, so mj_setConst
    # writes qpos0 into whatever data it is given and does not put back what was
    # there. Passing the live data would rewind every joint and body pose to its
    # declared value while leaving qvel as it was, so a resize issued after any
    # stepping would teleport the scene into a state it never occupied - the
    # order-dependent silent damage this refresh exists to remove.
    mj.mj_setConst(model, mj.MjData(model))
    return None


def persist_body_mass(world: SimWorld, body_id: int, *, mass_ratio: float) -> str | None:
    """Record a runtime body mass change in the spec the model is compiled from.

    Recording the change as a scale is what keeps the two representations equal:
    ``set_body_properties`` documents a mass change as a uniform density change at
    fixed geometry, and both mass and inertia are linear in density, so applying
    one ratio reproduces exactly the inertial the setter reported.

    A body's compiled inertial comes from one of two places, and only the one in
    force is writable. A body that declares an explicit ``<inertial>`` carries its
    own mass and inertia (``explicitinertial``); every other body has both
    integrated from its geoms, and there assigning ``mass`` on the body element is
    silently ignored by the compiler - the geoms' mass or density is what has to
    move.

    Args:
        world: The scene holding the live spec.
        body_id: Compiled body index, already resolved by the caller.
        mass_ratio: The new mass divided by the compiled mass. Finite and ``> 0``,
            which the caller guarantees by refusing a body with no mass to scale.

    Returns:
        ``None`` once the change is recorded, otherwise the reason it could not be.
    """
    spec = _get_spec(world)
    if spec is None:
        return _NO_SPEC_REASON
    spec_body, reason = _spec_element_by_id(spec.bodies, body_id, "body")
    if spec_body is None:
        return reason

    if spec_body.explicitinertial:
        spec_body.mass *= mass_ratio
        # A body states its inertia as either the principal diagonal or the six
        # unique components of the full tensor; whichever form is not in use is
        # all zeros, so scaling both is exact and needs no branch.
        spec_body.inertia[:] = [value * mass_ratio for value in spec_body.inertia]
        spec_body.fullinertia[:] = [value * mass_ratio for value in spec_body.fullinertia]
        return None

    spec_geoms = list(spec_body.geoms)
    if not spec_geoms:
        return (
            "the body declares no explicit inertial and owns no geom, so it holds "
            "nothing whose mass the change could scale"
        )
    for spec_geom in spec_geoms:
        # A geom states either an explicit mass or a density, and the compiler
        # uses the mass only when it is set (an unset mass reads as nan, which
        # fails every comparison). Scaling whichever one is in force scales that
        # geom's contribution, so the body total scales by the same ratio.
        if spec_geom.mass > 0:
            spec_geom.mass *= mass_ratio
        else:
            spec_geom.density *= mass_ratio
    return None


def persist_world_option(
    world: SimWorld,
    *,
    gravity: list[float] | None = None,
    timestep: float | None = None,
) -> str | None:
    """Record a runtime physics-option write in the spec the model is compiled from.

    ``model.opt`` is compiled from ``spec.option``, so a gravity or timestep
    written only into the model is restored to the scene's declared value by the
    next recompile - putting a lunar-gravity world back to 9.81 m/s^2 on the next
    ``add_object``.

    Args:
        world: The scene holding the live spec.
        gravity: The three gravity components, already validated.
        timestep: The integration step in seconds, already validated.

    Returns:
        ``None`` once the value is recorded, otherwise the reason it could not be.
    """
    spec = _get_spec(world)
    if spec is None:
        return _NO_SPEC_REASON
    if gravity is not None:
        spec.option.gravity[:] = gravity
    if timestep is not None:
        spec.option.timestep = timestep
    return None


# Inject


def inject_robot_into_scene(
    world: SimWorld,
    robot: SimRobot,
    robot_xml_path: str,
) -> bool:
    """Attach a robot to the scene via ``spec.attach(other, prefix=..., frame=...)``.

    MuJoCo handles name prefixing (bodies, joints, geoms, actuators, sensors,
    sites), asset deduplication (meshes, textures, materials), and default-
    class namespacing. No manual tree-walking required.

    Registers the robot's source joint names on ``robot.joint_names`` so
    downstream observation/policy code can resolve them via
    ``{robot.namespace}{joint_name}``.

    ``spec.attach`` mutates the scene spec before the recompile that validates
    the result, so a refused recompile - e.g. the robot model references a mesh
    file that cannot be opened - leaves the robot's whole namespaced subtree in
    the live spec. Every later scene mutation then recompiles that same broken
    spec and fails too, so one unloadable robot bricked the world: subsequent
    ``add_object`` / ``add_camera`` / ``add_robot`` calls all reported "spec
    recompile refused" with nothing wrong with them, and each failed retry left
    another orphan subtree behind. The attach is therefore rolled back out
    before returning ``False``, mirroring :func:`inject_object_into_scene` and
    :func:`inject_camera_into_scene`, so a refused robot costs exactly the add
    that was refused.

    The rollback reinstalls a snapshot taken before the attach
    (:func:`_snapshot_spec`) rather than rebuilding the scene from ``world``: a
    rebuild would drop the spec-only mutations the scene may already hold (weld
    equalities, added actuators, agent-authored bodies), turning a correctly
    refused add into corruption of a scene that was healthy before it. The
    refused ``spec.recompile`` leaves ``world._model`` and ``_data`` untouched,
    so putting the spec back is the whole rollback - no state restore, forward
    pass or ID rediscovery is needed.
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("inject_robot: no spec or model in world")
        return False

    # Snapshot BEFORE mutating. Without a way back the add is refused here,
    # while the scene is still exactly as it was found.
    context = f"inject_robot {robot.name!r}"
    backup_spec = _snapshot_spec(spec, context=context)
    if backup_spec is None:
        return False

    try:
        with filter_mujoco_attach_noise():
            joint_names = SpecBuilder.attach_robot(spec, robot, robot_xml_path)
        robot.joint_names = joint_names
    except (ValueError, RuntimeError, OSError) as e:
        # attach_robot can insert its worldbody frame before the call that
        # raised. That leftover compiles, so it never broke the scene, but the
        # snapshot is already in hand - restoring it costs nothing and puts
        # every failure path on one rule: the spec is left as it was found.
        logger.error("Robot attach failed for '%s': %s", robot.name, e)
        world._backend_state["spec"] = backup_spec
        return False

    if _recompile_preserving_state(world, spec):
        return True

    # The attach landed in the spec but the model it produced was refused, so
    # the robot's whole namespaced subtree is sitting in the live spec while the
    # caller is about to report a failed add. Put the pre-attach spec back.
    world._backend_state["spec"] = backup_spec
    return False


def inject_object_into_scene(world: SimWorld, obj: SimObject) -> bool:
    """Add a ``SimObject`` to the scene and recompile in place.

    A ``shape="mesh"`` object references a mesh asset that must be registered
    on the spec before the geom that names it can compile. The full-scene
    ``SpecBuilder.build`` registers those assets in its own pass, but the
    incremental path (``SpecBuilder.add_object``) does not, so this function
    registers the mesh (``spec.add_mesh(name=f"mesh_{obj.name}", ...)``) itself
    before adding the body. Without this, ``add_object(shape="mesh")`` at
    runtime always failed to recompile even for a valid mesh file.

    ``SpecBuilder.add_object`` mutates the spec (adds the body + geom) before
    the recompile that validates it. If that recompile is refused - e.g. the
    mesh file cannot be loaded - the just-added body AND its mesh asset are
    deleted again before returning ``False`` so the spec stays compilable.
    Without the rollback the orphan lingers and every later scene mutation,
    including a corrected retry under the same name, keeps failing to recompile
    (``repeated name`` collisions), bricking the whole scene after one bad add.

    The same rollback applies when ``SpecBuilder.add_object`` itself raises
    part-way through, which it can do *after* inserting the body (the geom's
    type lookup rejects an unsupported shape; the mass write rejects a
    non-numeric value). That error is then re-raised rather than folded into a
    ``False`` return, so the caller can report the actual reason - a swallowed
    ``ValueError`` left the caller with nothing but "spec recompile refused"
    while the actionable message went to the log.
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("inject_object: no spec or model in world")
        return False

    try:
        # Meshes need their asset registered before the geom references it.
        # build() registers meshes in a separate pass, so add_object does not;
        # the incremental path must register it here.
        if obj.shape == "mesh" and obj.mesh_path:
            spec.add_mesh(name=f"mesh_{obj.name}", file=obj.mesh_path)
        SpecBuilder.add_object(spec, obj)
    except (ValueError, RuntimeError):
        # add_object is atomic over its own body mutation: a raise there (an
        # unsupported shape, or a name that collides with an existing scene
        # body) rolls the half-built body back out itself, so only the mesh
        # asset registered here still needs undoing. Removing it keeps the spec
        # compilable, then the error propagates: the caller turns it into a
        # structured result, and the reason - e.g. the exact unsupported shape
        # and the supported list - is what a caller needs instead of a generic
        # recompile refusal.
        SpecBuilder.remove_mesh(spec, f"mesh_{obj.name}")
        raise

    # Roll the just-added body (and any mesh asset) back out so the spec
    # returns to its last good, compilable state (a worldbody body delete is
    # safe - the attach/delete segfault only affects spec.attach() child specs).
    #
    # Ask for the compiler's own reason rather than a bare False. Because the
    # object's mass is declared on its geom, MuJoCo integrates the inertia from
    # the shape, so its "mass and inertia of moving bodies must be larger than
    # mjMINVAL" floor is shape-dependent: a mass above mjMINVAL can still
    # integrate to an inertia below it on a small geom. add_object's numeric
    # pre-check cannot express that floor without duplicating the compiler's
    # per-shape integration, so the residual case has to arrive as the reason
    # the compiler gives. Folded into a False it became "spec recompile
    # refused." with the actionable text left in the log - the same dead end
    # the unsupported-shape path was fixed to stop producing.
    try:
        recompiled = _recompile_preserving_state(world, spec, raise_on_refusal=True)
    except (ValueError, RuntimeError):
        SpecBuilder.remove_body(spec, obj.name)
        SpecBuilder.remove_mesh(spec, f"mesh_{obj.name}")
        raise
    if not recompiled:
        SpecBuilder.remove_body(spec, obj.name)
        SpecBuilder.remove_mesh(spec, f"mesh_{obj.name}")
        return False
    return True


def inject_camera_into_scene(world: SimWorld, cam: SimCamera) -> bool:
    """Add a camera to the scene and recompile in place.

    Mirrors :func:`inject_object_into_scene`: ``SpecBuilder.add_camera`` mutates
    the spec before the validating recompile, so a refused recompile rolls the
    just-added camera back out to keep the spec compilable for later edits.
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("inject_camera: no spec or model in world")
        return False

    try:
        SpecBuilder.add_camera(spec, cam)
    except (ValueError, RuntimeError) as e:
        logger.error("Camera add failed for '%s': %s", cam.name, e)
        return False

    if not _recompile_preserving_state(world, spec):
        SpecBuilder.remove_camera(spec, cam.name)
        return False
    return True


# Eject


def eject_body_from_scene(world: SimWorld, body_name: str) -> bool:
    """Remove a body (by short name) and recompile."""
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("eject_body: no spec or model in world")
        return False

    if not SpecBuilder.remove_body(spec, body_name):
        logger.warning("Body '%s' not found in spec - nothing ejected", body_name)
        # Matching legacy behaviour: return True so scene state stays consistent
        # (caller has already popped the Python-side dict entry).
        return True

    # Objects added at runtime register a mesh asset named f"mesh_{name}".
    # Delete it too so the name is fully reusable and unused assets do not
    # accumulate across remove/re-add cycles (safe no-op for primitives).
    SpecBuilder.remove_mesh(spec, f"mesh_{body_name}")

    return _recompile_preserving_state(world, spec)


def reposition_body_in_scene(
    world: SimWorld,
    body_name: str,
    position: list[float] | None = None,
    orientation: list[float] | None = None,
) -> bool:
    """Reposition a body (by short name) by editing its spec pose and recompiling.

    Used for STATIC objects, which have no freejoint and therefore cannot be
    moved through ``data.qpos`` at runtime - MuJoCo welds a static body to the
    worldbody with no DOF. Editing the spec body ``pos``/``quat`` and
    recompiling (preserving other joints' state) is the only way to move a
    welded fixture, mirroring :func:`inject_object_into_scene` /
    :func:`eject_body_from_scene`.

    ``position`` / ``orientation`` are applied only when provided (a ``None``
    leaves that component untouched). Returns ``True`` on success, ``False`` if
    the spec/body is missing or the recompile fails (both logged), so the
    caller can surface a clean error instead of a silent no-op.
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("reposition_body: no spec or model in world")
        return False

    try:
        body = spec.body(body_name)
    except (KeyError, ValueError):
        body = None
    if body is None:
        logger.warning("Body '%s' not found in spec - nothing repositioned", body_name)
        return False

    if position is not None:
        body.pos = list(position)
    if orientation is not None:
        body.quat = list(orientation)

    return _recompile_preserving_state(world, spec)


def _snapshot_joint_state(world: SimWorld) -> dict[str, tuple[list[float], list[float]]]:
    """Snapshot per-joint ``(qpos, qvel)`` slices keyed by fully-qualified
    MuJoCo joint name.

    Used by :func:`eject_robot_from_scene` to preserve the state of surviving
    robots and object freejoints across a scene rebuild. Flat-index slicing
    is unsafe here because the body-tree order may shift when a robot is
    removed (see AGENTS.md "Per-name state copy" rule).

    Returns a dict mapping ``<joint_name> -> (qpos_slice, qvel_slice)`` where
    each slice has the appropriate width for the joint type (1 for hinge/
    slide, 4 for ball, 7 for free).
    """
    if world._model is None or world._data is None:
        return {}
    mj = _ensure_mujoco()
    model = world._model
    data = world._data
    snap: dict[str, tuple[list[float], list[float]]] = {}
    for jid in range(model.njnt):
        name = mj.mj_id2name(model, mj.mjtObj.mjOBJ_JOINT, jid)
        if not name:
            continue
        qpos_adr = int(model.jnt_qposadr[jid])
        qvel_adr = int(model.jnt_dofadr[jid])
        jtype = int(model.jnt_type[jid])
        # qpos width: free=7, ball=4, hinge/slide=1
        # qvel width: free=6, ball=3, hinge/slide=1
        if jtype == mj.mjtJoint.mjJNT_FREE:
            qpos_w, qvel_w = 7, 6
        elif jtype == mj.mjtJoint.mjJNT_BALL:
            qpos_w, qvel_w = 4, 3
        else:
            qpos_w, qvel_w = 1, 1
        snap[name] = (
            [float(x) for x in data.qpos[qpos_adr : qpos_adr + qpos_w]],
            [float(x) for x in data.qvel[qvel_adr : qvel_adr + qvel_w]],
        )
    return snap


def _restore_joint_state(
    world: SimWorld,
    snapshot: dict[str, tuple[list[float], list[float]]],
) -> int:
    """Restore per-joint state from a snapshot into ``world._data`` by name.

    Joints that no longer exist in the compiled model (e.g. those belonging
    to the ejected robot) are silently skipped. Joints that exist in the
    new model but were not in the snapshot keep their fresh-compile defaults
    (body pos/quat for freejoints, 0 for hinge/slide).

    Returns the number of joints actually restored, for logging.
    """
    if world._model is None or world._data is None:
        return 0
    mj = _ensure_mujoco()
    model = world._model
    data = world._data
    restored = 0
    for name, (qpos_vals, qvel_vals) in snapshot.items():
        jid = mj_name_to_id(model, mj.mjtObj.mjOBJ_JOINT, name)
        if jid < 0:
            continue  # joint no longer exists (expected for ejected robot)
        qpos_adr = int(model.jnt_qposadr[jid])
        qvel_adr = int(model.jnt_dofadr[jid])
        # Width sanity check: if joint type changed (should not happen for
        # same-name joints across an eject), skip to avoid corrupting state.
        jtype = int(model.jnt_type[jid])
        if jtype == mj.mjtJoint.mjJNT_FREE:
            expect_qp, expect_qv = 7, 6
        elif jtype == mj.mjtJoint.mjJNT_BALL:
            expect_qp, expect_qv = 4, 3
        else:
            expect_qp, expect_qv = 1, 1
        if len(qpos_vals) != expect_qp or len(qvel_vals) != expect_qv:
            logger.warning(
                "_restore_joint_state: width mismatch for %r (qpos %d!=%d or qvel %d!=%d), skipping",
                name,
                len(qpos_vals),
                expect_qp,
                len(qvel_vals),
                expect_qv,
            )
            continue
        for i, v in enumerate(qpos_vals):
            data.qpos[qpos_adr + i] = v
        for i, v in enumerate(qvel_vals):
            data.qvel[qvel_adr + i] = v
        restored += 1
    return restored


def eject_robot_from_scene(world: SimWorld, robot_name: str) -> bool:
    """Remove every spec element namespaced under ``{robot_name}/``.

    Implementation note: deleting a body that was added via ``spec.attach()``
    triggers a known MuJoCo 3.8 segfault at interpreter shutdown (the
    attached child spec's memory gets freed twice). To sidestep that bug
    we REBUILD the scene spec from scratch using the post-remove
    ``world.robots`` / ``world.objects`` / ``world.cameras`` state, then
    re-attach the remaining robots.

    Joint state preservation: before the rebuild we snapshot every joint's
    ``(qpos, qvel)`` keyed by fully-qualified name; after the fresh compile
    we restore state for every joint that still exists in the new model.
    Joints belonging to the ejected robot are naturally dropped (their name
    no longer resolves). This keeps surviving robots at their current pose
    and object freejoints at their current world pose - the behaviour the
    agent expects when calling ``remove_robot`` mid-scene.

    Flat-index slicing is **not** safe here: removing a robot shifts every
    body/joint index that comes after it in the kinematic tree, so
    ``data.qpos[:]`` copies across compiles would mis-assign DOFs. Per-name
    lookup is the only correct approach (see AGENTS.md).
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("eject_robot: no spec or model in world")
        return False

    mj = _ensure_mujoco()

    # Snapshot joint state BEFORE we rebuild. Keyed by the fully-qualified
    # MuJoCo joint name (prefix/joint for attached robots, bare name for
    # object freejoints).
    state_snapshot = _snapshot_joint_state(world)

    # First drop cameras that originated from the robot being ejected.
    # They're in world.cameras with origin_robot == robot_name. Without this,
    # SpecBuilder.build would skip them (via origin_robot), but stale entries
    # would linger in the registry and confuse observation code.
    stale_cam_names = [cname for cname, cam in world.cameras.items() if getattr(cam, "origin_robot", "") == robot_name]
    for cname in stale_cam_names:
        del world.cameras[cname]

    # Step 1: rebuild the base spec from world (objects + cameras +
    # lights + ground).
    new_spec = SpecBuilder.build(world)

    # Step 2: re-attach every remaining robot (the one being ejected is
    # already popped from ``world.robots`` by the caller).
    for robot in world.robots.values():
        # Re-discover joint names via the attach - they're stable per URDF.
        with filter_mujoco_attach_noise():
            joint_names = SpecBuilder.attach_robot(new_spec, robot, robot.urdf_path)
        robot.joint_names = joint_names

    # Step 2b: mount the body-mounted cameras ``SpecBuilder.build`` deferred,
    # now that every surviving robot's bodies exist in the spec. A camera whose
    # parent belonged to the robot being ejected has no parent left, so it is
    # dropped from the registry with a warning rather than aborting the eject -
    # the same treatment the robot's own URDF cameras get above. Without this
    # step the rebuild raised ``ValueError`` straight out of ``remove_robot``,
    # so a scene with a wrist camera could not remove any robot at all.
    for cname in SpecBuilder.add_deferred_cameras(new_spec, world):
        logger.warning(
            "eject_robot: dropping camera %r - its parent body %r belonged to the removed robot %r.",
            cname,
            getattr(world.cameras[cname], "parent_body", ""),
            robot_name,
        )
        del world.cameras[cname]

    # Step 3: compile fresh and install. No spec.recompile(model, data)
    # here - recompile implicitly preserves qpos state which doesn't
    # make sense across a scene rebuild, and forcing a fresh compile
    # avoids the attach/delete bug.
    try:
        with filter_mujoco_attach_noise():
            new_model = new_spec.compile()
        new_data = mj.MjData(new_model)
    except (ValueError, RuntimeError) as e:
        logger.error("eject_robot: fresh compile failed: %s", e)
        return False

    install_compiled_model(world, new_model, new_data)
    world._backend_state["spec"] = new_spec
    _sync_cached_xml(world, new_spec)

    # Step 4: restore state for every joint that survived the rebuild. Joints
    # belonging to the ejected robot simply don't resolve and get skipped.
    restored = _restore_joint_state(world, state_snapshot)

    # Step 5: run a forward pass so derived quantities (xpos, cam xforms)
    # reflect the restored state. Without this, the next render() call can
    # produce stale frames because MjData was freshly allocated in Step 3.
    mj.mj_forward(new_model, new_data)

    # Re-discover joint/actuator IDs for remaining robots.
    for robot in world.robots.values():
        pfx = robot.namespace or ""
        robot.joint_ids = []
        for jnt_name in robot.joint_names:
            jid = -1
            if pfx:
                jid = mj_name_to_id(new_model, mj.mjtObj.mjOBJ_JOINT, pfx + jnt_name)
            if jid < 0:
                jid = mj_name_to_id(new_model, mj.mjtObj.mjOBJ_JOINT, jnt_name)
            if jid >= 0:
                robot.joint_ids.append(jid)
        robot.actuator_ids = robot_owned_actuator_ids(new_model, robot, mj)
        # See _recompile_preserving_state for why the lone-robot case falls back.
        if not robot.actuator_ids and len(world.robots) == 1:
            robot.actuator_ids = list(range(new_model.nu))

    logger.debug(
        "eject_robot %r: scene rebuilt, restored state for %d/%d joints",
        robot_name,
        restored,
        len(state_snapshot),
    )
    return True


# Runtime attach / actuate primitives (GH #1533, PR 1)


def add_weld_constraint(
    world: SimWorld,
    *,
    name: str,
    parent: str,
    child: str,
    relpos: list[float],
    relquat: list[float],
    torquescale: float = 1.0,
) -> bool:
    """Add a named weld equality constraint between two bodies and recompile.

    ``relpos`` / ``relquat`` are the pose of ``child`` expressed in ``parent``'s
    frame - callers capture the CURRENT runtime relative pose so the weld holds
    the bodies exactly where they are (MuJoCo's all-zero ``relpose`` quat would
    instead bake in the compile-time ``qpos0`` pose, which is wrong for a
    runtime grasp-attach). The equality is stored on the live spec so it
    survives later recompiles (``add_object`` etc.).

    Returns ``True`` on success. On recompile failure the just-added equality
    is deleted again so the spec stays compilable, and ``False`` is returned.
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("add_weld_constraint: no spec or model in world")
        return False
    mj = _ensure_mujoco()

    eq = spec.add_equality()
    eq.name = name
    eq.type = mj.mjtEq.mjEQ_WELD
    eq.objtype = mj.mjtObj.mjOBJ_BODY
    eq.name1 = parent
    eq.name2 = child
    # mjEQ_WELD data layout (mjNEQDATA=11): [anchor(3), relpose pos(3),
    # relpose quat(4), torquescale(1)]. anchor stays zero - relpose fully
    # determines the held configuration.
    data = [0.0] * 11
    data[3:6] = [float(v) for v in relpos]
    data[6:10] = [float(v) for v in relquat]
    data[10] = float(torquescale)
    eq.data = data

    if not _recompile_preserving_state(world, spec):
        spec.delete(eq)
        return False
    return True


def remove_equality_constraint(world: SimWorld, name: str) -> bool:
    """Delete a named equality constraint from the live spec and recompile.

    Returns ``False`` (logged) when the constraint is missing or the recompile
    fails, so callers can surface a clean error instead of a silent no-op.
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("remove_equality_constraint: no spec or model in world")
        return False
    for eq in spec.equalities:
        if eq.name == name:
            spec.delete(eq)
            return _recompile_preserving_state(world, spec)
    logger.warning("Equality constraint '%s' not found in spec - nothing removed", name)
    return False


def actuate_robot_in_scene(
    world: SimWorld,
    robot: SimRobot,
    kp_by_joint: dict[str, float],
    *,
    damping: float,
    armature: float,
    gravity_compensation: bool,
    disable_self_collision: bool,
) -> bool:
    """Add position-servo actuators to a robot's joints and recompile.

    The supported form of the private-spec surgery the ``so101_curobo`` example
    performed by hand (GH #1533): converts an actuator-less (URDF-loaded) arm
    into a position-controlled one so ``send_action`` / ``run_policy`` can
    drive it. Per joint in ``kp_by_joint`` (SHORT joint names, values = kp):

    * a position actuator (``set_to_position``: fixed gain kp, affine bias with
      ``dampratio=1.0`` for ~critical damping, ctrlrange inherited from the
      joint's range when it declares one),
    * joint ``damping`` / ``armature`` floors (bare URDFs ship none, which
      blows up explicit integration),

    plus, scene-wide, the stable ``implicitfast`` integrator, and optionally
    gravity compensation on the robot's bodies and self-collision disable on
    the robot's own geoms (cuRobo-style planners ignore adjacent-link
    contacts, which otherwise block planned motion).

    Atomicity: the spec is snapshotted (XML round-trip) before surgery; any
    failure restores the snapshot so no partial edit (integrator flip, gravcomp,
    half the actuators) lingers on the live spec. Returns ``True`` on success.
    """
    spec = _get_spec(world)
    if spec is None or world._model is None:
        logger.error("actuate_robot_in_scene: no spec or model in world")
        return False
    mj = _ensure_mujoco()
    pfx = robot.namespace or ""

    # Snapshot for atomic rollback (mirrors patch_scene_mjcf): the surgery
    # touches option/bodies/joints/actuators/geoms, too many objects to undo
    # piecewise.
    backup_spec = _snapshot_spec(spec, context="actuate_robot_in_scene")
    if backup_spec is None:
        return False

    try:
        # Bare URDF chains (no damping/armature) diverge under the default
        # Euler integrator once stiff position servos are added; implicitfast
        # integrates joint damping implicitly and stays stable.
        spec.option.integrator = mj.mjtIntegrator.mjINT_IMPLICITFAST

        for body in spec.bodies:
            body_name = body.name or ""
            if pfx and body_name.startswith(pfx):
                if gravity_compensation:
                    body.gravcomp = 1.0
                if disable_self_collision:
                    for geom in body.geoms:
                        geom.contype = 0
                        geom.conaffinity = 0

        for joint in spec.joints:
            joint_name = joint.name or ""
            if not (pfx and joint_name.startswith(pfx)):
                continue
            short = joint_name[len(pfx) :]
            if short not in kp_by_joint:
                continue
            joint.damping[0] = max(float(joint.damping[0]), damping)
            joint.armature = max(float(joint.armature), armature)

        for short, kp in kp_by_joint.items():
            act = spec.add_actuator()
            act.name = f"{robot.name}_act_{short}"
            act.target = f"{pfx}{short}"
            act.trntype = mj.mjtTrn.mjTRN_JOINT
            # Find the joint's spec range: inheritrange clamps ctrl to the
            # joint's limits when it declares any (URDF limits map here);
            # unlimited joints keep an unlimited ctrlrange.
            jnt_range_defined = False
            for joint in spec.joints:
                if (joint.name or "") == f"{pfx}{short}":
                    jnt_range_defined = bool(float(joint.range[0]) < float(joint.range[1]))
                    break
            act.set_to_position(kp=float(kp), dampratio=1.0, inheritrange=jnt_range_defined)
    except (ValueError, RuntimeError, TypeError) as e:
        logger.error("actuate_robot_in_scene: spec surgery failed for '%s': %s", robot.name, e)
        world._backend_state["spec"] = backup_spec
        return False

    if not _recompile_preserving_state(world, spec):
        # Restore the pre-surgery spec so the failed edit doesn't poison
        # later scene mutations.
        world._backend_state["spec"] = backup_spec
        return False
    return True


# Agent-authored raw MJCF (Stage 6)


def replace_scene_mjcf(world: SimWorld, xml: str) -> bool:
    """Atomically swap the whole scene for agent-written MJCF.

    Validated by actually compiling it. On failure raises ``ValueError`` with
    MuJoCo's compiler error verbatim. On success, the old spec/model/data are
    replaced and all per-robot joint/actuator IDs re-discovered (but since
    the agent may have changed the whole scene, the ``world.robots`` dict
    is NOT touched - that's the caller's responsibility).
    """
    mj = _ensure_mujoco()
    new_spec = SpecBuilder.from_mjcf_string(xml)
    # Compile eagerly so malformed XML fails here rather than on the next
    # mj_step.
    with filter_mujoco_attach_noise():
        new_model = new_spec.compile()
    new_data = mj.MjData(new_model)

    world._backend_state["spec"] = new_spec
    install_compiled_model(world, new_model, new_data)

    # Run a single forward pass so geom positions / camera xforms are
    # populated. Without this, the very first sim.render() call after
    # replace_scene_mjcf hits `data.xpos == 0 for all bodies` and the
    # renderer dumps a 100% black frame. Matches the semantics of
    # _compile_world() which also calls mj_forward after MjData construction.
    mj.mj_forward(new_model, new_data)

    _sync_cached_xml(world, new_spec)
    return True


# Structured-op patching of the live spec (Stage 6, part 2 - GH #125)

# Supported ops for patch_scene_mjcf, each mapped to the complete set of keys
# that op reads. Kept narrow on purpose - adding unchecked attribute setters
# would make the tool an arbitrary-code hole. Agents that need exotic MJCF
# should go through replace_scene_mjcf with a full XML.
#
# This mapping is the single source of truth for the op vocabulary: it names
# the supported ops AND, per op, every key that reaches MuJoCo. Any other key
# is refused, because each field below is read with a fallback default - a
# misspelled key would otherwise apply that default and report success, so
# ``{"op": "set_body_pos", "name": "crate", "position": [...]}`` would move the
# body to the origin rather than to the requested pose.
_PATCH_OP_KEYS: dict[str, frozenset[str]] = {
    "add_body": frozenset({"op", "parent", "name", "pos", "quat"}),
    "add_geom": frozenset({"op", "body", "type", "size", "rgba", "name", "pos", "quat"}),
    "add_site": frozenset({"op", "body", "name", "pos", "size", "rgba"}),
    "set_body_pos": frozenset({"op", "name", "pos"}),
    "set_body_quat": frozenset({"op", "name", "quat"}),
    "delete_body": frozenset({"op", "name"}),
}


def _unknown_op_keys_error(kind: str, op: Mapping[str, Any]) -> str | None:
    """Reject keys a patch op does not read, before its defaults are applied.

    Every field of every op is read with a fallback default (``pos`` defaults
    to the origin, ``quat`` to identity, ``type`` to ``"box"``, ``parent`` to
    ``"world"``), so an unrecognised key is not an inert extra: the op runs
    with that default and reports success. A misspelled ``pos`` silently moves
    a body to the world origin, a misspelled ``parent`` re-parents it to the
    worldbody, and a misspelled ``type`` compiles a box where a sphere was
    asked for.

    Args:
        kind: The op name, already known to be in :data:`_PATCH_OP_KEYS`.
        op: The caller-supplied op dict.

    Returns:
        An error message naming the unknown key(s), a close-match suggestion
        where one exists, and the keys this op accepts - or ``None`` when
        every key is honored.
    """
    accepted = _PATCH_OP_KEYS[kind]
    unknown = [key for key in op if key not in accepted]
    if not unknown:
        return None
    known = sorted(accepted)
    # Suggest field names only. "op" selects the op and is already validated
    # above, so it is never what a misspelled field meant - offering it turns a
    # real MJCF attribute like "group" into a nonsense hint.
    candidates = [key for key in known if key != "op"]
    described = []
    for key in sorted(unknown, key=str):
        close = difflib.get_close_matches(str(key).lower(), candidates, n=1, cutoff=0.5)
        described.append(f"{key!r}" + (f" (did you mean {close[0]!r}?)" if close else ""))
    return f"{kind}: unknown op key(s): {', '.join(described)}. Accepted keys: {', '.join(known)}."


# The numeric vector fields each op writes into the compiled MJCF. Every name
# here is also an accepted key of the same op in :data:`_PATCH_OP_KEYS`; a
# parity test pins that, so the two tables cannot drift into disagreement about
# which fields an op reads.
_PATCH_OP_VECTOR_FIELDS: dict[str, frozenset[str]] = {
    "add_body": frozenset({"pos", "quat"}),
    "add_geom": frozenset({"pos", "quat", "size", "rgba"}),
    "add_site": frozenset({"pos", "size", "rgba"}),
    "set_body_pos": frozenset({"pos"}),
    "set_body_quat": frozenset({"quat"}),
    "delete_body": frozenset(),
}


def _rgba_field_domain(kind: str, field: str, value: Any) -> tuple[Any, str | None]:
    """Validate an op's colour and complete a 3-component RGB with opaque alpha.

    Thin adapter so :func:`~strands_robots.utils.coerce_rgba` - the single
    definition of the library's colour contract - can sit in
    :data:`_OP_FIELD_DOMAINS` beside the other field domains.

    It differs from that helper in one respect, because an op dict differs from a
    keyword argument: a key that is PRESENT carries a supplied value, so ``None``
    is a colour this op cannot apply rather than an omission asking for the
    default. Reading it as omitted would paint the fallback grey under a success
    result, which is the silent default the op-key check exists to prevent.

    Args:
        kind: The op name, for the message.
        field: The field name, always ``"rgba"``.
        value: The caller-supplied colour.

    Returns:
        ``(rgba, None)`` with exactly four components, or ``(None, message)``.
    """
    if value is None:
        return None, f"{kind}: '{field}' must be a sequence of numbers, got None"
    return coerce_rgba(kind, field, value)


# The domain each numeric op field is held to, keyed by field name. Every field
# any op declares in :data:`_PATCH_OP_VECTOR_FIELDS` has an entry here and a
# parity test pins that, so a field added to an op cannot reach the compiled
# model without a decided domain.
#
# The widths are the library's, not MuJoCo's: a ``pos`` is three components and a
# ``quat`` four whichever op writes it, and a colour is the contract
# :func:`~strands_robots.utils.coerce_rgba` defines for every backend. ``size``
# is the one field whose count is shape-dependent, so only its components are
# checked here and the count is left to ``_validate_size``, which knows the shape.
_OP_FIELD_DOMAINS: dict[str, Callable[[str, str, Any], tuple[Any, str | None]]] = {
    "pos": lambda kind, field, value: (value, pose_vector_error(kind, field, value, 3)),
    "quat": lambda kind, field, value: (value, pose_vector_error(kind, field, value, 4)),
    "rgba": _rgba_field_domain,
    "size": lambda kind, field, value: (value, finite_vector_error(kind, field, value)),
}


def _normalized_op_vector_fields(kind: str, op: Mapping[str, Any]) -> tuple[dict[str, Any] | None, str | None]:
    """Hold every numeric op field to its domain, before the spec is touched.

    MuJoCo's compiler does not reject a ``nan``/``inf`` pose, extent or colour
    (its one exception is a ``nan`` geom size), so an unchecked component is
    written verbatim into the model and the patch reports success. A non-finite
    ``pos`` on a body that owns a freejoint then poisons ``qpos``/``qvel`` on the
    next step, and every call in that chain - the patch, the step, the
    observation read - keeps reporting success, so a caller has no way to learn
    the scene is dead.

    Component count is part of the same contract. These ops write the very
    buffers ``add_object``, ``add_camera`` and ``move_object`` write, so a value
    either surface refuses has to be refused by the other - and MuJoCo alone
    cannot deliver that:

    * ``set_body_pos`` / ``set_body_quat`` assign the field as a spec ATTRIBUTE
      (``body.pos = ...``) rather than passing it as a constructor keyword, and
      pybind11 reports a width mismatch there by dumping its C++ overload table
      and the receiving object's address. That message names neither the op nor
      the field, so the one thing the caller needs is the one thing absent from
      it, while the sibling ops writing the same two fields report cleanly.
    * A three-component ``rgba`` is refused outright, though it is the RGB that
      ``add_object(color=...)`` accepts and completes with an opaque alpha. One
      backend, two surfaces, one ``geom_rgba`` buffer, opposite verdicts.

    Args:
        kind: The op name, already known to be in :data:`_PATCH_OP_KEYS`.
        op: The caller-supplied op dict.

    Returns:
        ``(op, None)`` with every numeric field present normalized to what the
        spec write consumes, or ``(None, message)`` naming the op, the field and
        the offending value.
    """
    normalized = dict(op)
    for field in sorted(_PATCH_OP_VECTOR_FIELDS[kind]):
        if field not in op:
            continue
        value, msg = _OP_FIELD_DOMAINS[field](kind, field, op[field])
        if msg is not None:
            return None, msg
        normalized[field] = value
    return normalized, None


def _find_body(spec: Any, name: str, new_bodies: dict[str, Any]) -> Any:
    """Locate a body by name in a live spec, checking batch-local additions.

    MuJoCo 3.8 ``spec.body(name)`` only resolves bodies that existed at the
    last ``compile()`` / ``recompile()`` call. Bodies added mid-batch are
    not visible through that lookup but ARE present on the spec - we track
    their handles in ``new_bodies`` so ``add_geom`` / ``add_site`` /
    ``set_body_pos`` etc. can reference them within the same patch.
    """
    if name == "world":
        return spec.worldbody
    if name in new_bodies:
        return new_bodies[name]
    b = spec.body(name)
    if b is not None:
        return b
    # Fallback: scan all bodies. Catches bodies introduced via spec.attach()
    # (e.g. robots composed into the scene) that aren't in new_bodies because
    # we didn't create them in this batch.
    for body in spec.bodies:
        if body.name == name:
            return body
    return None


def _apply_patch_op(spec: Any, op: dict[str, Any], new_bodies: dict[str, Any]) -> None:
    """Apply a single structured op to a live MjSpec.

    Raises ``ValueError`` with a human-readable message on bad input;
    MuJoCo compile errors surface on the enclosing ``recompile`` call.
    ``new_bodies`` is a batch-local cache of body handles added earlier
    in the same patch (see ``_find_body`` for why this is needed).
    """
    if not isinstance(op, dict):
        raise ValueError(f"each op must be a dict, got {type(op).__name__}")

    kind = op.get("op")
    if kind not in _PATCH_OP_KEYS:
        raise ValueError(f"unknown op '{kind}'. Supported: {sorted(_PATCH_OP_KEYS)}")
    if err := _unknown_op_keys_error(str(kind), op):
        raise ValueError(err)
    normalized, err = _normalized_op_vector_fields(str(kind), op)
    if normalized is None:
        raise ValueError(err)
    # Every numeric field below now holds the value the spec write consumes -
    # notably a three-component 'rgba' completed to RGBA.
    op = normalized

    if kind == "add_body":
        parent = op.get("parent", "world")
        name = op.get("name")
        if not name:
            raise ValueError("add_body requires 'name'")
        pos = op.get("pos", [0.0, 0.0, 0.0])
        quat = op.get("quat", [1.0, 0.0, 0.0, 0.0])
        parent_body = _find_body(spec, parent, new_bodies)
        if parent_body is None:
            raise ValueError(f"add_body: parent '{parent}' not found")
        new_body = parent_body.add_body(name=name, pos=pos, quat=quat)
        new_bodies[name] = new_body
        return

    if kind == "add_geom":
        body_name = op.get("body")
        if not body_name:
            raise ValueError("add_geom requires 'body'")
        body = _find_body(spec, body_name, new_bodies)
        if body is None:
            raise ValueError(f"add_geom: body '{body_name}' not found")

        shape = op.get("type", "box")
        from strands_robots.simulation.mujoco.spec_builder import (
            _geom_type,
            _normalize_size,
        )

        geom_kwargs: dict[str, Any] = {
            "type": _geom_type(shape),
            "size": _normalize_size(shape, op.get("size", [0.1, 0.1, 0.1])),
            "rgba": op.get("rgba", [0.5, 0.5, 0.5, 1.0]),
        }
        if "name" in op:
            geom_kwargs["name"] = op["name"]
        if "pos" in op:
            geom_kwargs["pos"] = op["pos"]
        if "quat" in op:
            geom_kwargs["quat"] = op["quat"]
        body.add_geom(**geom_kwargs)
        return

    if kind == "add_site":
        body_name = op.get("body", "world")
        body = _find_body(spec, body_name, new_bodies)
        if body is None:
            raise ValueError(f"add_site: body '{body_name}' not found")
        name = op.get("name")
        if not name:
            raise ValueError("add_site requires 'name'")
        site_kwargs: dict[str, Any] = {
            "name": name,
            "pos": op.get("pos", [0.0, 0.0, 0.0]),
        }
        if "size" in op:
            site_kwargs["size"] = op["size"]
        if "rgba" in op:
            site_kwargs["rgba"] = op["rgba"]
        body.add_site(**site_kwargs)
        return

    if kind == "set_body_pos":
        name = op.get("name")
        if not name:
            raise ValueError("set_body_pos requires 'name'")
        body = _find_body(spec, name, new_bodies)
        if body is None:
            raise ValueError(f"set_body_pos: body '{name}' not found")
        body.pos = op.get("pos", [0.0, 0.0, 0.0])
        return

    if kind == "set_body_quat":
        name = op.get("name")
        if not name:
            raise ValueError("set_body_quat requires 'name'")
        body = _find_body(spec, name, new_bodies)
        if body is None:
            raise ValueError(f"set_body_quat: body '{name}' not found")
        body.quat = op.get("quat", [1.0, 0.0, 0.0, 0.0])
        return

    if kind == "delete_body":
        name = op.get("name")
        if not name:
            raise ValueError("delete_body requires 'name'")
        body = _find_body(spec, name, new_bodies)
        if body is None:
            raise ValueError(f"delete_body: body '{name}' not found")
        spec.delete(body)
        new_bodies.pop(name, None)
        return


def patch_scene_mjcf(world: SimWorld, ops: list[dict[str, Any]]) -> int:
    """Apply a sequence of structured ops to the live spec in order.

    Each op is a small dict like::

        {"op": "add_body", "parent": "world", "name": "foo", "pos": [0,0,1]}
        {"op": "add_geom", "body": "foo", "type": "sphere", "size": [0.1]}
        {"op": "set_body_pos", "name": "foo", "pos": [1,0,1]}
        {"op": "delete_body", "name": "foo"}

    Each op accepts exactly the keys listed for it in :data:`_PATCH_OP_KEYS`;
    anything else is rejected, because an unread key would leave the op running
    on its fallback default (see :func:`_unknown_op_keys_error`). Every numeric
    field an op writes is held to its domain in :data:`_OP_FIELD_DOMAINS` - a
    ``pos`` of exactly 3 finite components, a ``quat`` of 4, a ``rgba`` of 3
    (RGB, completed with an opaque alpha) or 4 - because MuJoCo bakes a
    ``nan``/``inf`` component into the model without complaint, and reports a
    width mismatch on the attribute writes by dumping a C++ overload table that
    names neither the op nor the field.

    The list is applied atomically: if any op raises, the whole patch is
    rejected and the world is left in its original state. After all ops
    succeed, ``spec.recompile(model, data)`` is called once, so joint
    qpos/qvel for unchanged joints are preserved automatically.

    Returns the number of ops applied (same as ``len(ops)`` on success).
    """
    if not isinstance(ops, list):
        raise ValueError(f"ops must be a list, got {type(ops).__name__}")
    if not ops:
        return 0

    spec = world._backend_state.get("spec")
    if spec is None:
        raise RuntimeError("world has no spec; patch_scene_mjcf requires a compiled world")

    # Snapshot the spec so a failed op can be atomically rejected: the batch is
    # applied to the live spec and the snapshot put back if any op fails.
    backup_spec = _snapshot_spec(spec, context="patch_scene_mjcf")
    if backup_spec is None:
        raise RuntimeError("failed to snapshot spec before patch (see logs)")

    applied = 0
    new_bodies: dict[str, Any] = {}
    try:
        for op in ops:
            _apply_patch_op(spec, op, new_bodies)
            applied += 1
    except Exception as err:
        world._backend_state["spec"] = backup_spec
        raise ValueError(f"patch op #{applied + 1} failed: {err}") from err

    # One recompile for the whole batch - preserves qpos/qvel for unchanged joints.
    with filter_mujoco_attach_noise():
        new_model, new_data = spec.recompile(world._model, world._data)
    install_compiled_model(world, new_model, new_data)

    # Forward pass so new bodies' xpos / xquat / cam_xmat are populated for
    # the very next render() or get_body_state() call. Same reasoning as
    # replace_scene_mjcf.
    mj = _ensure_mujoco()
    mj.mj_forward(world._model, world._data)

    _sync_cached_xml(world, spec)
    return applied
