# ── Lung Cancer Segmentation — Single Container ──────────────────────────────
# HuggingFace Spaces Docker SDK
# Contains: FastAPI + nnU-Net + model checkpoint + DICOM converter
#
# Build args (set as HF Space secrets or pass to docker build):
#   HF_MODEL_REPO  — your HF model repo with the nnU-Net checkpoint
#   HF_TOKEN       — HF token (if model repo is private)
#
# docker build \
#   --build-arg HF_MODEL_REPO=yourname/lung-seg-checkpoint \
#   --build-arg HF_TOKEN=hf_... \
#   -t lung-seg .
# ─────────────────────────────────────────────────────────────────────────────

FROM python:3.10-slim

# ── System deps ───────────────────────────────────────────────────────────────
RUN apt-get update && apt-get install -y --no-install-recommends \
        git curl ca-certificates \
    && rm -rf /var/lib/apt/lists/*

# ── Python deps ───────────────────────────────────────────────────────────────
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# ── nnU-Net environment dirs ──────────────────────────────────────────────────
ENV nnUNet_raw=/nnunet/raw
ENV nnUNet_preprocessed=/nnunet/preprocessed
ENV nnUNet_results=/nnunet/results
RUN mkdir -p /nnunet/raw /nnunet/preprocessed /nnunet/results /outputs /tmp/work

# ── Bake checkpoint into the image at build time ──────────────────────────────
# This means inference containers start fast — no download at runtime.
ARG HF_MODEL_REPO
ARG HF_TOKEN

RUN if [ -n "$HF_MODEL_REPO" ]; then \
      python -c "\
from huggingface_hub import snapshot_download; \
import os; \
snapshot_download( \
    repo_id='${HF_MODEL_REPO}', \
    repo_type='model', \
    local_dir='/nnunet/results', \
    token='${HF_TOKEN}' if '${HF_TOKEN}' else None \
); \
open('/nnunet/results/.checkpoint_ready', 'w').close(); \
print('Checkpoint baked into image.') \
"; \
    else \
      echo "WARNING: HF_MODEL_REPO not set — checkpoint will be downloaded at runtime."; \
    fi

# ── App code ──────────────────────────────────────────────────────────────────
COPY app/ ./app/

# ── Runtime env ───────────────────────────────────────────────────────────────
ENV OUTPUT_DIR=/outputs
ENV DEVICE=cpu
ENV HF_DATASET_REPO=farrell236/NSCLC-Radiomics-NIFTI
# These must match what was used during YOUR training:
ENV NNUNET_TRAINER=nnUNetTrainer_10epochs
ENV NNUNET_PLANS=nnUNetResEncUNetMPlans
ENV NNUNET_FOLD=0
ENV NNUNET_DATASET_ID=1

# HF Spaces runs on port 7860
EXPOSE 7860

CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]