# ---------- llama.cpp Build Stage
FROM nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04 AS builder

ENV DEBIAN_FRONTEND=noninteractive

# Install build tools, CMake, Git, and documentation tools
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    cmake \
    git \
    ninja-build \
    ca-certificates \
    libgomp1 \
    pandoc \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /workspace

# Clone llama.cpp repository
ARG LLAMA_CPP_TAG=master
RUN git clone --depth 1 --branch ${LLAMA_CPP_TAG} https://github.com/ggerganov/llama.cpp.git .

# Build llama.cpp, the ldconfig is important
RUN echo "/usr/local/cuda-13.2/compat" > /etc/ld.so.conf.d/988-cuda-compat.conf && ldconfig
RUN cmake -B build -S . \
     -DCMAKE_BUILD_TYPE=Release \
     -DGGML_CUDA=ON \
     -DCMAKE_INSTALL_PREFIX=/install \
     -DCMAKE_CUDA_ARCHITECTURES=120 \
     -DLLAMA_BUILD_TESTS=OFF \
     -DLLAMA_BUILD_EXAMPLES=ON \
    && cmake --build build --config Release -j $(nproc) \
    && cmake --install build \
    && rm -rf build
# Generate or stage manpages into /install/share/man
# If manpages are prebuilt or markdown docs exist, compile/place them in standard sections:
RUN mkdir -p /install/share/man/man1 && \
    for doc in docs/*.md; do \
        [ -f "$doc" ] || continue; \
        base=$(basename "$doc" .md); \
        pandoc "$doc" -s -t man -o "/install/share/man/man1/${base}.1" 2>/dev/null || true; \
    done

# ---------- Python virtual environment builder

FROM nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04 AS venv-builder

RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt update && apt-get install -y --no-install-recommends \
    python3-pip \
    python3-dev \
    python3-venv \
    git

WORKDIR /workspace
RUN python3 -m venv /workspace/py-venv
ENV PATH="/workspace/py-venv/bin:${PATH}"
RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked pip3 install torch torchvision torchaudio --upgrade --index-url \
    https://download.pytorch.org/whl/cu132/

# Fetch convert_to_quant, and install it into our python venv.
RUN git clone --depth=1 https://github.com/silveroxides/convert_to_quant && \
    cd convert_to_quant && pip3 install -r requirements.txt && pip3 install .

# Install heretic-llm
RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked pip3 install heretic-llm

# Install zefie's nvfp4 quant requirements
COPY --chmod=0644 requirements.txt /workspace/requirements.txt
RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked pip3 install -r /workspace/requirements.txt

# ---------- Main image

FROM nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04 AS runner
ENV DEBIAN_FRONTEND=noninteractive

# "Unminify" nvidia container to enable man pages
RUN rm -f /etc/dpkg/dpkg.cfg.d/excludes && \
    rm -f /etc/apt/apt.conf.d/docker-clean && \
    rm -f /etc/apt/apt.conf.d/01-nodoc && \
    rm -f /usr/bin/man && \
    dpkg-divert --quiet --remove /usr/bin/man

RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt update && apt-get install -y \
    python3 \
    man \
    man-db \
    wget \
    bzip2 \
    curl

WORKDIR /opt/nvfp4

RUN mkdir nvfp4_gguf_quant
COPY --exclude=__pycache__ nvfp4gguf ./nvfp4_gguf_quant/nvfp4gguf
COPY --chmod=0644 gguf_to_nvfp4_gguf.py hf_to_nvfp4_gguf.py hf_to_nvfp4_safetensors.py safetensors_to_nvfp4_gguf.py hf_info.py gguf_info.py ./nvfp4_gguf_quant/

COPY --chmod=0644 docker/README.txt /workspace/README.txt
COPY --chmod=0755 --exclude=README.txt docker/* /usr/local/bin/

# Install manpages for zefie's nvfp4 quant scripts, ctq and heretic
RUN mkdir -p /usr/local/share/man/man1
COPY --chmod=0644 man/man1/* /usr/local/share/man/man1/

# Copy llama.cpp compiled binaries, manpages and shared libraries from builder
COPY --from=builder /install/bin/ /usr/local/bin/
COPY --from=builder /install/lib/ /usr/local/lib/
COPY --from=builder /install/share/man/ /usr/share/man/

# Set up a minimal llama.cpp source for conversions
RUN mkdir /opt/nvfp4/llama.cpp
COPY --from=builder /workspace/conversion /opt/nvfp4/llama.cpp/conversion/
COPY --from=builder --chmod=0644 /workspace/*.py /opt/nvfp4/llama.cpp/

COPY --from=venv-builder /workspace/py-venv /opt/nvfp4/py-venv
ENV PATH="/opt/nvfp4/py-venv/bin:${PATH}"

# Finalize and settle at /workspace
ENV DEBIAN_FRONTEND=
WORKDIR /workspace