FROM ubuntu:24.04 AS builder
ENV DEBIAN_FRONTEND=noninteractive

#install build deps
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 \
    python3-pip \
    curl \
    wget \
    golang-go \
    ca-certificates \
    vulkan-tools \
    spirv-headers \
    libshaderc-dev \
    glslc \
    && rm -rf /var/lib/apt/lists/*
RUN pip3 install --no-cache-dir --break-system-packages --upgrade meson

# Build libdrm
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 LLVMSPIRVLib
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
RUN cmake --install build

# Build Mesa-git
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
RUN ninja -C build install

# Build Ollama
RUN git clone --recursive https://github.com/ollama/ollama.git /app
WORKDIR /app
ENV CGO_ENABLED=1
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=builder /opt/mesa /opt/mesa

# Copy compiled Ollama binary from Stage 2
COPY --from=builder /app/ollama /app/ollama
COPY --from=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"]