#!/usr/bin/env bash
# Fetch every vLLM build sub-dependency for the image at the pins the
# cloned 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).
#
# Lock: after staging, the script verifies the staged state against
# $DEPS_LOCK (default /tmp/deps.lock, shipped in the image dir as
# locks/deps.lock and COPY'd by the Dockerfile). Every build therefore
# proves: each dep is a clean checkout at the locked commit, every vendored
# cutlass gitlink is populated at the locked SHA (r18's empty-stub failure
# mode), and the non-git triton_kernels tree has the locked content hash.
# Regenerate the lock after a pin cut: `--regen-lock <out>` (or in-build:
# --build-arg REGEN_LOCK=1). `--emit-lock <out>` writes the lock from an
# already-staged $DEPS without re-cloning (offline lock generation).
#
# 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}
LOCK=${DEPS_LOCK:-/tmp/deps.lock}
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 ----
stage_pins() {
    # Pins verified against the r39 tree (9e5d1793, 2026-09-21): the lab moved
    # every external pin in one cut, DeepGEMM to its vllm-project fork, cutlass
    # v4.4.2 -> v4.7.1, and added DeepSelect. See profiles/r39.env.
    gate CMakeLists.txt:nvidia/cutlass.git \
         CMakeLists.txt:v4.7.1 \
         deepgemm.cmake:vllm-project/DeepGEMM \
         deepgemm.cmake:ad1f1726aa540a76c1d26d6a120effb8de21eaa4 \
         deepselect.cmake:vllm-project/DeepSelect \
         deepselect.cmake:d96d33afe1fab0d6066da49cdc91e64c2bee65ea \
         flashkda.cmake:b59532f1f464fbd536272780e30df5bf6a2ccc02 \
         flashmla.cmake:0eee43b12f034b657133cf2afca6a72ebb6efccf \
         fmha_sm100.cmake:f355c37eb4e1413f21ee2ad8bbad25079e6bef9d \
         qutlass.cmake:e74319e3405ce6d71965732880f5dc1f52371f64 \
         tml_fa4.cmake:75765e76a9c2c012c1f6ecd64577eb646eb4d303 \
         vllm_flash_attn.cmake:506341a143fcabd4bb79052a7605ada727d6b3f5 \
         vllm_flash_attn.cmake:vllm-project/flash-attention \
         triton_kernels.cmake:triton-lang/triton \
         triton_kernels.cmake:v3.5.1
}

# ---- 2. Clone deps at the tree's pins ----
stage_clones() {
    mkdir -p "$DEPS"
    fetch https://github.com/nvidia/cutlass.git v4.7.1 "$DEPS/cutlass"
    fetch https://github.com/vllm-project/flash-attention.git \
          506341a143fcabd4bb79052a7605ada727d6b3f5 "$DEPS/flash-attention"
    fetch https://github.com/vllm-project/MSA.git \
          f355c37eb4e1413f21ee2ad8bbad25079e6bef9d "$DEPS/MSA"
    fetch https://github.com/vllm-project/FlashMLA.git \
          0eee43b12f034b657133cf2afca6a72ebb6efccf "$DEPS/FlashMLA"
    fetch https://github.com/vllm-project/FlashKDA.git \
          b59532f1f464fbd536272780e30df5bf6a2ccc02 "$DEPS/FlashKDA"
    fetch https://github.com/IST-DASLab/qutlass.git \
          e74319e3405ce6d71965732880f5dc1f52371f64 "$DEPS/qutlass"
    fetch https://github.com/vllm-project/tml-fa4.git \
          75765e76a9c2c012c1f6ecd64577eb646eb4d303 "$DEPS/tml-fa4"
    fetch https://github.com/vllm-project/DeepSelect.git \
          d96d33afe1fab0d6066da49cdc91e64c2bee65ea "$DEPS/DeepSelect"
    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.
stage_gitlinks() {
    local d path sha
    for d in flash-attention MSA FlashMLA FlashKDA qutlass tml-fa4 DeepSelect; 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"
}

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

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

# ---------------------------------------------------------------- lock ----
# sha256 over the sorted file list (relpath + per-file sha256) — the
# file-level fingerprint for non-git content (git checkouts fingerprint via
# commit + clean worktree instead: git objects already hash every byte).
blob_hash() {
    (cd "$1" && find . -type f | sort | while IFS= read -r f; do sha256sum "$f"; done | sha256sum | cut -d' ' -f1)
}

emit_lock() { # out — write the lock from the staged $DEPS state
    local out="$1" d path sha h
    {
        echo "# sm121 deps.lock — staged-content fingerprints for fetch-deps.sh"
        echo "# generated: $(date -u +%Y-%m-%dT%H:%M:%SZ)"
        echo "# format: dep <name> <commit> | gitlink <parent>/<path> <commit>"
        echo "#         | skip <parent>/<path> <commit> | blob <name> <sha256>"
        echo "# regenerate: bash fetch-deps.sh --regen-lock <out> (in-build: --build-arg REGEN_LOCK=1)"
        for d in cutlass flash-attention MSA FlashMLA FlashKDA qutlass tml-fa4 DeepSelect; do
            [ -e "$DEPS/$d/.git" ] || { echo "emit_lock: $d not staged" >&2; exit 1; }
            echo "dep $d $(git -C "$DEPS/$d" rev-parse HEAD)"
        done
        for d in flash-attention MSA FlashMLA FlashKDA qutlass tml-fa4 DeepSelect; do
            while IFS=$'\t' read -r path sha; do
                [ -n "${path:-}" ] || continue
                case "$path" in
                    */composable_kernel|composable_kernel|*/aiter|aiter)
                        echo "skip $d/$path $sha" ;;
                    *)
                        echo "gitlink $d/$path $sha" ;;
                esac
            done < <(git -C "$DEPS/$d" ls-tree -r -t HEAD | awk '$1=="160000" {print $4"\t"$3}')
        done
        [ -d "$DEPS/triton_kernels" ] || { echo "emit_lock: triton_kernels missing (run stage_triton_layout first)" >&2; exit 1; }
        h=$(blob_hash "$DEPS/triton_kernels")
        echo "blob triton_kernels $h"
    } > "$out"
    echo "[deps] lock written: $out"
}

