# RosDiff simulation worker image for RunPod serverless (SIM_BACKEND=runpod; Modal uses worker/modal_app.py).
#
#   docker build --platform linux/amd64 -f worker/Dockerfile -t <registry>/rosdiff-worker:0.1.0 .
#   docker push <registry>/rosdiff-worker:0.1.0
#
# Build from the repository root: the worker reuses api/mjcf_check.py, api/robots.py and api/storage.py.
FROM ubuntu:22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    # Headless GPU rendering through the NVIDIA driver's EGL (falls back to OSMesa on failure).
    MUJOCO_GL=egl \
    PYOPENGL_PLATFORM=egl \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=all \
    MENAGERIE_DIR=/opt/mujoco_menagerie

RUN apt-get update && apt-get install -y --no-install-recommends \
        python3 python3-pip git ca-certificates \
        libglvnd0 libgl1 libegl1 libgles2 libosmesa6 \
    && rm -rf /var/lib/apt/lists/* \
    # Tell libglvnd where the NVIDIA EGL implementation is (the library itself is mounted in from the host driver).
    && mkdir -p /usr/share/glvnd/egl_vendor.d \
    && printf '{"file_format_version":"1.0.0","ICD":{"library_path":"libEGL_nvidia.so.0"}}\n' \
        > /usr/share/glvnd/egl_vendor.d/10_nvidia.json

# Robot models (only the ones generated scenes may include).
RUN git clone --depth 1 --filter=blob:none --sparse https://github.com/google-deepmind/mujoco_menagerie.git \
        $MENAGERIE_DIR \
    && cd $MENAGERIE_DIR && git sparse-checkout set unitree_g1 unitree_h1 unitree_go2 franka_emika_panda \
    && rm -rf .git

COPY worker/requirements.txt /tmp/requirements.txt
RUN python3 -m pip install --no-cache-dir --upgrade pip && python3 -m pip install --no-cache-dir -r /tmp/requirements.txt

WORKDIR /app
COPY api/__init__.py api/robots.py api/mjcf_check.py api/storage.py api/
COPY worker/__init__.py worker/sim_runner.py worker/handler.py worker/sweep.py worker/

# Fail the build early if anything cannot be imported or the robot does not load.
RUN MUJOCO_GL=osmesa python3 -c "import mujoco, mujoco_warp, foxglove, imageio_ffmpeg, boto3, runpod; \
from worker.handler import run_job; \
mujoco.MjModel.from_xml_path('$MENAGERIE_DIR/unitree_g1/scene.xml'); print('worker image ok')"

CMD ["python3", "-u", "-m", "worker.handler"]