# Use the LTS Ubuntu repo
FROM ubuntu:26.04 AS reqs
# Set global build variables
ENV PIP_BREAK_SYSTEM_PACKAGES=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive

# Install libraries
RUN apt update -y && \
    apt upgrade -y && \
    apt install -y --no-install-recommends \
    ca-certificates \
    wget \
    curl \
    git \
    gfortran \
    ninja-build \
    clang \
    pkg-config \
    xxd \
    make \
    automake \
    texinfo \
    bison \
    cargo \
    rustc \
    flex \
    mirage \
    cmake \
    patchelf \
    build-essential \
    zlib1g-dev \
    protobuf-compiler \
    libtool \
    libegl1-mesa-dev \
    libncurses5-dev \
    libopenblas-dev \
    libjpeg-dev \
    libpng-dev \
    libavcodec-dev \
    libavformat-dev \
    libswscale-dev \
    libpthreadpool-dev \
    libgoogle-glog-dev \
    libgtest-dev \
    libbenchmark-dev \
    libgdbm-dev \
    libnss3-dev \
    libssl-dev \
    libreadline-dev \
    libffi-dev \
    libsqlite3-dev \
    libelf-dev \
    libpciaccess-dev \
    libprotobuf-dev \
    libbz2-dev \
    liblzma-dev \
    libzstd-dev \
    && rm -rf /var/lib/apt/lists/*

# Build python 3.11
FROM reqs AS python
WORKDIR /build/python
RUN wget https://www.python.org/ftp/python/3.11.15/Python-3.11.15.tgz && \
    tar -xf Python-3.11.15.tgz && \
    cd Python-3.11.15 && \
    ./configure --enable-optimizations --enable-shared LDFLAGS="-Wl,-rpath /usr/local/lib" && \
    make -j 80 && \
    make install && \
    ldconfig /usr/local/lib && \
    wget https://bootstrap.pypa.io/get-pip.py && \
    python3 get-pip.py && \
    ln -sf /usr/local/bin/pip3 /bin/pip && \
    ln -sf /usr/local/bin/pip3 /bin/pip3 && \
    ln -sf /usr/local/bin/python3 /bin/python && \
    ln -sf /usr/local/bin/python3 /bin/python3.11

# Build ROCm latest
FROM python AS rocm
WORKDIR /build/therock
ENV CFLAGS="-O2 -Wno-error -DSIMDE_ENABLE_NATIVE_ALIASES"
ENV CXXFLAGS="-O2 -Wno-error -fpermissive -DSIMDE_ENABLE_NATIVE_ALIASES"
ENV MAX_JOBS=16
ENV MAX_LINK=1

# ROC Profiler is not supported on GFX1100
# HIPSparSELT is not supported on GFX1100
# Disable GIT Version Checking Used as a speedup
# Use precompiled LIBDRM
RUN env INSTALL_PREFIX=/usr/local ./dockerfiles/install_pinned_patchelf.sh && \
    pip install --root-user-action=ignore -r requirements.txt && \
    cmake -S . -B build \
    -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
    -DTHEROCK_AMDGPU_TARGETS="gfx1100" \
    -DAMDGPU_TARGETS="gfx1100" \
    -DTHEROCK_AMDGPU_FAMILIES="gfx1100" \
    -DTHEROCK_DIST_AMDGPU_FAMILIES="gfx1100" \
    -DTHEROCK_AMDGPU_DIST_BUNDLE_NAME="gfx1100" \
    -DLLVM_HOST_TRIPLE="aarch64-unknown-linux-gnu" \
    -DLLVM_DEFAULT_TARGET_TRIPLE="aarch64-unknown-linux-gnu" \
    -DCMAKE_C_FLAGS="$CFLAGS" \
    -DCMAKE_CXX_FLAGS="$CXXFLAGS" \
    -DLLVM_ENABLE_RUNTIMES="compiler-rt;libcxx;libcxxabi" \
    -DLLVM_ENABLE_PROJECTS="clang;lld;flang" \
    -DFLANG_PARALLEL_COMPILE_JOBS=${MAX_JOBS} \
    -DFLANG_PARALLEL_LINK_JOBS=${MAX_JOBS} \
    -DLLVM_PARALLEL_LINK_JOBS=${MAX_LINK} \
    -DLLVM_PARALLEL_COMPILE_JOBS=${MAX_JOBS} \
    -DCLANG_PARALLEL_COMPILE_JOBS=${MAX_JOBS} \
    -DCLANG_PARALLEL_LINK_JOBS=${MAX_LINK} \
    -DCOMPILER_RT_BUILD_SANITIZERS=OFF \
    -DTHEROCK_ENABLE_ROCPROFILER_COMPUTE=OFF \
    -DTHEROCK_ENABLE_HIPSPARSELT=OFF \
    -DTHEROCK_FLAG_STAMP_LIBRARY_GIT_VERSIONS=OFF \
    -DTHEROCK_ENABLE_SYSDEPS=ON \
    -DTHEROCK_BUILD_SYSDEPS=ON \
    -DCMAKE_INSTALL_PREFIX=./install && \
    sed -i 's|-Wl,--version-script=\${CMAKE_CURRENT_SOURCE_DIR}/version.lds||' /build/therock/third-party/sysdeps/linux/amd-mesa/CMakeLists.txt && \
    sed -i 's|-Wl,--version-script=\${CMAKE_CURRENT_SOURCE_DIR}/version.lds||' /build/therock/third-party/sysdeps/linux/libdrm/CMakeLists.txt && \
    sed -i 's|-Wl,--version-script=\${CMAKE_CURRENT_SOURCE_DIR}/version.lds||' /build/therock/third-party/sysdeps/linux/libpciaccess/CMakeLists.txt && \
    cmake --build build --parallel 1
    # Work in progress [ Breaking due to x86 intrinsics not ported to Aarch64 ]

# Build pytorch
FROM rocm AS pytorch
# Setup Envronment
ENV PYTORCH_ROCM_ARCH=1100
ENV USE_ROCM=1
ENV USE_CUDA=0
ENV USE_AOTRITON=0
ENV USE_SYSTEM_PTHREADPOOL=1
ENV USE_SYSTEM_PROTOC=1
ENV BUILD_CUSTOM_PROTOBUF=OFF
ENV USE_PYTORCH_QNNPACK=OFF
ENV USE_XNNPACK=OFF
ENV USE_ONNX=OFF
ENV USE_QNNPACK=OFF
ENV USE_NNPACK=OFF
ENV USE_KINETO=OFF
ENV USE_OPENMP=OFF
ENV USE_TENSORPIPE=OFF
ENV ROCM_PATH=/usr
ENV DEVICE_LIB_PATH=/usr/lib/llvm-21/lib/clang/21/amdgcn/bitcode
ENV HIPFLAGS="--rocm-device-lib-path=$DEVICE_LIB_PATH"
ENV CMAKE_ARGS="-DCMAKE_HIP_COMPILER_WORKS=1 \
-DCMAKE_CXX_COMPILER_WORKS=1 \
-DCMAKE_C_COMPILER_WORKS=1 \
-DCMAKE_HIP_COMPILER_ROCM_LIB_PATH=$DEVICE_LIB_PATH \
-DCMAKE_HIP_FLAGS=--rocm-device-lib-path=$DEVICE_LIB_PATH \
-DCMAKE_C_COMPILER=/usr/lib/llvm-21/bin/clang \
-DCMAKE_CXX_COMPILER=/usr/lib/llvm-21/bin/clang++ \
-DCMAKE_HIP_COMPILER=/usr/lib/llvm-21/bin/clang++ \
-DPython3_ROOT_DIR=/usr/local \
-DPYTHON_EXECUTABLE=/usr/local/bin/python3.11 \
-DPYTHON_LIBRARY=/usr/local/lib/libpython3.11.so \
-DPYTHON_INCLUDE_DIR=/usr/local/include/python3.11"
# Pull and build pytorch
WORKDIR /build/pytorch
RUN git clone https://github.com/pytorch/pytorch.git && \
    cd pytorch && \
    git submodule update --init --recursive && \
    make triton && \
    python tools/amd_build/build_amd.py && \
    pip install -r requirements.txt && \
    mkdir -p /opt/rocm/bin && \
    ln -s /usr/bin/hipcc /opt/rocm/bin/hipcc && \
    ln -s /usr/lib/llvm-21/bin/clang++ /usr/bin/clang++ && \
    python3 setup.py build && \
    python3 setup.py install
    # Failing

# Build pytorch-audio
FROM pytorch AS torchaudio
WORKDIR /build/torchaudio
RUN wget https://github.com/pytorch/audio/archive/refs/tags/v2.11.0.tar.gz && \
    tar -xf v2.11.0.tar.gz && \
    cd audio-2.11.0 && \
    pip install --no-build-isolation --no-deps -v .
    # Tested Working via comfyui-cuda-dockerfile

# Build pytorch-vision
FROM torchaudio AS torchvision
WORKDIR /build/torchaudio
ENV FORCE_CUDA=0
RUN git clone https://github.com/pytorch/vision.git . && \
    pip install numpy pillow ninja setuptools wheel && \
    python setup.py bdist_wheel && \
    pip install dist/torchvision-*.whl
    # Not tested and not verified ROCM build

# Build final stripped down image
#FROM ubuntu:26.04
# Copy in rocm, pytorch, torchaudio, torchvision