"""Unit + integration tests for ``SpecBuilder`` - the MjSpec-based MJCF builder.

Tests cover:

* Module-level helpers (``_geom_type``, ``_normalize_size``, ``_target_quat``,
  ``_is_z0_ground_plane``).
* ``SpecBuilder.build`` produces a compile-valid spec for empty, object-only,
  and camera-only worlds.
* Mutation helpers (``add_object``, ``remove_body``, ``add_camera``,
  ``remove_camera``) produce specs that recompile cleanly.
* ``attach_robot`` prefixes names correctly and returns the source joint names.
* ``from_mjcf_string`` / ``from_file`` round-trip cleanly.

These tests use the actual mujoco AST - they require the ``mujoco`` package.
"""

from __future__ import annotations

import types

import pytest

pytest.importorskip("mujoco")

import mujoco  # noqa: E402
import numpy as np  # noqa: E402

from strands_robots.simulation.models import (  # noqa: E402
    SimCamera,
    SimObject,
    SimRobot,
    SimWorld,
)
from strands_robots.simulation.mujoco.spec_builder import (  # noqa: E402
    SpecBuilder,
    _geom_type,
    _is_z0_ground_plane,
    _normalize_size,
    _target_quat,
)

# Module-level helpers


class TestGeomType:
    def test_known_shapes_map_to_enum(self):
        assert _geom_type("box") == mujoco.mjtGeom.mjGEOM_BOX
        assert _geom_type("sphere") == mujoco.mjtGeom.mjGEOM_SPHERE
        assert _geom_type("cylinder") == mujoco.mjtGeom.mjGEOM_CYLINDER
        assert _geom_type("capsule") == mujoco.mjtGeom.mjGEOM_CAPSULE
        assert _geom_type("mesh") == mujoco.mjtGeom.mjGEOM_MESH
        assert _geom_type("plane") == mujoco.mjtGeom.mjGEOM_PLANE
        assert _geom_type("ellipsoid") == mujoco.mjtGeom.mjGEOM_ELLIPSOID

    def test_unknown_shape_raises_with_helpful_list(self):
        with pytest.raises(ValueError, match="Unsupported shape"):
            _geom_type("hyperboloid")


class TestNormalizeSize:
    def test_box_halves_full_extents(self):
        assert _normalize_size("box", [0.2, 0.4, 0.6]) == [0.1, 0.2, 0.3]

    def test_sphere_halves_first_coordinate(self):
        assert _normalize_size("sphere", [0.1])[0] == pytest.approx(0.05)

    def test_cylinder_radius_and_half_height(self):
        out = _normalize_size("cylinder", [0.1, 0, 0.4])
        assert out[0] == pytest.approx(0.05)  # radius
        assert out[1] == pytest.approx(0.2)  # half-height

    def test_capsule_same_as_cylinder(self):
        assert _normalize_size("capsule", [0.1, 0, 0.4]) == _normalize_size("cylinder", [0.1, 0, 0.4])

    def test_plane_size_duplicates_x_to_y_when_only_one_given(self):
        out = _normalize_size("plane", [2.0])
        assert out[0] == 2.0
        assert out[1] == 2.0

    def test_mesh_returns_zeros(self):
        assert _normalize_size("mesh", [0.1, 0.2, 0.3]) == [0.0, 0.0, 0.0]

    def test_ellipsoid_halves_all_three_full_extents(self):
        # An ellipsoid's three full diameters become per-axis radii (each /2).
        assert _normalize_size("ellipsoid", [0.1, 0.2, 0.3]) == [0.05, 0.1, 0.15]

    def test_ellipsoid_size_too_short_raises_instead_of_defaulting(self):
        # Corrected contract: a <3 component ellipsoid size cannot express the
        # request, and completing it from a hardcoded default compiled a
        # differently-sized geom while add_object reported success (echoing the
        # size the caller asked for). It is now rejected, so the caller learns
        # which components the shape consumes.
        with pytest.raises(ValueError, match=r"needs 3 'size' component"):
            _normalize_size("ellipsoid", [])

    def test_unknown_shape_raises(self):
        with pytest.raises(ValueError, match="Cannot normalize size"):
            _normalize_size("hyperboloid", [1.0])

    def test_rejects_non_positive_extent(self):
        # A zero or negative extent on an axis the shape actually consumes
        # is rejected up front (add_object surfaces this as a clear error
        # rather than compiling a degenerate geom). The message states the
        # full-extent convention so a caller can self-correct.
        with pytest.raises(ValueError, match="extent must be > 0"):
            _normalize_size("box", [0.2, 0.0, 0.4])
        with pytest.raises(ValueError, match="extent must be > 0"):
            _normalize_size("sphere", [-0.1])


