#!/usr/bin/env bash
# Fetch every vLLM build sub-dependency for the r27 image at the pins the
# r27 tree itself declares (cmake/external_projects/*.cmake + CMakeLists.txt).
#
# No fallbacks: each clone fetches the pinned ref directly and fails the
# build loudly on any mismatch. Before cloning, every pin printed below is
# grep-verified against the cloned vLLM tree — if a pin has drifted from
# what the tree expects, the mismatch aborts with the offending pin(s).
#
# Runs INSIDE the build container: /workspace/vllm must already contain the
# gated r27 tree (see Dockerfile `vllm-clone` stage). Populates /deps2.
set -euo pipefail
WORK=${WORK:-/workspace/vllm}
DEPS=${DEPS:-/deps2}
git config --global protocol.version 0

fetch() { # repo ref dst
    local repo="$1" ref="$2" dst="$3"
    echo "[deps] $dst @ $ref  ($repo)"
    git init -q "$dst"
    git -C "$dst" fetch -q --depth 1 "$repo" "$ref"
    git -C "$dst" checkout -q FETCH_HEAD
    git -C "$dst" config --global --add safe.directory "$dst"
}

gate() { # file:pattern ...  — every pin must appear verbatim in the tree
    local missing=0
    for g in "$@"; do
        local file="${g%%:*}" pattern="${g#*:}"
        local fp="$WORK/cmake/external_projects/$file"
        [ -f "$fp" ] || fp="$WORK/$file"   # CMakeLists.txt lives at the root
        if ! grep -Fq "$pattern" "$fp"; then
            echo "GATE FAILED: pin '$pattern' not found in $file (tree/pin drift)" >&2
            missing=1
        fi
    done
    [ "$missing" = 0 ] || { echo "aborting: staged pins drifted from the tree" >&2; exit 1; }
    echo "[deps] all pins verified against the r27 tree"
}

# ---- 1. Verify our pins against the tree's own declarations ----
gate CMakeLists.txt:nvidia/cutlass.git \
     CMakeLists.txt:v4.4.2 \
     deepgemm.cmake:deepseek-ai/DeepGEMM \
     deepgemm.cmake:8b1392b978f5a03c828dd1711090d7fb50958b8a \
     flashkda.cmake:3b225bf26bb8e218928a1fe14751cb48cf31d11b \
     flashmla.cmake:a8f794d1251cbfd88a5011445dd5582289c727e4 \
     fmha_sm100.cmake:087c161814d4d9c735b46c21212a09e5f8eb92fa \
     qutlass.cmake:e74319e3405ce6d71965732880f5dc1f52371f64 \
     tml_fa4.cmake:b206834606ed5b5f21f8eed6b0683f528ea9cf7d \
     vllm_flash_attn.cmake:f3e1a4f74c99145c0717709860bf765de1703779 \
     vllm_flash_attn.cmake:vllm-project/flash-attention \
     triton_kernels.cmake:triton-lang/triton \
     triton_kernels.cmake:v3.5.1

mkdir -p "$DEPS"

# ---- 2. Clone deps at the tree's pins ----
fetch https://github.com/nvidia/cutlass.git v4.4.2 "$DEPS/cutlass"
fetch https://github.com/vllm-project/flash-attention.git \
      f3e1a4f74c99145c0717709860bf765de1703779 "$DEPS/flash-attention"
fetch https://github.com/vllm-project/MSA.git \
      087c161814d4d9c735b46c21212a09e5f8eb92fa "$DEPS/MSA"
fetch https://github.com/vllm-project/FlashMLA.git \
      a8f794d1251cbfd88a5011445dd5582289c727e4 "$DEPS/FlashMLA"
fetch https://github.com/vllm-project/FlashKDA.git \
      3b225bf26bb8e218928a1fe14751cb48cf31d11b "$DEPS/FlashKDA"
fetch https://github.com/IST-DASLab/qutlass.git \
      e74319e3405ce6d71965732880f5dc1f52371f64 "$DEPS/qutlass"
fetch https://github.com/vllm-project/tml-fa4.git \
      b206834606ed5b5f21f8eed6b0683f528ea9cf7d "$DEPS/tml-fa4"
fetch https://github.com/triton-lang/triton.git v3.5.1 "$DEPS/triton"

# ---- 3. Sub-module populations the tree's install steps require ----
# Several deps vendor cutlass via a 160000 gitlink and their cmake install
# steps require the tree to EXIST (r18 lesson: an empty stub dies at cmake
# INSTALL — and r27's FlashMLA addendum died at compile). Generic rule: walk
# every dep's tree recursively and populate each gitlink at its pinned SHA.
# AMD-only trees (composable_kernel, aiter) are skipped explicitly — the
# CUDA cmake path never references them. All gitlinks so far vendor
# nvidia/cutlass; anything else aborts loudly, keeping the contract explicit.
populate_gitlinks() {
    local d path sha
    for d in flash-attention MSA FlashMLA FlashKDA qutlass tml-fa4; do
        while IFS=$'\t' read -r path sha; do
            [ -n "${path:-}" ] || continue
            case "$path" in
                */composable_kernel|composable_kernel|*/aiter|aiter)
                    echo "[deps] skip AMD-only gitlink $path (CUDA path unused)"; continue ;;
            esac
            case "$(basename "$path")" in
                cutlass) ;;    # expected vendor
                *) echo "ABORT: unhandled gitlink '$path' in $d (populate explicitly)" >&2; exit 1 ;;
            esac
            # a gitlink path EXISTS in the worktree as an EMPTY dir (git's
            # clone behavior) — the gate is "did WE populate it", .git marker.
            # Skipping on plain -d reproduces the r18 empty-stub failure.
            if [ ! -e "$DEPS/$d/$path/.git" ]; then
                fetch https://github.com/nvidia/cutlass.git "$sha" "$DEPS/$d/$path"
            fi
        done < <(git -C "$DEPS/$d" ls-tree -r -t HEAD | awk '$1=="160000" {print $4"\t"$3}')
    done
    echo "[deps] gitlink populations verified"
}
populate_gitlinks

# ---- 4. Triton into the layout the Dockerfile/build expects ----
mv "$DEPS/triton/python/triton_kernels/triton_kernels" "$DEPS/triton_kernels"
rm -rf "$DEPS/triton"

# ---- 5. Manifest ----
echo "[deps] staged:"
for d in cutlass flash-attention MSA FlashMLA FlashKDA qutlass tml-fa4 triton_kernels; do
    printf '  %-16s %s\n' "$d" "$(git -C "$DEPS/$d" rev-parse HEAD 2>/dev/null | cut -c1-12)"
done