# Copyright © Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier:  MIT

# Build type selection - must be declared before first use in FROM
ARG BUILD_TYPE=prebuilt

# Capture the ASAN mode the container is being created with, if provided.
ARG ASAN_MODE=unknown

# ============================================================================
# Common base image with all dependencies
# ============================================================================
FROM ubuntu:24.04 AS base

# hipDNN dev container version (version number set arbitrarily, update if/when needed)
ENV HIPDNN_DEV_CONTAINER_VERSION=1

ENV DEBIAN_FRONTEND=noninteractive
ENV HIP_PLATFORM=amd

RUN apt-get update && apt-get install -y \
    git \
    gdb \
    gfortran \
    git-lfs \
    ninja-build \
    cmake \
    g++ \
    pkg-config \
    xxd \
    patchelf \
    automake \
    python3-venv \
    python3-dev \
    libegl1-mesa-dev \
    curl \
    unzip \
    patchelf \
    libtool \
    python3-pip \
    vim \
    rsync \
    ca-certificates \
    wget \
    texinfo \
    bison \
    flex \
    gh \
    libsqlite3-dev \
    && rm -rf /var/lib/apt/lists/*


# Install DVC. The dvc.org/deb apt repo has been unreachable, so install the
# official .deb package directly. It bundles all remote backends (incl. S3),
# which is required for `dvc pull` from the rocm-libraries S3 remote.
# The only apt dependency (git) is already installed above, so use dpkg
# directly — no apt-get update or list cleanup needed.
RUN curl -fL https://dvc.org/download/linux-deb/dvc-3.67.0 -o /tmp/dvc.deb && \
    dpkg -i /tmp/dvc.deb && \
    rm /tmp/dvc.deb

# ============================================================================
# TheRock build environment - system tools + Python venv ready
#
# A minimal container with everything needed to build TheRock from source.
# Contains all build dependencies from base plus a Python venv ready for
# TheRock's requirements.txt. No source code or build artifacts included.
#
# Usage: Mount your TheRock source and output directory when running.
#
# Example: Build this target
#   docker build -f Dockerfile.ubuntu24 --target fullbuildshell -t hipdnn:fullbuildshell .
#
# Example: Run with mounts and build manually
#   docker run -it \
#       -v /path/to/TheRock:/TheRock \
#       -v /path/to/rocm:/opt/rocm \
#       hipdnn:fullbuildshell bash
#
# Inside the container, run:
#   cd /TheRock
#   pip install -r requirements.txt
#   python ./build_tools/fetch_sources.py --depth 1
#   mkdir -p build && cd build
#   cmake .. -GNinja -DCMAKE_INSTALL_PREFIX=/opt/rocm -DTHEROCK_AMDGPU_FAMILIES=gfx942 ...
#   cmake --build .
#   cmake --install .
# ============================================================================
FROM base AS fullbuildshell

RUN python3 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

WORKDIR /workspace

# ============================================================================
# Prebuilt stage - downloads and extracts dev env prebuilt binaries tarball
#
# New path: uses install_rocm_from_artifacts.py for downloading ROCm artifacts.
# Legacy path: uses wget/tar when deprecated args are set (backward compat).
#
# Example: Default build (latest nightly for gfx94X-dcgpu):
#   docker build -f Dockerfile.ubuntu24 -t hipdnn:prebuilt .
#
# Example: Specific release version:
#   docker build -f Dockerfile.ubuntu24 \
#       --build-arg THEROCK_RELEASE=7.12.0a20260202 \
#       -t hipdnn:prebuilt .
#
# Example: Different ASIC family:
#   docker build -f Dockerfile.ubuntu24 \
#       --build-arg THEROCK_ASIC=gfx90X \
#       --build-arg THEROCK_ASIC_VARIANT=dcgpu \
#       -t hipdnn:prebuilt .
#
# Example: Full artifact group override:
#   docker build -f Dockerfile.ubuntu24 \
#       --build-arg THEROCK_ARTIFACT_GROUP=gfx110X-all \
#       -t hipdnn:prebuilt .
# ============================================================================
FROM base AS install_therock_prebuilt

# Build args for release selection
ARG THEROCK_ASIC=gfx94X
ARG THEROCK_ASIC_VARIANT=dcgpu
ARG THEROCK_RELEASE=latest
# Full artifact group (combines ASIC and variant)
ARG THEROCK_ARTIFACT_GROUP=$THEROCK_ASIC-$THEROCK_ASIC_VARIANT

# Deprecated build args (kept for backward compatibility)
ARG THEROCK_GIT_TAG=
ARG THEROCK_PREBUILT_ID=
ARG THEROCK_TARBALL=

# Clone TheRock for build tools
RUN git clone --depth 1 https://github.com/ROCm/TheRock.git /TheRock

WORKDIR /TheRock

RUN python3 -m venv /venv && /venv/bin/pip install -r requirements.txt

# Install ROCm - two paths:
# 1. Legacy: any deprecated arg set → wget/tar (old method, no translation needed)
# 2. New: install_rocm_from_artifacts.py
# Priority: THEROCK_TARBALL > THEROCK_PREBUILT_ID > THEROCK_GIT_TAG
RUN if [ -n "$THEROCK_TARBALL" ] || [ -n "$THEROCK_PREBUILT_ID" ] || [ -n "$THEROCK_GIT_TAG" ]; then \
        echo "WARNING: Using deprecated build args. Migrate to THEROCK_RELEASE + THEROCK_ARTIFACT_GROUP." >&2; \
        tarball="${THEROCK_TARBALL:-therock-dist-linux-${THEROCK_PREBUILT_ID:-$THEROCK_ARTIFACT_GROUP-$THEROCK_GIT_TAG}.tar.gz}"; \
        echo "Downloading: $tarball"; \
        mkdir -p /opt/rocm && \
        wget --progress=dot:giga \
            "https://therock-nightly-tarball.s3.us-east-2.amazonaws.com/$tarball" \
            -O "/tmp/$tarball" && \
        tar -xf "/tmp/$tarball" -C /opt/rocm --strip-components=1 && \
        rm "/tmp/$tarball"; \
        version=$(echo "$tarball" | sed 's/^therock-dist-linux-//; s/\.tar\.gz$//' | rev | cut -d- -f1 | rev); \
        echo "$version" > /opt/rocm/THEROCK_VERSION; \
    elif [ "$THEROCK_RELEASE" = "latest" ]; then \
        echo "Installing latest nightly for $THEROCK_ARTIFACT_GROUP"; \
        /venv/bin/python build_tools/install_rocm_from_artifacts.py \
            --latest-release \
            --artifact-group "$THEROCK_ARTIFACT_GROUP" \
            --output-dir /opt/rocm \
            2>&1 | tee /tmp/install.log; \
        # Fragile: depends on install_rocm_from_artifacts.py printing
        # "Found latest release: <version>". Falls back if format changes. \
        version=$(grep -oP 'Found latest release: \K\S+' /tmp/install.log) && \
            echo "$version" > /opt/rocm/THEROCK_VERSION || \
            echo "latest (version details unavailable)" > /opt/rocm/THEROCK_VERSION; \
    else \
        echo "Installing release $THEROCK_RELEASE for $THEROCK_ARTIFACT_GROUP"; \
        /venv/bin/python build_tools/install_rocm_from_artifacts.py \
            --release "$THEROCK_RELEASE" \
            --artifact-group "$THEROCK_ARTIFACT_GROUP" \
            --output-dir /opt/rocm \
            2>&1 | tee /tmp/install.log; \
        echo "$THEROCK_RELEASE" > /opt/rocm/THEROCK_VERSION; \
    fi

RUN rm -rf /TheRock /venv /tmp/install.log

# ============================================================================
# Build from source stage - clones and builds TheRock
#
# Example build command for default source build (default branch, gfx94X):
# docker build -f Dockerfile.ubuntu24 --build-arg BUILD_TYPE=fullbuild -t hipdnn:fullbuild .
#
# Example build command for overriding ASIC and git hash:
# docker build -f Dockerfile.ubuntu24 --build-arg BUILD_TYPE=fullbuild --build-arg THEROCK_ASIC=gfx950 --build-arg THEROCK_GIT_HASH=abc123def456 -t hipdnn:custom_source_build .
#
# Build mode options:
# - THEROCK_BUILD_MODE can be: Preset, Debug, or Release (default)
# - When using "Preset", specify the preset with THEROCK_BUILD_PRESET
# - When using "Debug" or "Release", standard CMake build types are used
#
# Note: When THEROCK_BUILD_MODE=Preset, the THEROCK_BUILD_PRESET value is used.
#       When THEROCK_BUILD_MODE=Debug/Release, THEROCK_BUILD_PRESET is ignored.
#
# Build parallelism:
# - BUILD_JOBS controls the number of parallel build jobs
# - Default is 0, which means use all available CPU cores
#
# Example build command for Debug build with 4 cores:
# docker build -f Dockerfile.ubuntu24 --build-arg BUILD_TYPE=fullbuild --build-arg THEROCK_BUILD_MODE=Debug --build-arg BUILD_JOBS=4 -t hipdnn:debug .
#
# Example build command with custom TheRock preset:
# docker build -f Dockerfile.ubuntu24 --build-arg BUILD_TYPE=fullbuild --build-arg THEROCK_BUILD_MODE=Preset --build-arg THEROCK_BUILD_PRESET=linux-debug-package -t hipdnn:custom-preset .
# ============================================================================
FROM base AS install_therock_fullbuild

ARG THEROCK_ASIC=gfx94X
ARG THEROCK_GIT_HASH=default
ARG THEROCK_BUILD_MODE=Release
ARG THEROCK_BUILD_PRESET=linux-release-package
ARG BUILD_JOBS=0
ARG ROCM_LIBRARIES_REF=default
ARG ROCM_SYSTEMS_REF=default
ARG VERBOSE_BUILD=0

RUN git config --global user.email "docker@build.local" && \
    git config --global user.name "Docker Build"

RUN git clone --depth 1 https://github.com/ROCm/TheRock.git /TheRock

WORKDIR /TheRock

RUN if [ "$THEROCK_GIT_HASH" != "default" ]; then \
        echo "Checking out TheRock: $THEROCK_GIT_HASH"; \
        git fetch --depth 1 origin "$THEROCK_GIT_HASH" && \
        git checkout FETCH_HEAD; \
    else \
        echo "Using default branch for TheRock"; \
    fi

# Update rocm-libraries submodule to specific ref if provided.
# Commit the change so that fetch_sources.py does not overwrite it.
RUN if [ "$ROCM_LIBRARIES_REF" != "default" ]; then \
        echo "Updating rocm-libraries submodule to: $ROCM_LIBRARIES_REF"; \
        git submodule update --init --depth 1 rocm-libraries && \
        cd rocm-libraries && \
        git fetch --depth 1 origin "$ROCM_LIBRARIES_REF" && \
        git checkout FETCH_HEAD && \
        cd .. && \
        git add rocm-libraries && \
        git commit -m "Update rocm-libraries submodule to $ROCM_LIBRARIES_REF"; \
    fi

# Update rocm-systems submodule to specific ref if provided.
# Commit the change so that fetch_sources.py does not overwrite it.
RUN if [ "$ROCM_SYSTEMS_REF" != "default" ]; then \
        echo "Updating rocm-systems submodule to: $ROCM_SYSTEMS_REF"; \
        git submodule update --init --depth 1 rocm-systems && \
        cd rocm-systems && \
        git fetch --depth 1 origin "$ROCM_SYSTEMS_REF" && \
        git checkout FETCH_HEAD && \
        cd .. && \
        git add rocm-systems && \
        git commit -m "Update rocm-systems submodule to $ROCM_SYSTEMS_REF"; \
    fi

RUN python3 -m venv /opt/venv

# Adding the venv into the path activates the python virtual env in a way.  We cant source it the
# normal way due to how dockerfiles run
ENV PATH="/opt/venv/bin:$PATH"

RUN pip install -r requirements.txt

RUN git config --global url."https://github.com/".insteadOf git@github.com: && \
    git config --global credential.helper "!f() { echo username=anonymous; echo password=; }; f" && \
    python ./build_tools/fetch_sources.py --depth 1

RUN mkdir build
WORKDIR /TheRock/build

RUN if [ "$VERBOSE_BUILD" = "1" ]; then export THEROCK_INTERACTIVE=1; fi; \
    if [ "$THEROCK_BUILD_MODE" = "Preset" ]; then \
        echo "Building with TheRock preset: $THEROCK_BUILD_PRESET"; \
        cmake .. \
            --preset=$THEROCK_BUILD_PRESET \
            -DCMAKE_INSTALL_PREFIX=/opt/rocm \
            -DTHEROCK_ENABLE_ALL=OFF \
            -DTHEROCK_ENABLE_MIOPEN=ON \
            -DTHEROCK_ENABLE_COMPOSABLE_KERNEL=ON \
            -DTHEROCK_AMDGPU_FAMILIES=$THEROCK_ASIC; \
    else \
        echo "Building with CMake build type: $THEROCK_BUILD_MODE"; \
        cmake .. \
            -GNinja \
            -DCMAKE_BUILD_TYPE=$THEROCK_BUILD_MODE \
            -DCMAKE_INSTALL_PREFIX=/opt/rocm \
            -DTHEROCK_ENABLE_ALL=OFF \
            -DTHEROCK_ENABLE_MIOPEN=ON \
            -DTHEROCK_ENABLE_COMPOSABLE_KERNEL=ON \
            -DTHEROCK_AMDGPU_FAMILIES=$THEROCK_ASIC; \
    fi

RUN if [ "$VERBOSE_BUILD" = "1" ]; then \
        export THEROCK_INTERACTIVE=1; \
        export TEATIME_FORCE_INTERACTIVE=1; \
        echo "Verbose build enabled: subproject output will stream to terminal."; \
    fi; \
    if [ "$BUILD_JOBS" = "0" ]; then \
        cmake --build .; \
    else \
        echo "Building with $BUILD_JOBS parallel jobs"; \
        cmake --build . -j$BUILD_JOBS; \
    fi && \
    cmake --install .

# ============================================================================
# Select which build stage to use
# ============================================================================
FROM install_therock_$BUILD_TYPE AS install_therock

# ============================================================================
# hipDNN dev shell - development tools WITHOUT /opt/rocm baked in
#
# This lightweight image contains all hipDNN development tools but expects
# /opt/rocm to be mounted at runtime. Useful for:
# - Local development with shared ROCm installation
# - Faster container builds (no TheRock build/download)
# - Multiple ASIC targets using the same base image
#
# Example: Build dev shell only
#   docker build -f Dockerfile.ubuntu24 --target devshell -t hipdnn:devshell .
#
# Example: Run with mounted ROCm
#   docker run -v /path/to/rocm:/opt/rocm hipdnn:devshell bash
# ============================================================================
FROM base AS devshell

# Install explicit version of tools that are supported for hipDNN development
# (clang-format, clang-tidy, llvm)
RUN apt-get update && apt-get install -y \
    clang-format-18 \
    clang-tidy-20 \
    clang-tools-20 \
    doxygen \
    llvm-20 \
    && rm -rf /var/lib/apt/lists/*

# Install flatc version 25.9.23
RUN wget -q https://github.com/google/flatbuffers/releases/download/v25.9.23/Linux.flatc.binary.g++-13.zip -O /tmp/flatc.zip \
    && unzip /tmp/flatc.zip -d /usr/local/bin \
    && chmod +x /usr/local/bin/flatc \
    && rm /tmp/flatc.zip

# Use update-alternatives to set default versions
RUN update-alternatives --install /usr/bin/clang-format clang-format /usr/bin/clang-format-18 100

RUN update-alternatives --install /usr/bin/clang-tidy clang-tidy /usr/bin/clang-tidy-20 100

# Set up update-alternatives for all LLVM tools
RUN for tool in /usr/bin/llvm*-20; do \
        if [ -f "$tool" ]; then \
            base=$(basename "$tool" -20); \
            update-alternatives --install "/usr/bin/$base" "$base" "$tool" 100; \
            echo "Set alternative for $base -> $tool"; \
        fi \
    done

RUN python3 -m pip install --break-system-packages pre-commit

# Create mount point for /opt/rocm (will be empty unless mounted or built upon)
RUN mkdir -p /opt/rocm

ENV HIP_PLATFORM=amd
ENV PATH="/opt/rocm/bin:/opt/rocm:${PATH}"
ENV CMAKE_GENERATOR="Ninja"

WORKDIR /workspace

# ============================================================================
# hipDNN full image - development tools WITH /opt/rocm baked in
#
# This is the default target, providing a complete self-contained image with
# TheRock (ROCm) pre-installed. Use BUILD_TYPE to choose prebuilt vs fullbuild.
#
# Example: Build with prebuilt TheRock (default)
#   docker build -f Dockerfile.ubuntu24 -t hipdnn:prebuilt .
#
# Example: Build with source-built TheRock
#   docker build -f Dockerfile.ubuntu24 --build-arg BUILD_TYPE=fullbuild -t hipdnn:fullbuild .
# ============================================================================
FROM devshell AS hipdnn

# Import ARGs into this layer so they can be captured in env vars below.
ARG BUILD_TYPE
ARG ASAN_MODE

# Record some container build settings in environment variables in the container
ENV HIPDNN_DEV_CONTAINER_BUILD_TYPE=${BUILD_TYPE}
ENV HIPDNN_DEV_CONTAINER_ASAN=${ASAN_MODE}

COPY --from=install_therock /opt/rocm /opt/rocm