name: CI

on:
  push:
  pull_request:
  workflow_dispatch:

jobs:
  tests:
    name: Lint and tests
    runs-on: ubuntu-24.04
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"
          cache: pip
          cache-dependency-path: requirements-dev.txt
      - name: Install system libraries (CPU OpenGL for the video test)
        run: sudo apt-get update && sudo apt-get install -y --no-install-recommends libosmesa6 libgl1
      - run: pip install -r requirements-dev.txt
      - run: ruff check . && ruff format --check .
      - run: python -m pytest -q

  worker-image:
    name: Build and smoke-test the worker image
    runs-on: ubuntu-24.04
    steps:
      - uses: actions/checkout@v4
      - name: Build
        run: docker build --platform linux/amd64 -f worker/Dockerfile -t rosdiff-worker:ci .
      - name: Run a real job inside the image (CPU rendering; the cloud GPU uses EGL)
        run: |
          mkdir -p ci-out && chmod 777 ci-out
          docker run --rm -i -e MUJOCO_GL=osmesa -e PYOPENGL_PLATFORM=osmesa -e ROSDIFF_LOCAL_STORAGE=/out \
            -v "$PWD/ci-out:/out" rosdiff-worker:ci python3 - <<'EOF'
          import json
          from worker.handler import handler

          scene = """<mujoco model="ci smoke">
            <include file="g1.xml"/>
            <worldbody>
              <light name="sun" pos="0 0 4" dir="0 0 -1"/>
              <geom name="floor" type="plane" size="0 0 0.05"/>
              <geom name="step" type="box" size="0.15 0.5 0.075" pos="1.65 0 0.075"/>
              <body name="crate" pos="1 1 0.3"><freejoint/><geom type="box" size="0.2 0.2 0.2" mass="3"/></body>
            </worldbody>
          </mujoco>"""
          out = handler({"input": {"run_id": "ci1", "mjcf": scene, "robot": "unitree_g1", "duration_s": 2}})
          print(json.dumps(out, indent=2)[:3000])
          assert "error" not in out, out
          s = out["summary"]
          assert s["video"] and not s["render_error"], s
          assert s["simulated_seconds"] >= 1.99 and s["fell"] is False, s
          assert set(out["files"]) == {"mcap", "video", "video_webm", "summary", "scene", "trajectory"}, out["files"]
          EOF
          ls -la ci-out/runs/ci1
      - name: Check the handler rejects unsafe scenes
        run: |
          docker run --rm -e ROSDIFF_LOCAL_STORAGE=/tmp/o rosdiff-worker:ci python3 -c "
          from worker.handler import handler
          out = handler({'input': {'run_id': 'bad', 'mjcf': '<mujoco><include file=\"/etc/passwd\"/></mujoco>', 'robot': 'unitree_g1'}})
          assert 'scene rejected' in out['error'], out
          print('rejected:', out['error'])"
      - name: Robustness sweep inside the image (MuJoCo CPU and MuJoCo Warp on the CPU device)
        run: |
          docker run --rm -i -e MUJOCO_GL=osmesa rosdiff-worker:ci python3 - <<'EOF'
          from worker.handler import handler

          scene = """<mujoco><include file="go2.xml"/><worldbody><light pos="0 0 4" dir="0 0 -1"/>
            <geom name="floor" type="plane" size="0 0 0.05"/></worldbody></mujoco>"""
          job = {"mjcf": scene, "robot": "unitree_go2", "n_worlds": 3, "duration_s": 1.0, "push_n": [0, 40]}
          from worker.sweep import run_sweep_job
          heights = {}
          for engine in ("mujoco", "warp"):
              out = run_sweep_job(job, engine=engine)
              s = out["summary"]
              print(engine, s["headline"], s.get("engine_note"))
              assert s["worlds"] == 3 and s["fell"] == 0, s
              heights[engine] = [w["min_height"] for w in __import__("json").loads(out["artifacts"]["sweep.json"])["worlds"]]
          print(heights)
          assert all(abs(a - b) < 0.01 for a, b in zip(heights["mujoco"], heights["warp"])), heights
          out = handler({"input": {**job, "function": "sweep"}})
          assert "error" not in out, out
          EOF
      - uses: actions/upload-artifact@v4
        if: always()
        with:
          name: worker-smoke-output
          path: ci-out/runs/ci1
          if-no-files-found: warn

  ros-sandbox:
    name: ROS 2 sandbox (${{ matrix.distro }})
    runs-on: ubuntu-24.04
    strategy:
      fail-fast: false
      matrix:
        distro: [humble, jazzy]
    steps:
      - uses: actions/checkout@v4
      - name: Build
        run: docker build -f worker/ros/Dockerfile --build-arg ROS_DISTRO=${{ matrix.distro }} -t rosdiff-ros:ci .
      - name: Check the sample packages (build, lint, test, launch, bag)
        run: |
          mkdir -p ros-out && chmod 777 ros-out
          docker run --rm -i -v "$PWD/ros-out:/out" rosdiff-ros:ci python3 - <<'EOF'
          import json
          from pathlib import Path
          from worker.ros_runner import check

          def run(name):
              root = Path("/app/samples") / name
              files = {p.relative_to(root).as_posix(): p.read_text() for p in root.rglob("*") if p.is_file()}
              r = check(files, "${{ matrix.distro }}", Path("/out") / name)
              print(f"===== {name}: {r['headline']}")
              print(json.dumps(r, indent=2)[:4000])
              return r

          r = run("good_py")
          assert r["passed"] and r["build"]["ok"] and r["launch"]["ok"], r
          assert "/heartbeat" in r["launch"]["nodes"], r["launch"]
          assert any(t.startswith("/heartbeat ") for t in r["launch"]["topics"]), r["launch"]
          assert r.get("bag_bytes", 0) > 0 and Path("/out/good_py/bag.mcap").is_file(), r

          r = run("good_cpp")
          assert r["passed"] and "/twist_source" in r["launch"]["nodes"], r
          assert r["test"]["ran"] and r["test"]["tests"] > 0, r["test"]

          r = run("broken_cpp")
          assert not r["passed"] and not r["build"]["ok"], r
          assert any("speed" in e for e in r["build"]["errors"]), r["build"]

          r = run("crash_py")
          assert r["build"]["ok"] and not r["launch"]["ok"] and r["launch"]["crashed"], r["launch"]
          print("all sample checks behaved as expected")
          EOF
      - uses: actions/upload-artifact@v4
        if: always()
        with:
          name: ros-sandbox-${{ matrix.distro }}
          path: ros-out
          if-no-files-found: warn

  lerobot-export:
    name: LeRobot dataset export
    runs-on: ubuntu-24.04
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"
      - name: Free disk space (unused preinstalled toolchains)
        run: sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL && df -h /
      - name: Install (CPU-only PyTorch, LeRobot, the simulator)
        run: |
          sudo apt-get update && sudo apt-get install -y --no-install-recommends libosmesa6 libgl1 ffmpeg
          # One resolve with the CPU wheel index, so LeRobot's torch requirement is met without CUDA (~6 GB).
          pip install --extra-index-url https://download.pytorch.org/whl/cpu \
            lerobot==0.4.4 mujoco==3.14.0 foxglove-sdk==0.27.0 imageio-ffmpeg==0.6.0 boto3
          python -c "import torch; assert not torch.version.cuda, torch.__version__; print('torch', torch.__version__)"
      - name: Simulate two runs, export them, load the dataset back
        env:
          MUJOCO_GL: osmesa
        run: |
          git clone -q --depth 1 --filter=blob:none --sparse https://github.com/google-deepmind/mujoco_menagerie.git /tmp/men
          git -C /tmp/men sparse-checkout set unitree_go2
          python - <<'EOF'
          from pathlib import Path
          from api.mjcf_check import prepare_scene_dir
          from worker.lerobot_export import run_export_job
          from worker.sim_runner import config_for, simulate

          scene = """<mujoco><include file="go2.xml"/><worldbody><light pos="0 0 4" dir="0 0 -1"/>
            <geom name="floor" type="plane" size="0 0 0.05"/></worldbody></mujoco>"""
          episodes = []
          for i, controller in enumerate(("hold", "passive")):
              d = prepare_scene_dir(scene, "unitree_go2", Path("/tmp/men"))
              cfg = config_for("unitree_go2", 1.0, controller)
              cfg.width, cfg.height = 320, 180
              simulate(d / "scene.xml", Path(f"/tmp/run{i}"), cfg)
              episodes.append({"run_id": f"r{i}", "task": f"Go2 {controller}",
                               "trajectory_url": f"/tmp/run{i}/trajectory.npz", "video_url": f"/tmp/run{i}/video.mp4"})
          out = run_export_job({"repo_id": "ci/rosdiff-go2", "robot": "unitree_go2", "episodes": episodes})
          print(out["summary"])
          assert out["summary"]["episodes"] == 2, out["summary"]
          Path("/tmp/ds.zip").write_bytes(out["artifacts"]["dataset.zip"])

          import shutil
          shutil.unpack_archive("/tmp/ds.zip", "/tmp/ds")
          from lerobot.datasets.lerobot_dataset import LeRobotDataset
          ds = LeRobotDataset("ci/rosdiff-go2", root="/tmp/ds")
          item = ds[len(ds) - 1]
          print({k: tuple(v.shape) if hasattr(v, "shape") else v for k, v in item.items()})
          assert ds.num_episodes == 2 and item["task"] == "Go2 passive"
          assert tuple(item["observation.images.main"].shape) == (3, 360, 640)
          assert tuple(item["observation.state"].shape) == (19,) and tuple(item["action"].shape) == (12,)
          EOF