#!/usr/bin/env bash
# Merges the NVIDIA wheels of the venv (uv) into a single CUDA tree for CMake.
# Produces .venv/cuda/{bin,include,lib} with symlinks; does NOT require sudo.
# Usage: bash scripts/prepare-cuda.sh

set -euo pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
VENV="$REPO/.venv"
NV="$VENV/lib/python3.12/site-packages/nvidia"
DEST="$VENV/cuda"

[ -d "$NV" ] || { echo "ERROR: NVIDIA wheels not found in $NV (did uv sync fail?)"; exit 1; }

mkdir -p "$DEST/bin" "$DEST/include" "$DEST/lib"

# nvcc
find "$NV" -path '*/bin/nvcc' -exec ln -sf {} "$DEST/bin/nvcc" \;
# headers (cuda.h, cublas_v2.h, ...)
find "$NV" -path '*/include/*' -name '*.h' -exec ln -sf {} "$DEST/include/" \;
# libraries (.so and .so.*)
find "$NV" -path '*/lib/*' -name 'lib*.so*' -exec ln -sf {} "$DEST/lib/" \;

echo "[prepare-cuda] toolkit merged at $DEST"
"$DEST/bin/nvcc" --version | tail -2 | head -1