# Refer to https://github.com/RadeonOpenCompute/ROCm-docker/blob/master/dev/Dockerfile-ubuntu-24.04-complete
FROM ubuntu:24.04

ARG ROCM_VERSION=6.3.2
ARG PYTHON_VERSION=3.12
ARG CMAKE_VERSION=3.30.1
ARG CCACHE_VERSION=4.7.4
ARG USE_MIGRAPHX=0

LABEL ROCM_VERSION="${ROCM_VERSION}"
LABEL USE_MIGRAPHX="${USE_MIGRAPHX}"
LABEL PYTHON_VERSION="${PYTHON_VERSION}"

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive \
    LC_ALL=C.UTF-8 \
    LANG=C.UTF-8 \
    VIRTUAL_ENV=/opt/venv \
    LD_LIBRARY_PATH=/opt/rocm/lib:/usr/lib/x86_64-linux-gnu::/opt/amdgpu/lib/x86_64-linux-gnu \
    PATH=/opt/venv/bin:/opt/rocm/bin:/usr/bin:/bin:/usr/sbin:/usr/local/bin

# Set default shell
SHELL ["/bin/bash", "-c"]

# Add ROCm package pinning
RUN echo "Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600" > /etc/apt/preferences.d/rocm-pin-600

# Install dependencies and ROCm
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        ca-certificates \
        curl \
        libnuma-dev \
        gnupg && \
    curl -sL https://repo.radeon.com/rocm/rocm.gpg.key | apt-key add - && \
    printf "deb [arch=amd64] https://repo.radeon.com/rocm/apt/${ROCM_VERSION}/ noble main" | tee --append /etc/apt/sources.list.d/rocm.list && \
    printf "deb [arch=amd64] https://repo.radeon.com/amdgpu/${ROCM_VERSION}/ubuntu noble main" | tee /etc/apt/sources.list.d/amdgpu.list && \
    migraphx=$( [ "$USE_MIGRAPHX" -eq 1 ] && echo "migraphx" || echo "" ) && \
    apt-get update && apt-get install -y --no-install-recommends \
        sudo \
        libelf1 \
        kmod \
        file \
        git \
        python3-pip \
        python${PYTHON_VERSION}-dev \
        python${PYTHON_VERSION}-venv \
        hipcc=1.1.1.60302-66~24.04 \
        rocm-cmake=0.14.0.60302-66~24.04 \
        rocm-utils=6.3.2.60302-66~24.04 \
        rocminfo=1.0.0.60302-66~24.04 \
        rocm-dev \
        rocm-libs \
        $migraphx \
        build-essential \
        locales \
        cifs-utils \
        wget \
        half \
        libnuma-dev \
        lsb-release && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* /var/cache/apt/archives/* && \
    locale-gen en_US.UTF-8 && \
    update-locale LANG=en_US.UTF-8

# Create render group
RUN groupadd -g 109 render

# Install CMake
WORKDIR /usr/local
RUN wget -q https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-Linux-x86_64.tar.gz && \
    tar -zxf cmake-${CMAKE_VERSION}-Linux-x86_64.tar.gz --strip=1 -C /usr && \
    rm -f cmake-${CMAKE_VERSION}-Linux-x86_64.tar.gz

# Install ccache
RUN mkdir -p /tmp/ccache && \
    cd /tmp/ccache && \
    wget -q -O - https://github.com/ccache/ccache/releases/download/v${CCACHE_VERSION}/ccache-${CCACHE_VERSION}-linux-x86_64.tar.xz | tar --strip 1 -J -xf - && \
    cp /tmp/ccache/ccache /usr/bin/ && \
    rm -rf /tmp/ccache

# Set up virtual environment for Python and install dependencies
COPY scripts/requirements.txt /ort/
RUN python3 -m venv ${VIRTUAL_ENV} && \
    source ${VIRTUAL_ENV}/bin/activate && \
    pip install --upgrade pip setuptools wheel && \
    pip install -r /ort/requirements.txt && \
    pip install ml_dtypes pytest-xdist pytest-rerunfailures scipy


# Default command
CMD ["/bin/bash"]