class TestIsZ0GroundPlane:
    """The z=0 ground-plane detector used to strip a redundant robot floor
    when attaching into a world that already owns a ground plane (#363).

    Only an axis-aligned plane sitting at z~0 is treated as the ground; a
    plane a robot MJCF ships intentionally (an elevated ramp, or one tilted
    via a non-identity quat such as a wall) must NOT match, so it survives
    the attach instead of being silently deleted.
    """

    def test_axis_aligned_z0_plane_is_ground(self):
        geom = types.SimpleNamespace(pos=[0.0, 0.0, 0.0], quat=[1.0, 0.0, 0.0, 0.0])
        assert _is_z0_ground_plane(geom) is True

    def test_missing_pos_and_quat_defaults_to_ground(self):
        # A bare plane geom with neither attribute set defaults to +Z at the
        # body origin, which is the ground case.
        geom = types.SimpleNamespace()
        assert _is_z0_ground_plane(geom) is True

    def test_elevated_plane_is_not_ground(self):
        geom = types.SimpleNamespace(pos=[0.0, 0.0, 0.5], quat=[1.0, 0.0, 0.0, 0.0])
        assert _is_z0_ground_plane(geom) is False

    def test_rotated_plane_is_not_ground(self):
        # A plane rotated off the +Z normal (e.g. a wall or angled ramp
        # modeled at z=0) has a non-identity quat and must survive attach.
        geom = types.SimpleNamespace(pos=[0.0, 0.0, 0.0], quat=[0.7071, 0.7071, 0.0, 0.0])
        assert _is_z0_ground_plane(geom) is False


class TestTargetQuat:
    def test_returns_none_for_degenerate(self):
        assert _target_quat([1, 2, 3], [1, 2, 3]) is None

    def test_returns_normalised_quaternion(self):
        quat = _target_quat([1, 0, 0.5], [0, 0, 0])
        assert quat is not None
        assert len(quat) == 4
        norm = (quat[0] ** 2 + quat[1] ** 2 + quat[2] ** 2 + quat[3] ** 2) ** 0.5
        assert norm == pytest.approx(1.0, abs=1e-6)

    def test_top_down_camera_resolves_valid_quat(self):
        # Top-down overhead camera: forward (0,0,-1) is parallel to world +Z.
        # Previously this returned None (degenerate) and add_camera silently
        # dropped the orientation, leaving the camera mis-pointed. It must now
        # fall back to world +Y and produce a valid, normalised quaternion.
        quat = _target_quat([0.2, 0, 0.5], [0.2, 0, 0])
        assert quat is not None
        assert len(quat) == 4
        norm = sum(c * c for c in quat) ** 0.5
        assert norm == pytest.approx(1.0, abs=1e-6)

    def test_straight_up_camera_resolves_valid_quat(self):
        # Looking straight up: forward (0,0,+1), also parallel to world +Z.
        quat = _target_quat([0, 0, 0], [0, 0, 1])
        assert quat is not None
        norm = sum(c * c for c in quat) ** 0.5
        assert norm == pytest.approx(1.0, abs=1e-6)

    @pytest.mark.parametrize(
        ("pos", "target"),
        [
            ([0.2, 0.0, 0.5], [0.2, 0.0, 0.0]),  # top-down looking down
            ([0.0, 0.0, 0.0], [0.0, 0.0, 1.0]),  # straight up
        ],
    )
    def test_vertical_camera_forward_points_at_target(self, pos, target):
        """A vertical look-at must orient the camera's -Z toward the target."""
        quat = _target_quat(pos, target)
        assert quat is not None

        spec = mujoco.MjSpec()
        spec.worldbody.add_camera(name="c", pos=pos, quat=quat, fovy=60, mode=mujoco.mjtCamLight.mjCAMLIGHT_FIXED)
        model = spec.compile()
        rot = model.cam_mat0[0].reshape(3, 3)
        cam_forward = -rot[:, 2]
        expected = np.array(target, dtype=float) - np.array(pos, dtype=float)
        expected /= np.linalg.norm(expected)
        assert np.allclose(cam_forward, expected, atol=1e-3)

    def test_quat_rotates_camera_to_face_target(self):
        """Compile a spec with quat=, then check the resulting cam_mat0 actually
        points the camera's -Z axis toward the target.
        """
        pos = [1.0, 0.0, 0.5]
        target = [0.0, 0.0, 0.0]
        quat = _target_quat(pos, target)
        assert quat is not None

        spec = mujoco.MjSpec()
        spec.worldbody.add_camera(name="c", pos=pos, quat=quat, fovy=60, mode=mujoco.mjtCamLight.mjCAMLIGHT_FIXED)
        model = spec.compile()
        rot = model.cam_mat0[0].reshape(3, 3)
        # Camera forward is -Z in its local frame. Column 2 of rot is world
        # frame +Z of the camera, so camera forward = -rot[:, 2].
        cam_forward = -rot[:, 2]
        expected = np.array([-1.0, 0.0, -0.5])
        expected /= np.linalg.norm(expected)
        assert np.allclose(cam_forward, expected, atol=1e-3)


# SpecBuilder.build


@pytest.fixture
def sample_world() -> SimWorld:
    w = SimWorld()
    w.objects["cube"] = SimObject(
        name="cube",
        shape="box",
        position=[0, 0, 0.1],
        size=[0.1, 0.1, 0.1],
        color=[0.5, 0.5, 0.5, 1],
        is_static=False,
        mass=0.2,
    )
    w.objects["ball"] = SimObject(
        name="ball",
        shape="sphere",
        position=[0.5, 0, 0.1],
        size=[0.05, 0.05, 0.05],
        color=[1, 0, 0, 1],
        is_static=False,
        mass=0.1,
    )
    w.cameras["front"] = SimCamera(
        name="front",
        position=[1, 0, 0.5],
        target=[0, 0, 0],
        fov=60,
        width=640,
        height=480,
    )
    return w


