# syntax=docker/dockerfile:1

ARG PYTORCH_VERSION=2.5.1
ARG CUDA_VERSION=12.4
ARG CUDNN_VERSION=9
ARG DEEPLABCUT_VERSION=3.0.0

# ── core ──────────────────────────────────────────────────────────────────────
FROM pytorch/pytorch:${PYTORCH_VERSION}-cuda${CUDA_VERSION}-cudnn${CUDNN_VERSION}-runtime AS core

ARG DEEPLABCUT_VERSION
ARG CUDA_VERSION
ARG PYTORCH_VERSION
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update -yy && \
    apt-get install -yy --no-install-recommends \
      libgl1 \
      libglib2.0-0 \
      ffmpeg \
      build-essential \
      git \
    && pip install --upgrade pip \
    && pip install "deeplabcut[modelzoo,wandb]==${DEEPLABCUT_VERSION}" \
    && apt-get purge -yy --auto-remove build-essential git \
    && rm -rf /var/lib/apt/lists/*

# Create pretrained weights directory and make it writable
# (same default as HuggingFaceWeightsMixin in backbones/base.py: Path(__file__).parent / "pretrained_weights")
RUN PRETRAINED=$(python -c "from pathlib import Path; import deeplabcut.pose_estimation_pytorch.models.backbones.base as b; print(Path(b.__file__).parent / 'pretrained_weights')") && \
    mkdir -p "$PRETRAINED" && \
    chmod a+rwx -R "$PRETRAINED"

RUN mkdir -p /app && chmod a+rwx /app

COPY motd.sh /home/motd.sh
RUN echo "source /home/motd.sh" >> /etc/profile

ENV CUDA_VERSION=${CUDA_VERSION}
ENV PYTORCH_VERSION=${PYTORCH_VERSION}
ENV DEEPLABCUT_VERSION=${DEEPLABCUT_VERSION}

WORKDIR /app

# ── jupyter ───────────────────────────────────────────────────────────────────
# (runs as root intentionally; deeplabcut-docker adds a host-matched user at runtime)
FROM core AS jupyter

# Default Jupyter token (override with -e NOTEBOOK_TOKEN=<secret>).
# An empty value (-e NOTEBOOK_TOKEN=) disables token auth — see docker/README.md.
ENV NOTEBOOK_TOKEN=deeplabcut

RUN pip install "notebook<7"
EXPOSE 8888
HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \
    CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8888/api/', timeout=3).read()" || exit 1
ENTRYPOINT exec jupyter notebook --no-browser --NotebookApp.token="${NOTEBOOK_TOKEN}" --ip 0.0.0.0

# ── test (CI only, not published to Hub) ─────────────────────────────────────
FROM core AS test

RUN pip install --no-cache-dir pytest
