# Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
#
# SPDX-License-Identifier: MIT

ARG UBUNTU_CODENAME=noble
FROM ubuntu:${UBUNTU_CODENAME} AS base
LABEL org.opencontainers.image.source=https://github.com/ROCm/hipFile
ARG UBUNTU_CODENAME
ARG UBUNTU_MAJOR_VERSION=24
ARG UBUNTU_MINOR_VERSION=04
ARG ROCM_VERSION=7.2.2

# If the target ROCm Version on the package repository is non-standard, you can forcibly
# set the ROCm version used in the package repo URL.
# Ex. Package 7.0_alpha installs 7.0.0
ARG ROCM_PKG_REPO_OVERRIDE

# Nightly ROCm build id (YYYYMMDD-<gha_run_id>). When set, ROCm is installed from
# the unsigned nightly repo at rocm.nightlies.amd.com/packages-multi-arch/deb/<ROCM_NIGHTLY_BUILD>
# and ROCM_VERSION is expected to be the clean MAJOR.MINOR.PATCH that build pins.
ARG ROCM_NIGHTLY_BUILD

RUN apt update && \
    apt install -y \
        curl \
        gpg

# Cache MAJOR/MINOR and the ROCm build type for use by later RUN steps.
# Threshold: ROCm 7.11+ ships from repo.amd.com to /opt/rocm/core-<MAJOR>.<MINOR>;
# earlier releases ship from repo.radeon.com to /opt/rocm (legacy symlink).
# A non-empty ROCM_NIGHTLY_BUILD selects the nightly repo and shares the 7.11+
# install layout (ROCM_VERSION must be the clean version that build pins).
RUN ROCM_MAJOR=$(echo "${ROCM_VERSION}" | cut -d. -f1); \
    ROCM_MINOR=$(echo "${ROCM_VERSION}" | cut -d. -f2); \
    if [ -n "${ROCM_NIGHTLY_BUILD}" ]; then \
        ROCM_BUILD_TYPE=nightly; \
        ROCM_INSTALL_PATH="/opt/rocm/core-${ROCM_MAJOR}.${ROCM_MINOR}"; \
    elif [ "${ROCM_MAJOR}" -gt 7 ] || { [ "${ROCM_MAJOR}" -eq 7 ] && [ "${ROCM_MINOR}" -ge 11 ]; }; then \
        ROCM_BUILD_TYPE=new; \
        ROCM_INSTALL_PATH="/opt/rocm/core-${ROCM_MAJOR}.${ROCM_MINOR}"; \
    else \
        ROCM_BUILD_TYPE=legacy; \
        ROCM_INSTALL_PATH="/opt/rocm"; \
    fi; \
    printf 'ROCM_MAJOR=%s\nROCM_MINOR=%s\nROCM_BUILD_TYPE=%s\nROCM_INSTALL_PATH=%s\nROCM_NIGHTLY_BUILD=%s\n' \
        "${ROCM_MAJOR}" "${ROCM_MINOR}" "${ROCM_BUILD_TYPE}" "${ROCM_INSTALL_PATH}" "${ROCM_NIGHTLY_BUILD}" \
        > /etc/rocm-build.env

# Retrieve ROCm GPG key. The nightly repo is unsigned, so no key is fetched
# (the nightly sources.list entry uses [trusted=yes]).
RUN . /etc/rocm-build.env && \
    if [ "${ROCM_BUILD_TYPE}" = "new" ]; then \
        curl https://repo.amd.com/rocm/packages/gpg/rocm.gpg | \
            gpg --dearmor >> /etc/apt/keyrings/amdrocm.gpg; \
    elif [ "${ROCM_BUILD_TYPE}" = "legacy" ]; then \
        curl https://repo.radeon.com/rocm/rocm.gpg.key | \
            gpg --dearmor >> /etc/apt/keyrings/rocm.gpg; \
    fi

# Manually setup ROCm Package Repos
# ROCM_PKG_REPO_VERSION normalizes the legacy ROCm version used in the package repo URL.
# If the PATCH component is 0, the PATCH of the version is dropped. Otherwise no
# modification is made.
RUN . /etc/rocm-build.env && \
    if [ "${ROCM_BUILD_TYPE}" = "new" ]; then \
        echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://repo.amd.com/rocm/packages/ubuntu${UBUNTU_MAJOR_VERSION}${UBUNTU_MINOR_VERSION} stable main" \
            > /etc/apt/sources.list.d/amdrocm.list; \
    elif [ "${ROCM_BUILD_TYPE}" = "nightly" ]; then \
        echo "deb [arch=amd64 trusted=yes] https://rocm.nightlies.amd.com/packages-multi-arch/deb/${ROCM_NIGHTLY_BUILD} stable main" \
            > /etc/apt/sources.list.d/amdrocm-nightly.list; \
    else \
        ROCM_PKG_REPO_VERSION="$(echo ${ROCM_PKG_REPO_OVERRIDE:-$ROCM_VERSION} | sed -E 's|([^.]+\.[^.]+)\.0$|\1|')"; \
        echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_PKG_REPO_VERSION} ${UBUNTU_CODENAME} main" \
            > /etc/apt/sources.list.d/rocm.list; \
    fi