class TestBuild:
    def test_empty_world_compiles(self):
        spec = SpecBuilder.build(SimWorld())
        model = spec.compile()
        assert model.nbody >= 1  # world + ground + lights

    def test_gravity_and_timestep_propagate(self):
        w = SimWorld()
        w.gravity = [0.0, 0.0, -5.0]
        w.timestep = 0.004
        spec = SpecBuilder.build(w)
        model = spec.compile()
        assert model.opt.timestep == pytest.approx(0.004)
        assert np.allclose(model.opt.gravity, [0.0, 0.0, -5.0])

    def test_sample_world_compiles(self, sample_world):
        model = SpecBuilder.build(sample_world).compile()
        # 1 world + 1 ground is a geom; cube, ball are bodies
        assert model.nbody >= 3
        assert model.ncam == 1
        assert model.nu == 0

    def test_body_positions_and_masses(self, sample_world):
        model = SpecBuilder.build(sample_world).compile()
        for name in ("cube", "ball"):
            bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name)
            assert bid >= 0
            obj = sample_world.objects[name]
            assert np.allclose(model.body_pos[bid], obj.position)
            assert model.body_mass[bid] == pytest.approx(obj.mass, abs=1e-6)

    def test_ground_plane_present(self):
        w = SimWorld()
        spec = SpecBuilder.build(w)
        model = spec.compile()
        gid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "ground")
        assert gid >= 0

    def test_ground_plane_absent_when_disabled(self):
        w = SimWorld()
        w.ground_plane = False
        spec = SpecBuilder.build(w)
        model = spec.compile()
        gid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "ground")
        assert gid < 0

    def test_headlight_is_dimmed(self):
        """GH #373 (renderings too bright): MuJoCo's default headlight
        (diffuse 0.4, specular 0.5, active) stacks additively on our two
        explicit scene lights, washing out the scene. build() must dim it
        to a low, shadow-free term."""
        model = SpecBuilder.build(SimWorld()).compile()
        hl = model.vis.headlight
        # Specular off entirely (no head-on glare hotspots).
        assert np.allclose(hl.specular, [0.0, 0.0, 0.0])
        # Diffuse pulled well below the 0.4 default so explicit lights
        # provide the directional illumination.
        assert float(hl.diffuse[0]) <= 0.25
        assert np.allclose(hl.diffuse, [hl.diffuse[0]] * 3)

    def test_explicit_scene_lights_present(self):
        """A single directional ``main_light`` survives build(), matching the
        Menagerie SO101 scene.xml. The previous two-point-light setup
        (main diffuse 1.0 + fill 0.5) over-lit the scene and blew the floor
        to white."""
        model = SpecBuilder.build(SimWorld()).compile()
        assert model.nlight >= 1
        names = {mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_LIGHT, i) for i in range(model.nlight)}
        assert "main_light" in names
        # main_light must be directional (not a point light firing straight down).
        mid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, "main_light")
        assert int(model.light_type[mid]) == int(mujoco.mjtLightType.mjLIGHT_DIRECTIONAL)


# Mutation helpers


