FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    UV_LINK_MODE=copy \
    MAX_JOBS=4

RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        curl \
        ffmpeg \
        git \
        libglib2.0-0 \
        libgl1 \
        libusb-1.0-0-dev \
        python3.12 \
        python3.12-dev \
        python3.12-venv \
    && rm -rf /var/lib/apt/lists/*

RUN curl -LsSf https://astral.sh/uv/install.sh | sh
ENV PATH="/root/.local/bin:/app/cra/.venv/bin:${PATH}"

WORKDIR /app
COPY cra /app/cra
COPY cra-hold.txt /app/cra-hold.txt

RUN uv venv --python python3.12 --seed /app/cra/.venv \
    && uv sync --directory /app/cra --frozen --extra training

COPY armnet-package-requirements.txt /tmp/armnet-package-requirements.txt
RUN printf 'pynput>=0.0.0; sys_platform == "never"\n\
evdev>=0.0.0; sys_platform == "never"\n' > /tmp/overrides.txt \
    && uv pip install --python /app/cra/.venv/bin/python \
        --override /tmp/overrides.txt \
        --constraint /app/cra-hold.txt \
        deepdiff==8.6.2 \
        'datasets>=4,<5' \
        draccus==0.10.0 \
        feetech-servo-sdk==1.0.0 \
        jsonlines==4.0.0 \
        pyserial==3.5 \
        termcolor==3.3.0 \
        -r /tmp/armnet-package-requirements.txt \
    && uv pip install --python /app/cra/.venv/bin/python \
        --no-deps lerobot==0.5.1

COPY armnet-packages /app/armnet-packages
COPY install-armnet-packages.sh /app/install-armnet-packages.sh
RUN PIP="/app/cra/.venv/bin/python -m pip" sh /app/install-armnet-packages.sh

RUN uv pip install --upgrade --python /app/cra/.venv/bin/python \
        --override /tmp/overrides.txt \
        --constraint /app/cra-hold.txt \
        --extra-index-url https://download.pytorch.org/whl/cu128 \
        --index-strategy unsafe-best-match \
        torch==2.7.1 torchvision==0.22.1

COPY cra_eval_runtime.py /app/cra_eval_runtime.py
COPY cra_eval_recorder.py /app/cra_eval_recorder.py
COPY cra_ticket.py /app/cra_ticket.py

ENV PYTHONPATH="/app:/app/cra"

CMD ["/app/cra/.venv/bin/armnet-runtime", "/app/cra_eval_runtime.py"]
