#!/usr/bin/env bash

# Generates the CUDA driver link stub: $1 = a CUDA C header, $2 = the .so to write.
#
# vllm links its C++ kernels against CUDA::cuda_driver, the CUDA driver library
# (libcuda.so.1). The driver package provides it; the CUDA toolkit wheel does
# not, and neither does a build container without a GPU. NVIDIA ships the
# link-time stub for it in the runfile installer, not in the pip wheel, so this
# script builds an equivalent one from the header: one exported symbol per
# driver entry point, no bodies, 87 KB.
#
# Link time only. The soname is libcuda.so.1, so the dynamic linker opens the
# real driver from the host at run time and can never pick the stub.
#
# Writes nothing when the file already has every entry point, so a re-run does
# not invalidate the compiled-ops hash recorded by build.sh.

set -euo pipefail

H="${1:?usage: $0 <cuda.h> <out.so>}"
OUT="${2:?usage: $0 <cuda.h> <out.so>}"

[ -r "$H" ] || { echo "no readable CUDA header at $H"; exit 1; }

decls(){
  # Calls written as "CUresult CUDAAPI cuFoo(...)", split over lines.
  awk '/CUresult[ \t]+CUDAAPI/{cap=1;buf=""} cap{buf=buf" "$0; if(buf~/;/){
       if(match(buf,/CUresult[ \t]+CUDAAPI[ \t]+[A-Za-z_][A-Za-z0-9_]*/)){
         s=substr(buf,RSTART,RLENGTH); sub(/.*CUDAAPI[ \t]+/,"",s); print s} cap=0}}' "$H"
  # Versioned names hidden behind a macro: code writes cuFoo, the preprocessor
  # rewrites it to cuFoo_v2, so the object file asks the linker for cuFoo_v2.
  sed -nE 's/^#define[ \t]+(cu[A-Za-z0-9_]+)[ \t]+(cu[A-Za-z0-9_]+_v[0-9]+)([ \t]|$).*/\2/p' "$H"
}

want="$(decls | sort -u)"
n="$(printf '%s\n' "$want" | sed '/^$/d' | wc -l)"

[ "$n" -gt 100 ] || { echo "only $n driver entry points found in $H"; exit 1; }

if [ -e "$OUT" ]; then
  have="$(nm -D --defined-only "$OUT" 2>/dev/null | awk '$2=="T"{print $3}' | sort -u)"
  miss="$(LC_ALL=C comm -23 <(printf '%s\n' "$want") <(printf '%s\n' "$have") | sed '/^$/d')"

  if [ -z "$miss" ]; then
    echo "driver stub is current: $(basename "$OUT"), $n entry points"
    exit 0
  fi

  echo "driver stub is missing $(printf '%s\n' "$miss" | wc -l) entry points, regenerating"
fi

mkdir -p "$(dirname "$OUT")"
src="$(mktemp --suffix=.c)"

cleanup(){ rm -f "$src"; }

trap cleanup EXIT
: > "$src"
printf '%s\n' "$want" | while read -r f; do
  [ -n "$f" ] && printf 'unsigned int %s(void){ return 0; }\n' "$f" >> "$src"
done

cc -shared -fPIC -O0 -o "$OUT" "$src" -Wl,-soname,libcuda.so.1

have="$(nm -D --defined-only "$OUT" | awk '$2=="T"{print $3}' | sort -u)"
miss="$(LC_ALL=C comm -23 <(printf '%s\n' "$want") <(printf '%s\n' "$have") | sed '/^$/d')"

[ -z "$miss" ] || { echo "the stub is still missing:"; printf '%s\n' "$miss"; exit 1; }

echo "driver stub generated: $OUT ($(printf '%s\n' "$want" | wc -l) entry points, soname libcuda.so.1)"