verify_lock() { # lock — abort the build if staged content != lock
    local lock="$1"
    if [ ! -f "$lock" ]; then
        echo "[deps] no lock at $lock — nothing to verify against (staging unverified)"
        return 0
    fi
    local kind name want got d fail=0
    while read -r kind name want; do
        case "$kind" in
            dep | gitlink)
                d="$DEPS/$name"
                if [ ! -e "$d/.git" ]; then
                    echo "LOCK FAIL: $kind $name — missing or not a git checkout (empty-stub r18 failure mode)" >&2
                    fail=1
                    continue
                fi
                got=$(git -C "$d" rev-parse HEAD 2>/dev/null) || got="<no HEAD>"
                if [ "$got" != "$want" ]; then
                    echo "LOCK FAIL: $kind $name @ ${got:0:12} != locked ${want:0:12}" >&2
                    fail=1
                    continue
                fi
                if [ -n "$(git -C "$d" status --porcelain 2>/dev/null)" ]; then
                    echo "LOCK FAIL: $kind $name — dirty worktree (expected a clean checkout)" >&2
                    fail=1
                fi
                ;;
            skip)
                d="$DEPS/${name%%/*}"
                got=$(git -C "$d" ls-tree -r -t HEAD 2>/dev/null | awk -v p="${name#*/}" '$1=="160000" && $4==p {print $3}')
                if [ "$got" != "$want" ]; then
                    echo "LOCK FAIL: skip $name — gitlink changed (${got:-gone} vs locked ${want:0:12})" >&2
                    fail=1
                fi
                ;;
            blob)
                got=$(blob_hash "$DEPS/$name" 2>/dev/null) || got="<unreadable>"
                if [ "$got" != "$want" ]; then
                    echo "LOCK FAIL: blob $name — content hash $got != locked $want" >&2
                    fail=1
                fi
                ;;
            *)
                echo "LOCK FAIL: unknown line kind '$kind' in $lock" >&2
                fail=1
                ;;
        esac
    done < <(grep -v '^#' "$lock" | grep -v '^[[:space:]]*$' || true)
    if [ "$fail" != 0 ]; then
        echo "aborting: staged deps do not match $lock" >&2
        exit 1
    fi
    echo "[deps] lock verified: staged content matches $lock"
}

case "${1:-}" in
    --regen-lock)
        stage_all
        emit_lock "${2:?usage: fetch-deps.sh --regen-lock <out>}"
        ;;
    --emit-lock)
        emit_lock "${2:?usage: fetch-deps.sh --emit-lock <out>}"
        ;;
    *)
        stage_all
        verify_lock "$LOCK"
        ;;
esac