class TestMutation:
    def test_add_object_then_recompile(self):
        w = SimWorld()
        spec = SpecBuilder.build(w)
        model = spec.compile()
        data = mujoco.MjData(model)

        SpecBuilder.add_object(
            spec,
            SimObject(
                name="ball",
                shape="sphere",
                position=[0, 0, 0.1],
                size=[0.05, 0.05, 0.05],
                color=[1, 0, 0, 1],
                is_static=False,
                mass=0.1,
            ),
        )
        new_model, _new_data = spec.recompile(model, data)
        ball_id = mujoco.mj_name2id(new_model, mujoco.mjtObj.mjOBJ_BODY, "ball")
        assert ball_id >= 0

    def test_remove_body_then_recompile(self):
        w = SimWorld()
        w.objects["victim"] = SimObject(
            name="victim",
            shape="box",
            position=[0, 0, 0.1],
            size=[0.1, 0.1, 0.1],
            color=[1, 0, 0, 1],
            is_static=False,
            mass=0.1,
        )
        spec = SpecBuilder.build(w)
        model = spec.compile()
        data = mujoco.MjData(model)
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "victim") >= 0

        assert SpecBuilder.remove_body(spec, "victim") is True
        new_model, _ = spec.recompile(model, data)
        assert mujoco.mj_name2id(new_model, mujoco.mjtObj.mjOBJ_BODY, "victim") < 0

    def test_remove_missing_body_returns_false(self):
        spec = SpecBuilder.build(SimWorld())
        assert SpecBuilder.remove_body(spec, "ghost") is False

    def test_remove_body_added_since_the_last_compile(self):
        """A body inserted after the last compile is still removable.

        ``spec.body(name)`` resolves only names present at the last
        ``compile()``/``recompile()``, so a just-added body is invisible to it
        and reachable only through the ``spec.bodies`` enumeration. Without the
        fallback scan, rolling back a half-added body removed nothing and the
        orphan made every later recompile fail on the duplicate name.
        """
        spec = SpecBuilder.build(SimWorld())
        model = spec.compile()
        data = mujoco.MjData(model)

        spec.worldbody.add_body(name="latecomer")
        assert spec.body("latecomer") is None, "precondition: invisible to the typed accessor"

        assert SpecBuilder.remove_body(spec, "latecomer") is True
        assert [b.name for b in spec.bodies] == ["world"]
        # The spec still compiles, and the name is free for a real object.
        spec.recompile(model, data)

    def test_add_camera_then_recompile(self):
        w = SimWorld()
        spec = SpecBuilder.build(w)
        model = spec.compile()
        data = mujoco.MjData(model)

        SpecBuilder.add_camera(
            spec,
            SimCamera(
                name="top",
                position=[0, 0, 2.0],
                target=[0, 0, 0],
                fov=60,
                width=640,
                height=480,
            ),
        )
        new_model, _ = spec.recompile(model, data)
        assert mujoco.mj_name2id(new_model, mujoco.mjtObj.mjOBJ_CAMERA, "top") >= 0

    def test_add_camera_mounted_on_body(self):
        """GH #373 (SO101 wrist cam): a camera with parent_body set must be
        attached to that body so it tracks the body's motion (realistic
        wrist/gripper camera), not fixed in the world."""
        # Build a world with a single movable body to mount onto.
        w = SimWorld()
        w.objects["holder"] = SimObject(
            name="holder", shape="box", position=[0.2, 0.0, 0.3], size=[0.05, 0.05, 0.05], mass=0.2
        )
        spec = SpecBuilder.build(w)
        model = spec.compile()
        data = mujoco.MjData(model)

        SpecBuilder.add_camera(
            spec,
            SimCamera(
                name="wrist",
                position=[0.0, 0.0, 0.1],
                target=[0.0, 0.0, 0.0],
                fov=60,
                width=320,
                height=240,
                parent_body="holder",
            ),
        )
        new_model, _ = spec.recompile(model, data)
        cam_id = mujoco.mj_name2id(new_model, mujoco.mjtObj.mjOBJ_CAMERA, "wrist")
        assert cam_id >= 0
        # The camera's parent body must be 'holder', not worldbody (id 0).
        body_id = int(new_model.cam_bodyid[cam_id])
        assert mujoco.mj_id2name(new_model, mujoco.mjtObj.mjOBJ_BODY, body_id) == "holder"

    def test_add_camera_unknown_parent_body_raises(self):
        w = SimWorld()
        spec = SpecBuilder.build(w)
        with pytest.raises(ValueError, match="parent_body"):
            SpecBuilder.add_camera(
                spec,
                SimCamera(name="wrist", position=[0, 0, 0.1], target=[0, 0, 0], parent_body="does_not_exist"),
            )

    def test_remove_camera(self):
        w = SimWorld()
        w.cameras["c"] = SimCamera(name="c", position=[1, 0, 0.5], target=[0, 0, 0], fov=60, width=640, height=480)
        spec = SpecBuilder.build(w)
        model = spec.compile()
        data = mujoco.MjData(model)

        assert SpecBuilder.remove_camera(spec, "c") is True
        new_model, _ = spec.recompile(model, data)
        assert mujoco.mj_name2id(new_model, mujoco.mjtObj.mjOBJ_CAMERA, "c") < 0

    def test_remove_camera_missing_is_safe_no_op(self):
        # Removing a camera that was never added is a documented safe no-op:
        # return False without touching the (still compile-valid) spec.
        w = SimWorld()
        w.cameras["keep"] = SimCamera(
            name="keep", position=[1, 0, 0.5], target=[0, 0, 0], fov=60, width=640, height=480
        )
        spec = SpecBuilder.build(w)

        assert SpecBuilder.remove_camera(spec, "ghost") is False
        # The pre-existing camera survives and the spec still compiles.
        model = spec.compile()
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, "keep") >= 0

    def test_readd_object_after_remove_reuses_material_and_texture_names(self):
        # Regression guard for the re-add-after-remove asset collision.
        #
        # remove_body() deletes only the body; a textured object's
        # ``<name>_mat`` / ``<name>_tex`` assets are left behind (remove_body
        # docstring: actuators/sensors/assets are cleaned up separately). If a
        # later add_object reuses the same object name, MuJoCo would raise
        # "repeated name" for the duplicate material/texture at compile unless
        # add_object first drops the orphaned assets. This pins that cleanup.
        def make_tile() -> SimObject:
            return SimObject(
                name="tile",
                shape="box",
                position=[0, 0, 0.1],
                size=[0.1, 0.1, 0.1],
                color=[1, 0, 0, 1],
                is_static=False,
                mass=0.1,
                material={
                    "builtin": "checker",
                    "rgb1": [0.2, 0.2, 0.2],
                    "rgb2": [0.8, 0.8, 0.8],
                },
            )

        w = SimWorld()
        w.objects["tile"] = make_tile()
        spec = SpecBuilder.build(w)
        spec.compile()

        # remove_body drops the body but leaves the material/texture orphaned.
        assert SpecBuilder.remove_body(spec, "tile") is True
        assert "tile_mat" in [m.name for m in spec.materials]
        assert "tile_tex" in [t.name for t in spec.textures]

        # Re-adding the same-named object must drop the stale assets, not
        # duplicate them, so the spec recompiles cleanly.
        SpecBuilder.add_object(spec, make_tile())
        assert [m.name for m in spec.materials].count("tile_mat") == 1
        assert [t.name for t in spec.textures].count("tile_tex") == 1
        model = spec.compile()
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "tile") >= 0
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MATERIAL, "tile_mat") >= 0


