#!/usr/bin/env bash

set -euo pipefail

failures=0

check_command() {
  local name="$1"
  if command -v "$name" >/dev/null 2>&1; then
    echo "OK   command: $name -> $(command -v "$name")"
  else
    echo "FAIL command: $name"
    failures=$((failures + 1))
  fi
}

check_file() {
  local path="$1"
  if [[ -f "$path" ]]; then
    echo "OK   file: $path"
  else
    echo "FAIL file: $path"
    failures=$((failures + 1))
  fi
}

echo "===== COMMANDS ====="
check_command python
check_command vllm
check_command gcc
check_command g++
check_command ninja
check_command nvidia-smi

echo
echo "===== SYSTEM FILES ====="
check_file /usr/include/python3.12/Python.h
check_file /usr/lib/wsl/lib/libcuda.so

echo
echo "===== PYTHON PACKAGES ====="
python - <<'PY'
import importlib.metadata

for package in ("vllm", "torch", "transformers", "compressed-tensors", "nvidia-cuda-nvrtc"):
    try:
        print(f"OK   {package}: {importlib.metadata.version(package)}")
    except importlib.metadata.PackageNotFoundError:
        print(f"MISS {package}")
PY

echo
echo "===== NVRTC FILES ====="
python - <<'PY'
import site
from pathlib import Path

found = []
for site_dir in site.getsitepackages():
    for path in (Path(site_dir) / "nvidia").glob("cu*/lib/libnvrtc*.so*"):
        found.append(path)

for path in sorted(found):
    print(path)

if not found:
    raise SystemExit("No NVRTC libraries found")
PY

echo
echo "===== CUDA DRIVER LINK ====="
cat >/tmp/internvl_cuda_link_test.cpp <<'CPP'
extern "C" int cuInit(unsigned int flags);
int main() { return 0; }
CPP

LIBRARY_PATH="/usr/lib/wsl/lib:${LIBRARY_PATH:-}" \
  g++ /tmp/internvl_cuda_link_test.cpp -Wl,--no-as-needed -lcuda \
  -o /tmp/internvl_cuda_link_test

echo "OK   CUDA driver link"

if [[ "$failures" -ne 0 ]]; then
  echo "Runtime check failed: $failures required item(s) missing." >&2
  exit 1
fi

echo
echo "WSL RUNTIME CHECK: SUCCESS"