# Builder stage
ARG CUDA_VERSION=12.8.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04 AS builder

ENV DEBIAN_FRONTEND=noninteractive
ENV HOME=/workspace
WORKDIR ${HOME}

# Install build tools, compilers, CMake, vcpkg, and project dependencies
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        build-essential \
        ca-certificates \
        curl \
        git \
        gnupg2 \
        libglu1-mesa-dev \
        libxcursor-dev \
        libxinerama-dev \
        locales \
        lsb-release \
        ninja-build \
        pkg-config \
        python3 \
        python3-dev \
        python3-pip \
        software-properties-common \
        sudo \
        unzip \
        wget \
        xorg-dev \
        zip && \
    add-apt-repository ppa:ubuntu-toolchain-r/test && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
        gcc-14 \
        g++-14 \
        gfortran-14 && \
    update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 60 && \
    update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 60 && \
    update-alternatives --install /usr/bin/gfortran gfortran /usr/bin/gfortran-14 60 && \
    update-alternatives --set g++ /usr/bin/g++-14 && \
    update-alternatives --set gcc /usr/bin/gcc-14 && \
    update-alternatives --set gfortran /usr/bin/gfortran-14 && \
    rm -rf /var/lib/apt/lists/*

# Bootstrap vcpkg
RUN git clone https://github.com/microsoft/vcpkg.git ${HOME}/vcpkg && \
    cd ${HOME}/vcpkg && ./bootstrap-vcpkg.sh -disableMetrics

ENV VCPKG_ROOT=${HOME}/vcpkg
ENV PATH=${VCPKG_ROOT}:$PATH

# Install specific CMake version
RUN wget https://github.com/Kitware/CMake/releases/download/v4.0.3/cmake-4.0.3-linux-x86_64.sh && \
    chmod +x cmake-4.0.3-linux-x86_64.sh && \
    ./cmake-4.0.3-linux-x86_64.sh --skip-license --prefix=/usr/local && \
    rm cmake-4.0.3-linux-x86_64.sh

# Fetch Libtorch
WORKDIR ${HOME}/LichtFeld-Studio
RUN wget https://download.pytorch.org/libtorch/cu128/libtorch-cxx11-abi-shared-with-deps-2.7.1%2Bcu128.zip && \
    unzip libtorch-cxx11-abi-shared-with-deps-2.7.1+cu128.zip -d external/ && \
    rm libtorch-cxx11-abi-shared-with-deps-2.7.1+cu128.zip

# Build the project. Disable OpenGL-CUDA Interop since this isn't supported by
# WSL yet https://docs.nvidia.com/cuda/wsl-user-guide/index.html
COPY LichtFeld-Studio ${HOME}/LichtFeld-Studio
RUN cmake -B build \
        -DCMAKE_BUILD_TYPE=Release \
        -DENABLE_CUDA_GL_INTEROP=OFF \
        -G Ninja && \
    cmake --build build -- -j8

# Runtime stage
FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive
ENV HOME=/workspace
WORKDIR ${HOME}

# Install only the runtime dependencies
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        libgl1 \
        libglu1-mesa \
        libgomp1 \
        openssh-client \
        unzip \
        wget \
        x11-apps \
        zenity && \
    rm -rf /var/lib/apt/lists/*

# Copy built artifacts from the builder stage
COPY --from=builder ${HOME}/LichtFeld-Studio/build     ${HOME}/LichtFeld-Studio/build
COPY --from=builder ${HOME}/LichtFeld-Studio/external  ${HOME}/LichtFeld-Studio/external
COPY --from=builder ${HOME}/LichtFeld-Studio/parameter ${HOME}/LichtFeld-Studio/parameter
COPY --from=builder ${HOME}/LichtFeld-Studio/resources ${HOME}/LichtFeld-Studio/resources
COPY --from=builder ${HOME}/LichtFeld-Studio/weights   ${HOME}/LichtFeld-Studio/weights

ENV PATH="${HOME}/LichtFeld-Studio/build:$PATH"

ENTRYPOINT ["/bin/bash"]