# Mesh-shaped objects + robot-owned camera skip


# A minimal but valid closed tetrahedron OBJ usable as a MuJoCo mesh asset.
_TETRA_OBJ = """\
v 0 0 0
v 0.1 0 0
v 0 0.1 0
v 0 0 0.1
f 1 3 2
f 1 2 4
f 1 4 3
f 2 3 4
"""


@pytest.fixture
def tetra_mesh_path(tmp_path):
    path = tmp_path / "tetra.obj"
    path.write_text(_TETRA_OBJ)
    return str(path)


class TestMeshObjects:
    """``shape="mesh"`` SimObjects: build() must register the mesh asset and
    add_object must reference it by ``meshname`` instead of a primitive size.

    This is the only path that turns a user-supplied mesh file into a geom; a
    regression here silently drops custom-mesh objects from every scene.
    """

    def test_static_mesh_object_registers_asset_and_compiles(self, tetra_mesh_path):
        w = SimWorld()
        w.objects["widget"] = SimObject(
            name="widget",
            shape="mesh",
            mesh_path=tetra_mesh_path,
            position=[0, 0, 0.2],
            is_static=True,
            color=[0.2, 0.6, 0.9, 1.0],
        )
        model = SpecBuilder.build(w).compile()

        # The mesh asset was registered under the documented ``mesh_<name>`` id.
        assert model.nmesh == 1
        assert mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_MESH, 0) == "mesh_widget"
        # The object's body exists and carries a mesh geom (not a primitive).
        bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "widget")
        assert bid >= 0
        gid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "widget_geom")
        assert gid >= 0
        assert int(model.geom_type[gid]) == int(mujoco.mjtGeom.mjGEOM_MESH)
        # A static mesh has no freejoint, so the model has no DOFs.
        assert model.nq == 0

    def test_dynamic_mesh_object_gets_freejoint(self, tetra_mesh_path):
        w = SimWorld()
        w.objects["m"] = SimObject(
            name="m",
            shape="mesh",
            mesh_path=tetra_mesh_path,
            position=[0, 0, 0.3],
            is_static=False,
            mass=0.3,
        )
        model = SpecBuilder.build(w).compile()
        # A dynamic body carries a 7-DOF freejoint (3 translation + 4 quat).
        assert model.nq == 7
        assert model.nmesh == 1


class TestEllipsoidObject:
    """An ellipsoid SimObject compiles to a single ellipsoid geom whose
    half-extents are the per-axis radii from ``_normalize_size``."""

    def test_ellipsoid_object_compiles_with_expected_radii(self):
        w = SimWorld()
        w.objects["egg"] = SimObject(
            name="egg",
            shape="ellipsoid",
            size=[0.1, 0.2, 0.3],
            position=[0, 0, 0.5],
            is_static=True,
        )
        model = SpecBuilder.build(w).compile()
        gid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "egg_geom")
        assert gid >= 0
        assert int(model.geom_type[gid]) == int(mujoco.mjtGeom.mjGEOM_ELLIPSOID)
        # geom_size holds the per-axis radii (full extents halved).
        assert np.allclose(model.geom_size[gid], [0.05, 0.1, 0.15])


class TestRobotOwnedCameraSkip:
    """A camera whose ``origin_robot`` is set was discovered inside a robot's
    URDF and is re-introduced by ``spec.attach``; build() must NOT re-add it at
    the top level or the attach would collide at compile time (see add_robot)."""

    def test_origin_robot_camera_is_skipped_world_camera_kept(self):
        w = SimWorld()
        w.cameras["robotcam"] = SimCamera(name="robotcam", origin_robot="so101")
        w.cameras["worldcam"] = SimCamera(name="worldcam", position=[1, 0, 1], target=[0, 0, 0])
        model = SpecBuilder.build(w).compile()
        names = {mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_CAMERA, i) for i in range(model.ncam)}
        assert "worldcam" in names
        assert "robotcam" not in names


# attach_robot


ARM_XML = """
<mujoco model="arm">
  <compiler angle="radian"/>
  <worldbody>
    <body name="base" pos="0 0 0.1">
      <joint name="pan" type="hinge" axis="0 0 1"/>
      <geom type="cylinder" size="0.05 0.05"/>
      <body name="link1" pos="0 0 0.1">
        <joint name="lift" type="hinge" axis="0 1 0"/>
        <geom type="capsule" size="0.03" fromto="0 0 0 0 0 0.2"/>
      </body>
    </body>
  </worldbody>
  <actuator>
    <position name="pan_act" joint="pan" kp="50"/>
    <position name="lift_act" joint="lift" kp="50"/>
  </actuator>
</mujoco>
"""


