# =============================================================================
# Dockerfile — DemIA HF Trainer Starter
#
# Two-stage build:
#   1. builder  — installs all Python dependencies.
#   2. runtime  — minimal image with only installed packages.
#
# Build (CPU):
#   docker build -t demia-hf-trainer-starter .
#
# Build (GPU, CUDA 12.1):
#   docker build --build-arg BASE_IMAGE=pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime \
#                -t demia-hf-trainer-starter:gpu .
#
# Run training:
#   docker run --rm -v $(pwd)/outputs:/app/outputs \
#              -e HF_TOKEN=<your-token> \
#              demia-hf-trainer-starter python scripts/run.py train
#
# Run API server:
#   docker run --rm -p 8000:8000 demia-hf-trainer-starter python scripts/run.py serve
#
# Run Gradio UI:
#   docker run --rm -p 7860:7860 demia-hf-trainer-starter python app/app.py
# =============================================================================

ARG BASE_IMAGE=python:3.11-slim
FROM ${BASE_IMAGE} AS builder

WORKDIR /build

# System deps: git (for Hub SSH cloning), libgomp (OpenMP for PyTorch CPU).
RUN apt-get update && apt-get install -y --no-install-recommends \
        git \
        libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# Upgrade pip and install build tools.
RUN pip install --upgrade pip setuptools wheel --no-cache-dir

# Install CPU-only PyTorch first (keeps the image reasonable in size).
RUN pip install torch --index-url https://download.pytorch.org/whl/cpu --no-cache-dir

# Copy and install project requirements.
COPY requirements.txt .
RUN pip install -r requirements.txt --no-cache-dir

# =============================================================================
# Runtime stage
# =============================================================================
FROM ${BASE_IMAGE} AS runtime

LABEL org.opencontainers.image.title="DemIA HF Trainer Starter" \
      org.opencontainers.image.description="Reusable Hugging Face Trainer template for DemIA Living Lab." \
      org.opencontainers.image.source="https://github.com/demia-usal/hf-trainer-starter" \
      org.opencontainers.image.licenses="MIT"

WORKDIR /app

# Copy installed packages from builder.
COPY --from=builder /usr/local/lib/python3.11 /usr/local/lib/python3.11
COPY --from=builder /usr/local/bin /usr/local/bin

# System runtime libs.
RUN apt-get update && apt-get install -y --no-install-recommends \
        libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# Copy source code.
COPY . .

# HF model cache: mount a volume here to persist downloads across runs.
ENV HF_HOME=/app/.cache/huggingface
RUN mkdir -p /app/.cache/huggingface

# Expose ports for the API server and Gradio UI.
EXPOSE 8000 7860

# Default command: run training.
CMD ["python", "scripts/run.py", "train"]
