#!/usr/bin/env bash
# Purpose: check an NVIDIA desktop or laptop against everything the rest of the course
#          assumes: the driver and card (and, under WSL2, that no Linux GPU driver was
#          installed inside WSL), the CUDA toolkit, Docker with the NVIDIA runtime and
#          Compose, the toolchain, the Part 1 environment and notebook, the hf CLI and the
#          model storage. It changes nothing.
# Platform: nvidia (Linux, or Windows 11 with WSL2: run it inside the WSL2 Ubuntu shell)
# Minimum memory: 8 GB
# Assumes: the NVIDIA driver on Linux from the distribution, or on Windows from NVIDIA's
#          Windows installer; ~/llm-course from Part 1 (override with COURSE=) and ~/models
#          from Part 4 (override with MODELS=)
#
# Usage: bash prepare-nvidia.sh [--help]
# Output: "ok" (in place), "MISSING" (a later part fails without it; the text names the
#         task that fixes it) or "NOTE" (read it, not blocking), then
#         "summary: N MISSING, M NOTE". Exit status 1 when anything is MISSING.
set -euo pipefail

MODELS="${MODELS:-$HOME/models}"
COURSE="${COURSE:-$HOME/llm-course}"
UV_PIN="0.12.11"
HF_MIN="1.30.0"
INSTALL_HINT="sudo apt install"

for arg in "$@"; do
  case "$arg" in
    -h|--help) sed -n '2,16p' "$0"; exit 0 ;;
    *) echo "prepare-nvidia: unknown argument '$arg' (the script only reports; try --help)" >&2; exit 2 ;;
  esac
done

MISSING=0
NOTES=0
have()    { command -v "$1" >/dev/null 2>&1; }
section() { printf '\n== %s\n' "$1"; }
ok()      { printf '   ok       %-20s %s\n' "$1" "$2"; }
missing() { MISSING=$((MISSING + 1)); printf '   MISSING  %-20s %s\n' "$1" "$2"; }
note()    { NOTES=$((NOTES + 1)); printf '   NOTE     %-20s %s\n' "$1" "$2"; }
# version_ge A B: true when version A >= version B
version_ge() {
  awk -v a="$1" -v b="$2" 'BEGIN {
    na = split(a, x, /[.-]/); nb = split(b, y, /[.-]/); n = (na > nb) ? na : nb
    for (i = 1; i <= n; i++) { xi = (i <= na) ? x[i] + 0 : 0; yi = (i <= nb) ? y[i] + 0 : 0
      if (xi > yi) exit 0; if (xi < yi) exit 1 }
    exit 0 }'
}
first_match() { awk -v pat="$1" '$0 ~ pat && !done { print; done = 1 }'; }

IS_WSL="no"
if grep -qi microsoft /proc/version 2>/dev/null; then IS_WSL="yes"; fi

printf 'prepare-nvidia.sh on %s, %s\n' "$(hostname)" "$(date '+%Y-%m-%d %H:%M')"

section "1. Machine and operating system"
ok "architecture" "$(uname -m)"
kernel="$(uname -r)"
ok "kernel" "$kernel"
if [ "$IS_WSL" = "yes" ]; then
  # The CUDA on WSL guide recommends WSL kernel 5.10.16.3 or later
  if version_ge "$kernel" "5.10.16.3"; then
    ok "WSL2" "yes, kernel at or above the 5.10.16.3 the CUDA on WSL guide recommends"
  else
    missing "WSL2" "kernel $kernel is older than 5.10.16.3: run wsl --update in PowerShell (task 2)"
  fi
else
  ok "WSL2" "no, native Linux"
fi
if [ -r /etc/os-release ]; then
  # shellcheck disable=SC1091
  . /etc/os-release
  ok "distribution" "${PRETTY_NAME:-unknown}"
fi
mem_kib="$(awk '/^MemTotal:/ {print $2}' /proc/meminfo)"
ok "memory (MemTotal)" "$(awk -v k="$mem_kib" 'BEGIN {printf "%.1f GiB of system memory (not VRAM)", k / 1048576}')"

section "2. Driver and card"
smi=""
if have nvidia-smi; then smi="nvidia-smi"; elif [ -x /usr/lib/wsl/lib/nvidia-smi ]; then smi="/usr/lib/wsl/lib/nvidia-smi"; fi
if [ -n "$smi" ]; then
  gpus="$("$smi" --query-gpu=name,driver_version,memory.total --format=csv,noheader 2>/dev/null || true)"
  if [ -n "$gpus" ]; then
    while IFS= read -r gpu; do ok "GPU, driver, VRAM" "$gpu"; done <<< "$gpus"
    driver="$(printf '%s\n' "$gpus" | awk -F', ' 'NR == 1 {print $2}')"
    if ! version_ge "$driver" "580"; then
      note "driver" "$driver is older than 580, the minimum for CUDA 13 (task 2)"
    fi
  else
    missing "GPU" "$smi runs but lists no GPU: see Troubleshooting"
  fi
  [ "$smi" = "nvidia-smi" ] || note "nvidia-smi" "found only at $smi; add /usr/lib/wsl/lib to PATH"
else
  if [ "$IS_WSL" = "yes" ]; then
    missing "nvidia-smi" "not found: install or update the NVIDIA driver on Windows, then wsl --update (task 2)"
  else
    missing "nvidia-smi" "not found: install the driver (task 2)"
  fi