@pytest.fixture
def arm_path(tmp_path):
    path = tmp_path / "arm.xml"
    path.write_text(ARM_XML)
    return str(path)


class TestAttachRobot:
    def test_attach_prefixes_names(self, arm_path):
        scene = SpecBuilder.build(SimWorld())
        robot = SimRobot(name="arm1", urdf_path=arm_path, position=[0.3, 0, 0])
        joint_names = SpecBuilder.attach_robot(scene, robot, arm_path)
        assert joint_names == ["pan", "lift"]

        model = scene.compile()
        # Prefixed joint lookup
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "arm1/pan") >= 0
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "arm1/lift") >= 0
        # Prefixed actuator lookup
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, "arm1/pan_act") >= 0

    def test_attach_two_of_same_robot(self, arm_path):
        """Two attaches of the same source MJCF - each needs a fresh MjSpec."""
        scene = SpecBuilder.build(SimWorld())
        robot_a = SimRobot(name="armA", urdf_path=arm_path, position=[0.3, 0, 0])
        robot_b = SimRobot(name="armB", urdf_path=arm_path, position=[-0.3, 0, 0])

        SpecBuilder.attach_robot(scene, robot_a, arm_path)
        SpecBuilder.attach_robot(scene, robot_b, arm_path)

        model = scene.compile()
        for prefix in ("armA", "armB"):
            assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, f"{prefix}/pan") >= 0
            assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, f"{prefix}/lift") >= 0

    def test_attach_frame_uses_robot_position(self, arm_path):
        scene = SpecBuilder.build(SimWorld())
        robot = SimRobot(name="arm1", urdf_path=arm_path, position=[0.7, 0.2, 0.0])
        SpecBuilder.attach_robot(scene, robot, arm_path)

        model = scene.compile()
        # The arm1 robot's body root gets child-of-frame semantics - body_pos
        # for arm1/base is relative to the attach frame.
        bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "arm1/base")
        # Walk up the body tree to verify the frame is offset
        # (simpler: compile works at all - spec.attach validates positions)
        assert bid >= 0

    def test_attach_robot_strips_robot_scene_ground_plane(self, tmp_path):
        """A robot scene that ships its own ``floor`` plane must not add a
        second coplanar ground plane (issue #320).

        Many menagerie scenes (e.g. franka_emika_panda/scene.xml) include a
        ``floor`` plane at z=0. Attached alongside the world's own ``ground``
        plane (also z=0) it caused two coplanar infinite planes with different
        materials -> depth-buffer Z-fighting / broken floor render.
        ``attach_robot`` now strips plane geoms from the robot scene so exactly
        one world-owned ``ground`` plane survives.
        """
        robot_with_floor = """
        <mujoco model="arm_with_floor">
          <compiler angle="radian"/>
          <worldbody>
            <geom name="floor" size="0 0 0.05" type="plane"/>
            <body name="base" pos="0 0 0.1">
              <joint name="pan" type="hinge" axis="0 0 1"/>
              <geom type="cylinder" size="0.05 0.05"/>
            </body>
          </worldbody>
          <actuator>
            <position name="pan_act" joint="pan" kp="50"/>
          </actuator>
        </mujoco>
        """
        path = tmp_path / "arm_with_floor.xml"
        path.write_text(robot_with_floor)

        scene = SpecBuilder.build(SimWorld())  # ground_plane=True by default
        robot = SimRobot(name="arm1", urdf_path=str(path), position=[0.0, 0.0, 0.0])
        SpecBuilder.attach_robot(scene, robot, str(path))

        model = scene.compile()
        plane_labels = [
            mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, g)
            for g in range(model.ngeom)
            if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE
        ]
        assert len(plane_labels) == 1, f"expected exactly one ground plane, got {plane_labels}"
        assert plane_labels == ["ground"], f"the surviving plane must be the world ground, got {plane_labels}"

        # The robot's own body/joint/actuator must still be present.
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "arm1/pan") >= 0
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, "arm1/pan_act") >= 0

    def test_attach_robot_without_floor_is_unchanged(self, arm_path):
        """A robot scene with no plane geom attaches normally - the strip is a
        no-op and the single world ``ground`` plane survives."""
        scene = SpecBuilder.build(SimWorld())
        robot = SimRobot(name="arm1", urdf_path=arm_path, position=[0.0, 0.0, 0.0])
        SpecBuilder.attach_robot(scene, robot, arm_path)
        model = scene.compile()
        planes = [
            mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, g)
            for g in range(model.ngeom)
            if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE
        ]
        assert planes == ["ground"]

    def test_attach_robot_keeps_floor_when_world_ground_disabled(self, tmp_path):
        """#363 guard 1: under ``ground_plane=False`` the world owns no ground
        plane, so the robot's own floor must NOT be stripped -- otherwise the
        opt-out scene is left with zero ground planes."""
        robot_with_floor = """
        <mujoco model="arm_with_floor">
          <compiler angle="radian"/>
          <worldbody>
            <geom name="floor" size="0 0 0.05" type="plane"/>
            <body name="base" pos="0 0 0.1">
              <joint name="pan" type="hinge" axis="0 0 1"/>
              <geom type="cylinder" size="0.05 0.05"/>
            </body>
          </worldbody>
        </mujoco>
        """
        path = tmp_path / "arm_with_floor.xml"
        path.write_text(robot_with_floor)

        scene = SpecBuilder.build(SimWorld(ground_plane=False))
        robot = SimRobot(name="arm1", urdf_path=str(path), position=[0.0, 0.0, 0.0])
        SpecBuilder.attach_robot(scene, robot, str(path))
        model = scene.compile()
        planes = [
            mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, g)
            for g in range(model.ngeom)
            if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE
        ]
        # The robot floor is the ONLY floor source -- it must survive.
        assert len(planes) == 1, f"opt-out scene must keep the robot floor, got {planes}"

    def test_attach_robot_logs_debug_on_strip(self, tmp_path, caplog):
        """#363 guard 2: the plane strip emits a DEBUG record naming what was
        removed so a disappearing floor is diagnosable."""
        import logging

        robot_with_floor = """
        <mujoco model="arm_with_floor">
          <compiler angle="radian"/>
          <worldbody>
            <geom name="floor" size="0 0 0.05" type="plane"/>
            <body name="base" pos="0 0 0.1">
              <joint name="pan" type="hinge" axis="0 0 1"/>
              <geom type="cylinder" size="0.05 0.05"/>
            </body>
          </worldbody>
        </mujoco>
        """
        path = tmp_path / "arm_with_floor.xml"
        path.write_text(robot_with_floor)
        scene = SpecBuilder.build(SimWorld())
        robot = SimRobot(name="arm1", urdf_path=str(path), position=[0.0, 0.0, 0.0])
        with caplog.at_level(logging.DEBUG, logger="strands_robots.simulation.mujoco.spec_builder"):
            SpecBuilder.attach_robot(scene, robot, str(path))
        msgs = [r.getMessage() for r in caplog.records if "stripped" in r.getMessage()]
        assert msgs, f"expected a DEBUG record about the stripped plane; got {[r.getMessage() for r in caplog.records]}"

    def test_attach_robot_keeps_non_z0_plane(self, tmp_path):
        """#363 guard 3: a robot plane that is NOT z=0 axis-aligned (e.g. an
        elevated/angled ramp modeled as a plane) must survive the strip; only
        the z=0 ground floor is removed."""
        robot_with_ramp = """
        <mujoco model="arm_with_ramp">
          <compiler angle="radian"/>
          <worldbody>
            <geom name="floor" size="0 0 0.05" type="plane"/>
            <geom name="ramp" size="1 1 0.05" type="plane" pos="0 0 0.5"/>
            <body name="base" pos="0 0 0.1">
              <joint name="pan" type="hinge" axis="0 0 1"/>
              <geom type="cylinder" size="0.05 0.05"/>
            </body>
          </worldbody>
        </mujoco>
        """
        path = tmp_path / "arm_with_ramp.xml"
        path.write_text(robot_with_ramp)
        scene = SpecBuilder.build(SimWorld())
        robot = SimRobot(name="arm1", urdf_path=str(path), position=[0.0, 0.0, 0.0])
        SpecBuilder.attach_robot(scene, robot, str(path))
        model = scene.compile()
        plane_labels = [
            mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, g)
            for g in range(model.ngeom)
            if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE
        ]
        # World ground survives + the elevated ramp survives; the z=0 robot
        # floor is the only plane stripped.
        assert "ground" in plane_labels
        assert any("ramp" in (lbl or "") for lbl in plane_labels), (
            f"the non-z=0 ramp plane must survive the strip; got {plane_labels}"
        )
        assert not any(lbl and lbl.endswith("floor") for lbl in plane_labels), (
            f"the z=0 robot floor should have been stripped; got {plane_labels}"
        )


