# ============================================================================ #
# Copyright (c) 2022 - 2026 NVIDIA Corporation & Affiliates.                   #
# All rights reserved.                                                         #
#                                                                              #
# This source code and the accompanying materials are made available under     #
# the terms of the Apache License 2.0 which accompanies this distribution.     #
# ============================================================================ #

ARG base_image=nvcr.io/nvidia/nightly/cuda-quantum:cu12-latest-base
FROM $base_image

USER root

# Optionally replace the base image documentation with the documentation
# produced by the selected deployment run.
ARG documentation=./documentation
ARG include_documentation=false
COPY "$documentation" "$CUDA_QUANTUM_PATH/publishing-documentation/"
RUN if [ "$include_documentation" = "true" ]; then \
        rm -rf "$CUDA_QUANTUM_PATH/docs" && mkdir -p "$CUDA_QUANTUM_PATH/docs" && \
        cp -r "$CUDA_QUANTUM_PATH/publishing-documentation/." "$CUDA_QUANTUM_PATH/docs"; \
    fi && \
    rm -rf "$CUDA_QUANTUM_PATH/publishing-documentation"

# Install additional runtime dependencies.
RUN cuda_version_suffix=$(echo ${CUDA_VERSION} | tr . -) && \
    for cudart_dependency in libcusolver libcusparse libcublas libcurand cuda-cudart cuda-nvrtc; do \
        if [ -z "$(apt list --installed | grep -o ${cudart_dependency}-${cuda_version_suffix})" ]; then \
            apt-get install -y --no-install-recommends \
                ${cudart_dependency}-${cuda_version_suffix}; \
        fi \
    done && \
    if [ $(echo ${CUDA_VERSION} | cut -d . -f1) -gt 11 ]; then \
        apt-get install -y --no-install-recommends \
            libnvjitlink-${cuda_version_suffix}; \
    fi && \
    # just here for convenience:
    apt-get install -y --no-install-recommends curl jq 
ADD ./requirements.txt /tmp/requirements.txt
RUN if [ -x "$(command -v pip)" ]; then \
        apt-get install -y --no-install-recommends gcc libpython3-dev \
        && pip install --no-cache-dir -r /tmp/requirements.txt; \
        if [ -n "$MPI_ROOT" ]; then \
            pip install --no-cache-dir mpi4py~=4.1; \
        fi; \
    fi \
    && rm /tmp/requirements.txt
# Install CUDA Python packages based on CUDA version
RUN cuda_major_version=$(echo ${CUDA_VERSION} | cut -d . -f1) && \
    if [ "$cuda_major_version" = "12" ]; then \
        pip install --no-cache-dir \
            nvidia-cuda-runtime-cu12 \
            nvidia-cuda-nvrtc-cu12 \
            nvidia-curand-cu12; \
    elif [ "$cuda_major_version" = "13" ]; then \
        pip install --no-cache-dir \
            nvidia-cuda-runtime \
            nvidia-cuda-nvrtc \
            nvidia-curand; \
    else \
        echo "Unsupported CUDA major version: ${cuda_major_version}" && exit 1; \
    fi

# Make sure that apt-get remains updated at the end!;
# If we don't do that, then apt-get will get confused when some CUDA
# components are already installed but not all of them.

# Include helper scripts and configurations to facilitate 
# development with VS Code and JupterLab.
# See also https://github.com/microsoft/vscode/issues/60#issuecomment-161792005
ARG vscode_config=.vscode
COPY "${vscode_config}" /home/cudaq/.vscode
RUN echo -e '#! /bin/bash \n\
    if [ ! -x "$(command -v code)" ]; then \n\
        os=$([ "$(uname -m)" == "aarch64" ] && echo cli-alpine-arm64 || echo cli-alpine-x64) \n\
        curl -Lk "https://code.visualstudio.com/sha/download?build=stable&os=$os" --output vscode_cli.tar.gz \n\
        tar -xf vscode_cli.tar.gz && rm vscode_cli.tar.gz && sudo mv code /usr/bin/ \n\
    fi \n\
    code "$@"' > "$CUDA_QUANTUM_PATH/bin/vscode-setup" \
    && chmod +x "$CUDA_QUANTUM_PATH/bin/vscode-setup"
RUN echo -e '#! /bin/bash \n\
    jupyter-lab --no-browser --ip=* --ServerApp.allow_origin=* --IdentityProvider.token="$@" \n\
    ' > "$CUDA_QUANTUM_PATH/bin/jupyter-lab-setup" \
    && chmod +x "$CUDA_QUANTUM_PATH/bin/jupyter-lab-setup"

RUN chown -R cudaq /home/cudaq && chgrp -R cudaq /home/cudaq
USER cudaq