fi
if [ "$IS_WSL" = "yes" ] && have dpkg; then
  linux_drivers="$(dpkg -l 2>/dev/null | awk '$1 == "ii" && ($2 ~ /^nvidia-driver-/ || $2 == "nvidia-open" || $2 ~ /^cuda-drivers/) {printf "%s ", $2}')"
  if [ -n "$linux_drivers" ]; then
    missing "Linux GPU driver" "installed inside WSL (${linux_drivers% }); the CUDA on WSL guide forbids it: Troubleshooting"
  else
    ok "Linux GPU driver" "none inside WSL, as the CUDA on WSL guide requires"
  fi
fi

section "3. CUDA toolkit"
if have nvcc; then
  ok "nvcc" "$(nvcc --version | first_match release)"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
  note "nvcc" "installed at /usr/local/cuda/bin but not on PATH: task 3"
else
  missing "nvcc" "not found; Part 6 compiles llama.cpp with it: task 3"
fi

section "4. Containers"
if have docker; then
  ok "docker" "$(docker --version 2>/dev/null)"
  # shellcheck disable=SC2016  # a Go template for docker, not a shell expansion
  if runtimes="$(docker info --format '{{range $name, $rt := .Runtimes}}{{$name}} {{end}}' 2>/dev/null)"; then
    case " $runtimes " in
      *" nvidia "*) ok "nvidia runtime" "registered (runtimes: ${runtimes% })" ;;
      *)
        if [ "$IS_WSL" = "yes" ]; then
          note "nvidia runtime" "not listed (${runtimes% }); with Docker Desktop, task 6's --gpus test is the check"
        else
          missing "nvidia runtime" "not among the runtimes docker reports (${runtimes% }): task 6"
        fi ;;
    esac
  else
    missing "docker access" "'docker info' failed for $(id -un): daemon stopped, or not in the docker group (task 6)"
  fi
  if compose="$(docker compose version --short 2>/dev/null)"; then
    ok "docker compose" "$compose"
  else
    missing "docker compose" "the Compose plugin is not installed; Part 7 needs it (task 6)"
  fi
else
  missing "docker" "not found: task 6"
fi
if [ "$IS_WSL" = "yes" ]; then
  case "$(cd "$HOME" && pwd -P)/" in
    /mnt/*) missing "home directory" "is on the Windows filesystem; keep the course and models in the WSL2 filesystem" ;;
    *) ok "home directory" "inside the WSL2 filesystem" ;;
  esac
fi

section "5. Toolchain and the Part 1 course directory"
for tool in git python3 curl; do
  if have "$tool"; then
    ok "$tool" "$("$tool" --version 2>&1 | awk 'NR == 1 {print $1, $2, $3}')"
  else
    missing "$tool" "install it with: $INSTALL_HINT $tool"
  fi
done
if have uv; then
  uv_version="$(uv --version | awk '{print $2}')"
  if [ "$uv_version" = "$UV_PIN" ]; then ok "uv" "$uv_version (the course pin)"; else note "uv" "$uv_version; the course is pinned to $UV_PIN (task 5)"; fi
else
  missing "uv" "not found on PATH: task 5"
fi
if [ -x "$COURSE/.venv/bin/python" ]; then ok "course environment" "$COURSE/.venv"; else missing "course environment" "$COURSE/.venv not found: Part 1's lab, tasks 2 and 3"; fi
if [ -f "$COURSE/labbook.md" ]; then ok "lab notebook" "$COURSE/labbook.md"; else missing "lab notebook" "$COURSE/labbook.md not found: Part 1's lab, task 1"; fi

section "6. Hugging Face CLI, login and model storage"
hf_bin=""
if have hf; then hf_bin="$(command -v hf)"; elif [ -x "$COURSE/.venv/bin/hf" ]; then hf_bin="$COURSE/.venv/bin/hf"; fi
if [ -n "$hf_bin" ]; then
  hf_version="$("$hf_bin" version --format json 2>/dev/null | sed -n 's/.*"version": *"\([^"]*\)".*/\1/p')"
  if [ -n "$hf_version" ] && version_ge "$hf_version" "$HF_MIN"; then
    ok "hf" "$hf_version ($hf_bin)"
  else
    missing "hf" "version '${hf_version:-unknown}' is older than $HF_MIN: task 7"
  fi
else
  missing "hf" "not on PATH and not in $COURSE/.venv: task 7"
fi
if [ -d "$MODELS" ]; then ok "model library" "$MODELS"; else missing "model library" "$MODELS not found: Part 4's lab, task 7, or task 7 here"; fi
if [ -f "$MODELS/README.md" ]; then ok "library README" "$MODELS/README.md"; else note "library README" "missing: Part 4's lab, task 7, writes it"; fi
if [ -d "$MODELS/hf" ]; then ok "HF_HOME directory" "$MODELS/hf"; else missing "HF_HOME directory" "$MODELS/hf not found: task 7"; fi
if [ "${HF_HOME:-}" = "$MODELS/hf" ]; then
  ok "HF_HOME" "$HF_HOME"
else
  missing "HF_HOME" "is '${HF_HOME:-unset}' in this shell, expected $MODELS/hf: task 7"
fi
if [ -f "$MODELS/hf/token" ]; then
  ok "stored login" "$MODELS/hf/token"
elif [ -f "$HOME/.cache/huggingface/token" ]; then
  note "stored login" "only under ~/.cache/huggingface, which hf stops reading once HF_HOME moves: task 7"
else
  note "stored login" "none; needed only for gated models (task 7)"
fi
df_target="$MODELS"; [ -d "$df_target" ] || df_target="$HOME"
ok "free disk" "$(df -Pk "$df_target" | awk 'NR == 2 {printf "%d GiB", $4 / 1048576}') on the filesystem holding $df_target"

printf '\nsummary: %d MISSING, %d NOTE\n' "$MISSING" "$NOTES"
[ "$MISSING" -eq 0 ]