"""ROSClaw v1.0 Physical Simulation E2E Test — MuJoCo + GPU.

Runs on Dell Precision 7960 Tower with 4× RTX A6000.
Full pipeline with actual physics:
    e-URDF load → MuJoCo simulation → trajectory validation
    → Provider routing → physics execution → Memory recording
"""


import pytest


def _cuda_available() -> bool:
    """Probe CUDA without crashing collection when torch is absent."""
    try:
        import torch
    except ImportError:
        return False
    return bool(torch.cuda.is_available())


class TestV1_0PhysicalSimulation:
    """End-to-end test with actual MuJoCo physics simulation."""

    def test_01_mujoco_environment_creation(self):
        """Create MuJoCo environment from UR5e e-URDF."""
        import mujoco

        xml = """
        <mujoco model="ur5e">
          <compiler angle="radian" meshdir="."/>
          <option timestep="0.002" integrator="RK4"/>
          <worldbody>
            <light diffuse=".5 .5 .5" pos="0 0 3" dir="0 0 -1"/>
            <geom type="plane" size="1 1 0.1" rgba=".9 .9 .9 1"/>
            <body name="base" pos="0 0 0.5">
              <freejoint/>
              <geom type="cylinder" size="0.05 0.1" rgba="0.2 0.5 0.8 1"/>
              <body name="shoulder" pos="0 0 0.1">
                <joint name="shoulder_pan_joint" type="hinge" axis="0 0 1"
                       range="-6.28 6.28" damping="1"/>
                <geom type="capsule" fromto="0 0 0 0 0 0.2" size="0.04" rgba="0.8 0.3 0.2 1"/>
                <body name="upper_arm" pos="0 0 0.2">
                  <joint name="shoulder_lift_joint" type="hinge" axis="0 1 0"
                         range="-6.28 6.28" damping="1"/>
                  <geom type="capsule" fromto="0 0 0 0 0 0.4" size="0.035" rgba="0.8 0.3 0.2 1"/>
                  <body name="forearm" pos="0 0 0.4">
                    <joint name="elbow_joint" type="hinge" axis="0 1 0"
                           range="-3.14 3.14" damping="0.5"/>
                    <geom type="capsule" fromto="0 0 0 0 0 0.3" size="0.03" rgba="0.8 0.3 0.2 1"/>
                    <body name="wrist_1" pos="0 0 0.3">
                      <joint name="wrist_1_joint" type="hinge" axis="0 1 0"
                             range="-6.28 6.28" damping="0.3"/>
                      <geom type="capsule" fromto="0 0 0 0 0 0.1" size="0.025" rgba="0.6 0.6 0.6 1"/>
                      <body name="wrist_2" pos="0 0 0.1">
                        <joint name="wrist_2_joint" type="hinge" axis="0 0 1"
                               range="-6.28 6.28" damping="0.3"/>
                        <geom type="capsule" fromto="0 0 0 0 0 0.08" size="0.02" rgba="0.6 0.6 0.6 1"/>
                        <body name="wrist_3" pos="0 0 0.08">
                          <joint name="wrist_3_joint" type="hinge" axis="0 1 0"
                                 range="-6.28 6.28" damping="0.3"/>
                          <geom type="sphere" size="0.02" rgba="0.9 0.9 0.1 1"/>
                        </body>
                      </body>
                    </body>
                  </body>
                </body>
              </body>
            </body>
          </worldbody>
          <actuator>
            <motor joint="shoulder_pan_joint" gear="100" ctrlrange="-150 150"/>
            <motor joint="shoulder_lift_joint" gear="100" ctrlrange="-150 150"/>
            <motor joint="elbow_joint" gear="80" ctrlrange="-100 100"/>
            <motor joint="wrist_1_joint" gear="30" ctrlrange="-28 28"/>
            <motor joint="wrist_2_joint" gear="30" ctrlrange="-28 28"/>
            <motor joint="wrist_3_joint" gear="20" ctrlrange="-10 10"/>
          </actuator>
        </mujoco>
        """
        model = mujoco.MjModel.from_xml_string(xml)
        data = mujoco.MjData(model)

        assert model.nq == 13
        assert model.nv == 12
        assert model.nu == 6
        assert model.nbody == 8

        for _ in range(100):
            mujoco.mj_step(model, data)

        assert data.time > 0
        print(f"MuJoCo physics: {data.time:.3f}s simulated, {model.nbody} bodies")

    def test_02_joint_trajectory_physics_simulation(self):
        """Simulate a joint trajectory in MuJoCo physics."""
        import mujoco

        xml = """
        <mujoco model="ur5e_simple">
          <option timestep="0.001" gravity="0 0 -9.81"/>
          <worldbody>
            <geom type="plane" size="1 1 0.1" rgba="0.9 0.9 0.9 1"/>
            <body name="base" pos="0 0 1.0">
              <geom type="sphere" size="0.05" mass="1" rgba="0.2 0.5 0.8 1"/>
              <joint name="shoulder_pan" type="hinge" axis="0 0 1" range="-6.28 6.28" damping="2"/>
              <body name="upper_arm" pos="0 0 0">
                <geom type="capsule" fromto="0 0 0 0 0 0.2" size="0.03" mass="0.5" rgba="0.8 0.3 0.2 1"/>
                <joint name="shoulder_lift" type="hinge" axis="0 1 0" pos="0 0 0.2" range="-6.28 6.28" damping="2"/>
                <body name="forearm" pos="0 0 0.2">
                  <geom type="capsule" fromto="0 0 0 0 0 0.15" size="0.025" mass="0.3" rgba="0.8 0.3 0.2 1"/>
                </body>
              </body>
            </body>
          </worldbody>
          <actuator>
            <motor joint="shoulder_pan" gear="10" ctrlrange="-20 20"/>
            <motor joint="shoulder_lift" gear="10" ctrlrange="-20 20"/>
          </actuator>
        </mujoco>
        """
        model = mujoco.MjModel.from_xml_string(xml)
        data = mujoco.MjData(model)

        # Apply small torques to create motion
        data.ctrl[:] = [2.0, -1.0]
        duration_steps = 500
        positions_recorded = []

        for step in range(duration_steps):
            mujoco.mj_step(model, data)
            positions_recorded.append(data.qpos.copy())

        final_pos = positions_recorded[-1]
        # Verify motion occurred (non-zero final position)
        assert abs(final_pos[0]) > 0.01 or abs(final_pos[1]) > 0.01
        assert len(positions_recorded) == duration_steps
        assert data.time > 0
        print(f"Trajectory: {len(positions_recorded)} steps, final: [{final_pos[0]:.3f}, {final_pos[1]:.3f}], time: {data.time:.3f}s")

    def test_03_sandbox_collision_detection(self):
        """Test collision detection in MuJoCo physics."""
        import mujoco

        xml = """
        <mujoco model="collision_test">
          <option timestep="0.002"/>
          <worldbody>
            <geom type="plane" size="2 2 0.1"/>
            <body name="robot" pos="0 0 0.5">
              <freejoint/>
              <geom type="sphere" size="0.1" rgba="0.2 0.5 0.8 1"/>
            </body>
            <body name="obstacle" pos="0.3 0 0.5">
              <geom type="sphere" size="0.15" rgba="0.8 0.2 0.2 1"/>
            </body>
          </worldbody>
        </mujoco>
        """
        model = mujoco.MjModel.from_xml_string(xml)
        data = mujoco.MjData(model)
        data.qvel[0] = 0.5

        collision_detected = False
        for _ in range(1000):
            mujoco.mj_step(model, data)
            if data.ncon > 0:
                collision_detected = True
                break

        assert collision_detected or data.qvel[0] < 0.1
        print(f"Collision: detected={collision_detected}, contacts={data.ncon}")

    @pytest.mark.skipif(
        not _cuda_available(),
        reason="CUDA not available on this machine",
    )
    def test_04_gpu_acceleration_available(self):
        """Verify GPU acceleration for MuJoCo and PyTorch."""
        import torch

        assert torch.cuda.is_available()
        assert torch.cuda.device_count() >= 1

        device = torch.device("cuda:0")
        x = torch.randn(1000, 1000, device=device)
        y = torch.matmul(x, x.T)
        assert y.device.type == "cuda"

        gpu_name = torch.cuda.get_device_name(0)
        mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"GPU: {gpu_name}, Memory: {mem:.1f} GB")

    @pytest.mark.asyncio
    async def test_05_provider_inference_with_physics(self):
        """Test provider routing with physics-backed validation."""

        from rosclaw.provider.core.request import ProviderRequest
        from rosclaw.provider.core.response import ProviderResponse
        from rosclaw.provider.core.router import CapabilityRouter
        from rosclaw.provider.core.registry import ProviderRegistry
        from rosclaw.provider.core.provider import Provider
        from rosclaw.provider.core.manifest import ProviderManifest
        from rosclaw.core.event_bus import EventBus

        event_bus = EventBus()
        provider_reg = ProviderRegistry(event_bus=event_bus)

        class PhysicsSkillProvider(Provider):
            name = "physics_skill"
            version = "0.1.0"
            capabilities = ["skill.pick_and_place"]

            async def infer(self, request):
                trajectory = [
                    [0.0, -1.57, 1.57, 0.0, 0.0, 0.0],
                    [0.1, -1.4, 1.4, 0.1, 0.0, 0.0],
                    [0.2, -1.2, 1.2, 0.2, 0.0, 0.0],
                ]
                return ProviderResponse(
                    request_id=request.request_id,
                    provider=self.name,
                    capability=request.capability,
                    result={"trajectory": trajectory, "gripper": "close"},
                    status="ok",
                )

        manifest = ProviderManifest.from_dict({
            "name": "physics_skill",
            "version": "0.1.0",
            "type": "skill",
            "capabilities": ["skill.pick_and_place"],
            "embodiment": {"supported_robots": ["ur5e"]},
            "safety": {"executable": True, "requires_guard": True},
        })
        provider_reg._providers["physics_skill"] = PhysicsSkillProvider(manifest)
        provider_reg._health["physics_skill"] = {"ok": True}
        provider_reg._manifests["physics_skill"] = manifest

        router = CapabilityRouter(provider_reg)

        async def run_test():
            request = ProviderRequest(
                request_id="phys_001",
                capability="skill.pick_and_place",
                inputs={"object": "red_cup", "location": "table_center"},
                context={"robot": "ur5e"},
            )
            return await router.invoke(request)

        response = await run_test()
        assert response.is_ok
        trajectory = response.result["trajectory"]

        for i, point in enumerate(trajectory):
            assert len(point) >= 2
            assert -6.28 <= point[0] <= 6.28
            assert -6.28 <= point[1] <= 6.28

        print(f"Physics-backed provider: {len(trajectory)} waypoints validated")

    def test_06_eurdf_to_mujoco_conversion(self):
        """Load e-URDF and create corresponding MuJoCo model."""
        import xml.etree.ElementTree as ET
        import mujoco

        from rosclaw.runtime.eurdf_loader import EURDFLoader

        loader = EURDFLoader()
        profile = loader.load("ur5e")

        assert profile.robot_id == "universal_robots_ur5e"
        assert profile.embodiment.dof == 6

        mujoco_root = ET.Element("mujoco", {"model": profile.robot_id})
        ET.SubElement(mujoco_root, "compiler", {"angle": "radian"})
        ET.SubElement(mujoco_root, "option", {"timestep": "0.002"})
        worldbody = ET.SubElement(mujoco_root, "worldbody")
        ET.SubElement(worldbody, "geom", {
            "type": "plane", "size": "1 1 0.1", "rgba": "0.9 0.9 0.9 1"
        })

        actuator = ET.SubElement(mujoco_root, "actuator")
        parent_body = worldbody

        for i, joint in enumerate(profile.embodiment.joints):
            body = ET.SubElement(parent_body, "body", {
                "name": joint["name"], "pos": f"0 0 {0.1 + i * 0.05}"
            })
            limits = joint.get("limits", {})
            ET.SubElement(body, "joint", {
                "name": joint["name"],
                "type": "hinge",
                "axis": "0 0 1" if i == 0 else "0 1 0",
                "range": f"{limits.get('lower', -6.28)} {limits.get('upper', 6.28)}",
                "damping": "0.5",
            })
            ET.SubElement(body, "geom", {
                "type": "capsule",
                "fromto": f"0 0 0 0 0 {0.1 + i * 0.02}",
                "size": str(0.04 - i * 0.003),
                "rgba": "0.8 0.3 0.2 1",
            })
            ET.SubElement(actuator, "motor", {
                "joint": joint["name"],
                "gear": str(100 - i * 10),
            })
            parent_body = body

        xml_str = ET.tostring(mujoco_root, encoding="unicode")
        model = mujoco.MjModel.from_xml_string(xml_str)
        data = mujoco.MjData(model)

        assert model.nq >= 6
        assert model.nu == 6

        for _ in range(100):
            mujoco.mj_step(model, data)

        print(f"e-URDF→MuJoCo: {profile.embodiment.dof} DOF, {len(profile.embodiment.joints)} joints")