# Apt pinning is required for the legacy repo only; 7.11+ does not need it.
# Outer EOT is a Dockerfile-level heredoc delimiter, kept distinct from the inner shell EOF.
RUN <<"EOT"
. /etc/rocm-build.env
if [ "${ROCM_BUILD_TYPE}" = "legacy" ]; then
    cat <<EOF > /etc/apt/preferences.d/rocm-pin-600
Package: *
Pin: release o=repo.radeon.com
Pin-Priority: 600
EOF
fi
EOT

# Common development dependencies
RUN apt update && \
    apt install -y \
        clang \
        cmake \
        doxygen \
        gdb \
        git \
        libmount-dev \
        llvm-dev \
        python3.12-dev \
        python3.12-venv

# ROCm dev packages. The new and nightly repos both ship the runtime + dev as a
# single amdrocm-runtime-dev<MAJOR>.<MINOR> meta-package.
RUN . /etc/rocm-build.env && \
    if [ "${ROCM_BUILD_TYPE}" = "new" ] || [ "${ROCM_BUILD_TYPE}" = "nightly" ]; then \
        apt install -y \
            amdrocm-runtime-dev${ROCM_MAJOR}.${ROCM_MINOR}; \
    else \
        apt install -y \
            hip-dev \
            rocm-llvm-dev; \
    fi

# Configure system linker for ROCm
RUN . /etc/rocm-build.env && \
    { echo "${ROCM_INSTALL_PATH}/lib"; \
      echo "${ROCM_INSTALL_PATH}/lib64"; } \
        > /etc/ld.so.conf.d/rocm.conf && \
    ldconfig

# Write a /etc/profile.d script exporting ROCM_PATH and PATH for the install layout.
RUN . /etc/rocm-build.env && \
    printf 'export ROCM_PATH=%s\nexport PATH=%s/bin:$PATH\n' \
        "${ROCM_INSTALL_PATH}" "${ROCM_INSTALL_PATH}" \
        > /etc/profile.d/rocm.sh && \
    chmod +x /etc/profile.d/rocm.sh

# /etc/profile already iterates /etc/profile.d/*.sh for login shells. The
# iterator script below applies the same behavior to interactive non-login
# bash (via /etc/bash.bashrc) and non-interactive bash (via BASH_ENV).
RUN printf '%s\n' \
        'for _f in /etc/profile.d/*.sh; do [ -r "$_f" ] && . "$_f"; done' \
        'unset _f' \
        > /etc/bash_env.sh && \
    echo '. /etc/bash_env.sh' >> /etc/bash.bashrc
ENV BASH_ENV="/etc/bash_env.sh"

# Environment Build Arg for ROCm
ENV ROCM_VERSION="${ROCM_VERSION}"

# AIS Code Repository Mount Point
VOLUME /mnt/ais

# Generic Mount point for container to access AIS-capable storage.
VOLUME /mnt/ais-fs

# ---- Stage 2: NVIDIA (extends base with CUDA SDK) ----
# The HIP runtime must be present to use hipFile with NVIDIA GPUs,
# which is why this stage inherits from base rather than a bare CUDA image.
FROM base AS nvidia
ARG UBUNTU_MAJOR_VERSION=24
ARG UBUNTU_MINOR_VERSION=04

# Install a minimal CUDA SDK
RUN curl https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_MAJOR_VERSION}${UBUNTU_MINOR_VERSION}/x86_64/cuda-keyring_1.1-1_all.deb \
    -o cuda-repo-install.deb
RUN apt install ./cuda-repo-install.deb
# https://gitlab.com/nvidia/container-images/cuda/blob/master/dist/12.6.1/ubuntu2404/devel/Dockerfile
# Attempt a "minimal" CUDA Dev install.
RUN apt update && \
    apt install -y \
    cuda-minimal-build-12-9 \
    libcufile-dev-12-9 \
    libnvidia-compute-570-server

# Add CUDA bin directory to the path
ENV PATH=/usr/local/cuda/bin:${PATH}