# from_mjcf_string / from_file


class TestFromSources:
    def test_from_mjcf_string_parses(self):
        spec = SpecBuilder.from_mjcf_string(
            '<mujoco><worldbody><body name="x"><geom size="0.1"/></body></worldbody></mujoco>'
        )
        model = spec.compile()
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "x") >= 0

    def test_from_mjcf_string_raises_on_invalid(self):
        with pytest.raises(ValueError):
            SpecBuilder.from_mjcf_string("not valid xml at all")

    def test_from_file_reads_mjcf(self, tmp_path):
        p = tmp_path / "scene.xml"
        p.write_text('<mujoco><worldbody><body name="y"><geom size="0.1"/></body></worldbody></mujoco>')
        spec = SpecBuilder.from_file(str(p))
        model = spec.compile()
        assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "y") >= 0


# MjSpec accessor-drift resilience (defensive branches)


class _RecordingBody:
    """A minimal ``MjsBody``-like stub that records ``add_camera`` calls.

    Only the attributes ``SpecBuilder.add_camera`` touches are implemented:
    a ``name`` and an ``add_camera(**kwargs)`` sink. This lets the fallback
    body-scan branch be asserted by its observable effect (the camera landing
    on the intended parent) rather than by inspecting internal state.
    """

    def __init__(self, name: str) -> None:
        self.name = name
        self.added_cameras: list[dict] = []

    def add_camera(self, **kwargs) -> None:
        self.added_cameras.append(kwargs)


