FROM ubuntu:24.04 AS mesa-builder
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    git \
    python3-mako \
    python3-yaml \
    python3-packaging \
    meson \
    ninja-build \
    build-essential \
    pkg-config \
    flex \
    bison \
    libpciaccess-dev \
    libexpat1-dev \
    libelf-dev \
    libvulkan-dev \
    glslang-tools \
    spirv-tools \
    llvm-dev \
    libclang-dev \
    cmake \
    libwayland-dev \
    wayland-protocols \
    libwayland-bin \
    libxcb1-dev \
    libx11-xcb-dev \
    libxcb-dri3-dev \
    libxcb-present-dev \
    libxcb-sync-dev \
    libxshmfence-dev \
    libx11-dev \
    libxcb-shm0-dev \
    libxrandr-dev \
    git \
    python3-mako \
    python3-yaml \
    python3-packaging \
    python3-pip \
    ninja-build \
    build-essential \
    pkg-config \
    flex \
    bison \
    libpciaccess-dev \
    libexpat1-dev \
    libelf-dev \
    libvulkan-dev \
    glslang-tools \
    spirv-tools \
    llvm-dev \
    libclang-dev \
    cmake \
    && rm -rf /var/lib/apt/lists/*

RUN pip3 install --no-cache-dir --break-system-packages --upgrade meson

# Build required newer libdrm from source
RUN git clone --depth 1 https://gitlab.freedesktop.org/mesa/drm.git /drm
WORKDIR /drm
RUN meson setup build -Dbuildtype=release -Dprefix=/opt/mesa
RUN ninja -C build -j16 && ninja -C build install

# Build matching LLVMSPIRVLib (SPIRV-LLVM-Translator) for LLVM 18
RUN git clone --branch llvm_release_180 --depth 1 https://github.com/KhronosGroup/SPIRV-LLVM-Translator.git /spirv-llvm-translator
WORKDIR /spirv-llvm-translator
RUN cmake -B build \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_INSTALL_PREFIX=/opt/mesa \
    -DLLVM_DIR=/usr/lib/llvm-18/cmake
RUN cmake --build build --parallel 16 --verbose
RUN cmake --install build

# Build Mesa-git pointing to the custom libdrm and LLVMSPIRVLib paths
RUN git clone --depth 1 https://gitlab.freedesktop.org/mesa/mesa.git /mesa
WORKDIR /mesa

ENV PKG_CONFIG_PATH=/opt/mesa/lib/pkgconfig:/opt/mesa/lib/aarch64-linux-gnu/pkgconfig:$PKG_CONFIG_PATH
ENV CMAKE_PREFIX_PATH=/opt/mesa:$CMAKE_PREFIX_PATH

RUN meson setup build \
    -Dbuildtype=release \
    -Dgallium-drivers= \
    -Dvulkan-drivers=amd \
    -Dplatforms=x11,wayland \
    -Dprefix=/opt/mesa

RUN ninja -C build -j16
RUN ninja -C build install

# Stage 2: Build Ollama with Vulkan support
FROM ubuntu:24.04 AS ollama-builder
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    build-essential \
    git \
    curl \
    wget \
    cmake \
    golang-go \
    ca-certificates \
    ninja-build \
    libvulkan-dev \
    vulkan-tools \
    glslang-tools \
    spirv-headers \
    spirv-tools \
    libshaderc-dev \
    glslc \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app
RUN git clone --recursive https://github.com/ollama/ollama.git .

# Configure and compile Ollama using the Vulkan backend
RUN cmake -B build . \
    -DOLLAMA_LLAMA_BACKENDS="vulkan" \
    -DGGML_VULKAN=ON \
    -DCMAKE_BUILD_TYPE=Release

RUN cmake --build build --parallel 16
RUN go build -o ollama .

# Stage 3: Minimal Runtime Image
FROM ubuntu:24.04

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    ca-certificates \
    libvulkan1 \
    vulkan-tools \
    libexpat1 \
    libelf1 \
    libdrm2 \
    && rm -rf /var/lib/apt/lists/*

# Copy compiled Mesa-git libraries and ICD files from Stage 1
COPY --from=mesa-builder /opt/mesa /opt/mesa

# Copy compiled Ollama binary from Stage 2
COPY --from=ollama-builder /app/ollama /app/ollama
COPY --from=ollama-builder /app/build /app/build

EXPOSE 11434

ENV OLLAMA_HOST=0.0.0.0:11434
ENV GGML_VULKAN_DISABLE_FMA=1
ENV RADV_PERFOPTS="no_outoforder"
ENV AMD_DEBUG="notcc"
ENV OLLAMA_MODELS=/root/.ollama/models

# Point the loader directly to the custom Mesa-git ICD driver
ENV VK_ICD_FILENAMES=/opt/mesa/share/vulkan/icd.d/radeon_icd.aarch64.json
ENV LD_LIBRARY_PATH=/opt/mesa/lib/aarch64-linux-gnu:/opt/mesa/lib:$LD_LIBRARY_PATH

ENTRYPOINT ["/app/ollama", "serve"]