#!/usr/bin/env bash
# Purpose: check a DGX Spark against everything the rest of the course assumes: DGX OS,
#          the GPU driver, 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; print one line per item and a summary. It changes nothing.
# Platform: spark (DGX Spark and the OEM systems on the same GB10 board, aarch64)
# Minimum memory: 8 GB
# Assumes: DGX OS; the course directory ~/llm-course from Part 1 (override with COURSE=)
#          and the model library ~/models from Part 4 (override with MODELS=)
#
# Usage: bash prepare-spark.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,14p' "$0"; exit 0 ;;
    *) echo "prepare-spark: 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 }'; }

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

section "1. Machine and operating system"
arch="$(uname -m)"
if [ "$arch" = "aarch64" ]; then
  ok "architecture" "$arch"
else
  missing "architecture" "$arch is not a GB10 machine: use the script for your own track"
fi
ok "kernel" "$(uname -r)"
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, shared by CPU and GPU", k / 1048576}')"

section "2. GPU and driver"
if have nvidia-smi; then
  gpu="$(nvidia-smi --query-gpu=name,driver_version --format=csv,noheader 2>/dev/null | first_match . || true)"
  if [ -n "$gpu" ]; then
    ok "GPU, driver" "$gpu"
    driver="${gpu##*, }"
    if ! version_ge "$driver" "580"; then
      note "driver" "$driver is older than the 580 series CUDA 13 needs: task 2 updates DGX OS"
    fi
  else
    missing "GPU" "nvidia-smi runs but lists no GPU: see Troubleshooting"
  fi
else
  missing "nvidia-smi" "not found; DGX OS ships the driver, so see the DGX Spark user guide"
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; DGX OS includes the CUDA toolkit: task 2, then task 3"
fi

section "4. Containers"
if have docker; then
  ok "docker" "$(docker --version 2>/dev/null)"
  docker_cmd=""
  if docker info >/dev/null 2>&1; then
    docker_cmd="docker"
    ok "docker access" "'docker info' works for $(id -un)"
  elif sudo -n docker info >/dev/null 2>&1; then
    docker_cmd="sudo -n docker"
    note "docker access" "works with sudo only; fine on DGX OS, or join the docker group (task 6)"
  else
    note "docker access" "'docker info' failed for $(id -un) without a password; run as a docker group member (task 6)"
  fi
  if [ -n "$docker_cmd" ]; then
    # shellcheck disable=SC2016  # a Go template for docker, not a shell expansion
    runtimes="$($docker_cmd info --format '{{range $name, $rt := .Runtimes}}{{$name}} {{end}}' 2>/dev/null || true)"
    case " $runtimes " in
      *" nvidia "*) ok "nvidia runtime" "registered (runtimes: ${runtimes% })" ;;
      # Inferred, not observed on a Spark: NVIDIA's own check is the --gpus=all container
      *) note "nvidia runtime" "not listed (${runtimes% }); task 6's --gpus=all container test is the check, but Part 7's --runtime=nvidia needs it: Troubleshooting" ;;
    esac
    if compose="$($docker_cmd 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
  fi
else
  missing "docker" "not found, which is unusual on DGX OS: task 6"
fi
if have nvidia-ctk; then
  ok "nvidia-ctk" "$(nvidia-ctk --version 2>/dev/null | first_match .)"
else
  note "nvidia-ctk" "not found; the user guide says the container toolkit is preinstalled"
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 ]