FROM nvidia/cuda:12.4.1-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive

# System deps + Python 3.10 + ffmpeg
RUN apt-get update && apt-get install -y --no-install-recommends \
    git git-lfs wget curl ca-certificates \
    libgl1-mesa-glx libglib2.0-0 \
    software-properties-common \
    && add-apt-repository ppa:deadsnakes/ppa -y \
    && add-apt-repository ppa:ubuntuhandbook1/ffmpeg6 -y \
    && apt-get update \
    && apt-get install -y --no-install-recommends \
    python3.10 python3.10-venv python3.10-dev \
    ffmpeg \
    && git lfs install \
    && rm -rf /var/lib/apt/lists/*

# Make python3.10 the default
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1 \
    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1

# Install pip
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python

# Install PyTorch with CUDA 12.4 support
RUN pip install --no-cache-dir \
    torch torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/cu124

# Clone our repo (includes our modified lerobot)
RUN huggingface-cli download StrongRoboticsLab/pi05-so100-diverse \
    --local-dir /workspace/pi05-so100-diverse \
    --exclude "logs/*"

# Install our modified lerobot + other deps
RUN pip install --no-cache-dir \
    -e /workspace/pi05-so100-diverse/lerobot \
    "transformers==4.54.1" \
    "accelerate>=0.34" \
    wandb \
    huggingface_hub \
    hf_xet

# Install Node.js + Claude Code
RUN curl -fsSL https://deb.nodesource.com/setup_22.x | bash - \
    && apt-get install -y nodejs \
    && rm -rf /var/lib/apt/lists/* \
    && npm install -g @anthropic-ai/claude-code

WORKDIR /workspace/pi05-so100-diverse

ENTRYPOINT ["/bin/bash", "-c"]
CMD ["bash"]