#!/usr/bin/env bash
# Build GLM-5.2 NF3-Hybrid Docker image for DGX Spark (aarch64/sm_121a)
#
# Build harness: m9e/blackwell-llm-docker  →  build-and-copy.sh
#   https://github.com/m9e/blackwell-llm-docker
#
# Pins:
#   vLLM:  local-inference-lab/vllm  @ dev/eldritch-enlightenment  45c1582e9b80ba83e71c3a6458e71da4736fbdc4
#   b12x:  voipmonitor/b12x           @ f3686b555d639823b276c2080f173145eed7f007
#   CUDA:  13.2 / torch 2.12.0+cu132
#   arch:  TORCH_CUDA_ARCH_LIST=12.1a, CUTE_DSL_ARCH=sm_121a
#   FlashInfer: 0.6.14 aarch64 wheels
#
# Usage:
#   cd recipes/glm52-nf3-hybrid
#   ./build.sh [--tag vllm-nf3-hybrid:probe]

set -euo pipefail

TAG="${1:-vllm-nf3-hybrid:probe}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# ── Prerequisites ──────────────────────────────────────────────────────────
command -v git  >/dev/null 2>&1 || { echo "Need git";  exit 1; }
command -v docker >/dev/null 2>&1 || { echo "Need docker"; exit 1; }

BUILD_ROOT="$(mktemp -d)"
trap 'rm -rf "$BUILD_ROOT"' EXIT

echo "==> Build root: $BUILD_ROOT"

# ── Clone build harness ────────────────────────────────────────────────────
HARNESS_REPO="https://github.com/m9e/blackwell-llm-docker"
echo "==> Cloning build harness from $HARNESS_REPO ..."
git clone --depth=1 "$HARNESS_REPO" "$BUILD_ROOT/harness"

# ── Build base image using build-and-copy.sh ───────────────────────────────
# build-and-copy.sh clones vLLM and b12x, applies patches, builds, and
# copies the artifacts (wheels + venv overlay) into a Docker image.
echo "==> Building base image with build-and-copy.sh ..."
cd "$BUILD_ROOT/harness"

# The harness vLLM ref includes NF3-hybrid loader, DCP ag_rs dispatch,
# low-bit MLA KV-cache dtypes (nvfp4_ds_mla, nf3_ds_mla), MTP/spec-decode
# cherries, and the patched b12x_sparse_indexer_warmup.py.
./build-and-copy.sh \
  --vllm-repo "https://github.com/local-inference-lab/vllm" \
  --vllm-ref "45c1582e9b80ba83e71c3a6458e71da4736fbdc4" \
  --b12x-repo "https://github.com/voipmonitor/b12x" \
  --b12x-ref "f3686b555d639823b276c2080f173145eed7f007" \
  --gpu-arch "12.1a" \
  -t "$TAG" \
  --clean

# The harness emits a Docker image tagged $TAG (no --push flag; local only).

# ── Bake NF3 v2 overlay + CuteDSL runtime ─────────────────────────────────
echo "==> Baking NF3 v2 overlay into $TAG ..."

# Create a Dockerfile that adds the layer on top of the base.
cat > "$BUILD_ROOT/Dockerfile.bake" << 'DOCKERFILE'
ARG BASE_TAG
FROM ${BASE_TAG}

# NF3 v2 overlay files (hybrid loader, NF3/NVFP4 kernels, replan, mxfp8 tier)
COPY overlay/hybrid_loader.py /usr/local/lib/python3.12/dist-packages/
COPY overlay/nf3_kernel.py    /usr/local/lib/python3.12/dist-packages/
COPY overlay/nvfp4_kernel.py  /usr/local/lib/python3.12/dist-packages/
COPY overlay/nf3_replan.py    /usr/local/lib/python3.12/dist-packages/
COPY overlay/mxfp8_tier.json  /usr/local/lib/python3.12/dist-packages/

# Install CuteDSL runtime (NF3 kernel dependency)
RUN pip install --no-cache-dir \
    nvidia-cutlass-dsl==4.5.2 \
    nvidia-cutlass-dsl-libs-cu13==4.5.2

# Install FlashInfer 0.6.14 aarch64 wheel
# Build from source for aarch64/sm_121a or use published wheel.
# Published aarch64 wheels are available at flashinfer.ai/whl.
RUN pip install --no-cache-dir \
    --find-links https://flashinfer.ai/whl/cu124/flashinfer_pypi/ \
    flashinfer==0.6.14

# Ensure CuteDSL .pth path hook is present
RUN python3 -c "import cutlass; print(f'CuteDSL OK: {cutlass.__version__}')"

# Set arch envs so JIT compile targets sm_121a
ENV TORCH_CUDA_ARCH_LIST=12.1a
ENV CUTE_DSL_ARCH=sm_121a
ENV FLASHINFER_CUDA_ARCH_LIST=12.1a

# Verify imports
RUN python3 -c "
from hybrid_loader import NF3HybridLoader
from nf3_kernel import NF3Kernel
from nvfp4_kernel import NVFP4Kernel
from nf3_replan import NF3Replan
print('NF3 hybrid overlay OK')
"
DOCKERFILE

# Build the overlay layer
docker build \
  -f "$BUILD_ROOT/Dockerfile.bake" \
  --build-arg BASE_TAG="$TAG" \
  -t "$TAG" \
  "$SCRIPT_DIR"

echo "==> Done: $TAG"
echo "To push: docker push $TAG (or use as local image)"