class _StubSpec:
    """A tiny ``MjSpec``-like stub with configurable accessor behaviour.

    ``SpecBuilder``'s mutation helpers defensively handle two MjSpec API
    shapes that differ across mujoco versions: a typed accessor
    (``spec.body(name)`` / ``spec.mesh(name)``) that either returns ``None``
    for a missing name OR raises ``KeyError``/``ValueError``. Current mujoco
    returns ``None``, so the raising branch is never taken with a real spec;
    this stub drives it deterministically.

    Args:
        bodies: bodies enumerable via ``spec.bodies`` (the fallback scan).
        body_raises: if True, ``body()`` raises ``KeyError`` (drift shape).
        body_hidden: if True, ``body()`` always returns ``None`` even for a
            name present in ``bodies`` - mirrors a post-``spec.attach()`` body
            that the typed accessor cannot resolve yet but which is enumerable.
        mesh_raises: if True, ``mesh()`` raises ``ValueError`` (drift shape).
        cameras: value returned for the ``cameras`` attribute (``None`` models
            an MjSpec build that never exposes the camera list).
    """

    def __init__(
        self,
        *,
        bodies=(),
        body_raises: bool = False,
        body_hidden: bool = False,
        mesh_raises: bool = False,
        cameras=(),
    ) -> None:
        self._bodies = list(bodies)
        self._body_raises = body_raises
        self._body_hidden = body_hidden
        self._mesh_raises = mesh_raises
        self.cameras = cameras
        self.worldbody = _RecordingBody("world")
        self.deleted: list = []

    def body(self, name: str):
        if self._body_raises:
            raise KeyError(name)
        if self._body_hidden:
            return None
        for b in self._bodies:
            if b.name == name:
                return b
        return None

    @property
    def bodies(self):
        return self._bodies

    def mesh(self, name: str):
        if self._mesh_raises:
            raise ValueError(name)
        return None

    def delete(self, obj) -> None:  # pragma: no cover - not reached in these paths
        self.deleted.append(obj)


class TestSpecAccessorDriftResilience:
    """SpecBuilder mutation helpers stay safe when the backing MjSpec accessor
    raises ``KeyError``/``ValueError`` instead of returning ``None``, and mount
    a camera on a body that only the enumeration (not the typed accessor)
    can resolve. These defensive branches are dead on mujoco versions whose
    accessors return ``None``, so they are pinned here to guard against a
    well-meaning "remove the unreachable except" regression."""

    def test_remove_body_returns_false_when_accessor_raises(self):
        spec = _StubSpec(body_raises=True)
        assert SpecBuilder.remove_body(spec, "anything") is False
        assert spec.deleted == []

    def test_remove_mesh_returns_false_when_accessor_raises(self):
        spec = _StubSpec(mesh_raises=True)
        assert SpecBuilder.remove_mesh(spec, "mesh_thing") is False
        assert spec.deleted == []

    def test_remove_camera_returns_false_when_camera_list_absent(self):
        spec = _StubSpec(cameras=None)
        assert SpecBuilder.remove_camera(spec, "cam") is False

    def test_add_camera_raises_actionable_error_when_accessor_raises(self):
        spec = _StubSpec(body_raises=True)
        with pytest.raises(ValueError, match="parent_body"):
            SpecBuilder.add_camera(
                spec,
                SimCamera(name="wrist", position=[0, 0, 0.1], target=[0, 0, 0], parent_body="ghost"),
            )

    def test_add_camera_mounts_on_body_found_only_by_fallback_scan(self):
        gripper = _RecordingBody("so101/gripper")
        # body_hidden: the typed accessor cannot resolve the parent (as with a
        # freshly attached robot body) but it is present in ``spec.bodies``.
        spec = _StubSpec(bodies=[gripper], body_hidden=True)
        SpecBuilder.add_camera(
            spec,
            SimCamera(name="wrist", position=[0, 0, 0.1], target=[0, 0, 0], parent_body="so101/gripper"),
        )
        assert len(gripper.added_cameras) == 1
        assert gripper.added_cameras[0]["name"] == "wrist"
        # The camera went to the scanned parent, not the worldbody fallback.
        assert spec.worldbody.added_cameras == []


class TestPublicApiSurface:
    """``__all__`` declares the public API; private helpers must stay out of it.

    The underscore-prefixed module helpers (``_geom_type`` et al.) are internal
    implementation detail consumed by sibling modules via explicit-name imports,
    not through ``from spec_builder import *``. Listing them in ``__all__`` both
    contradicts their private naming and pollutes the documented public surface.
    """

    def test_all_declares_only_public_names(self):
        from strands_robots.simulation.mujoco import spec_builder

        assert spec_builder.__all__ == ["SpecBuilder"]
        assert all(not name.startswith("_") for name in spec_builder.__all__)

    def test_private_helpers_remain_importable_by_name(self):
        # Removing them from __all__ must not remove them as module attributes;
        # scene_ops and simulation import them by explicit name.
        from strands_robots.simulation.mujoco import spec_builder

        for name in (
            "_is_z0_ground_plane",
            "_geom_type",
            "_normalize_size",
            "_validate_size",
            "_target_quat",
        ):
            assert callable(getattr(spec_builder, name)), name
