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

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# System deps
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 python3.11-venv python3.11-dev python3-pip \
    git wget curl \
    # FreeCAD headless deps
    freecad \
    libgl1-mesa-glx libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1

# Create venv
RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

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

# Install PyG
RUN pip install --no-cache-dir \
    torch-geometric \
    pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv \
    -f https://data.pyg.org/whl/torch-2.4.0+cu124.html

WORKDIR /workspace

# Install project
COPY pyproject.toml .
RUN pip install --no-cache-dir -e ".[train,dev]" || true

COPY . .
RUN pip install --no-cache-dir -e ".[train,dev]"

# -------------------------------------------------------------------
FROM base AS cpu

# CPU-only variant (for CI and non-GPU environments)
FROM python:3.11-slim AS cpu-only

ENV PYTHONUNBUFFERED=1

RUN apt-get update && apt-get install -y --no-install-recommends \
    git freecad libgl1-mesa-glx libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /workspace

COPY pyproject.toml .
RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
RUN pip install --no-cache-dir torch-geometric

COPY . .
RUN pip install --no-cache-dir -e ".[train,dev]"

CMD ["pytest", "tests/", "-v"]