FROM debian:12

ARG GO_VERSION=1.24.0

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update -qq && apt-get install -y \
    ca-certificates \
    live-build \
    debootstrap \
    squashfs-tools \
    xorriso \
    grub-pc-bin \
    grub-efi-amd64-bin \
    mtools \
    git \
    wget \
    curl \
    tar \
    xz-utils \
    rsync \
    build-essential \
    gcc \
    make \
    perl \
    linux-headers-amd64 \
    && rm -rf /var/lib/apt/lists/*

# Add NVIDIA CUDA repo and install nvcc (needed to compile nccl-tests)
RUN wget -qO /tmp/cuda-keyring.gpg \
        https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/3bf863cc.pub \
    && gpg --dearmor < /tmp/cuda-keyring.gpg \
        > /usr/share/keyrings/nvidia-cuda.gpg \
    && rm /tmp/cuda-keyring.gpg \
    && echo "deb [signed-by=/usr/share/keyrings/nvidia-cuda.gpg] \
https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/ /" \
        > /etc/apt/sources.list.d/cuda.list \
    && apt-get update -qq \
    && apt-get install -y cuda-nvcc-12-8 \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sfn /usr/local/cuda-12.8 /usr/local/cuda

RUN arch="$(dpkg --print-architecture)" \
    && case "$arch" in \
        amd64) goarch=amd64 ;; \
        arm64) goarch=arm64 ;; \
        *) echo "unsupported architecture: $arch" >&2; exit 1 ;; \
    esac \
    && wget -q -O /tmp/go.tar.gz "https://go.dev/dl/go${GO_VERSION}.linux-${goarch}.tar.gz" \
    && rm -rf /usr/local/go \
    && tar -C /usr/local -xzf /tmp/go.tar.gz \
    && rm -f /tmp/go.tar.gz

ENV PATH=/usr/local/go/bin:${PATH}
WORKDIR /work

CMD ["/bin/bash"]