#!/usr/bin/env bash
# Purpose: check a Ryzen AI Max+ 395 (Strix Halo) machine against everything the rest of
#          the course assumes: the kernel AMD's Ryzen ROCm page names, the render and video
#          groups, how much memory the GPU may address (VRAM, GTT, the TTM page limit), the
#          Vulkan and ROCm paths, Docker with Compose, the toolchain, the Part 1 environment
#          and notebook, the hf CLI and the model storage. It changes nothing.
# Platform: strix (AMD Ryzen AI Max+ 395, gfx1151, on Linux; Ubuntu 24.04 is the path
#          AMD's Ryzen installation page documents)
# Minimum memory: 8 GB
# Assumes: Linux with the inbox amdgpu kernel driver; ~/llm-course from Part 1 (override
#          with COURSE=) and ~/models from Part 4 (override with MODELS=)
#
# Usage: bash prepare-strix.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"
PAGE_BYTES=4096   # the TTM limit is counted in pages of this size

for arg in "$@"; do
  case "$arg" in
    -h|--help) sed -n '2,16p' "$0"; exit 0 ;;
    *) echo "prepare-strix: 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-strix.sh on %s, %s\n' "$(hostname)" "$(date '+%Y-%m-%d %H:%M')"

section "1. Machine, operating system and kernel"
ok "architecture" "$(uname -m)"
kernel="$(uname -r)"
VERSION_ID=""
if [ -r /etc/os-release ]; then
  # shellcheck disable=SC1091
  . /etc/os-release
  ok "distribution" "${PRETTY_NAME:-unknown}"
fi
# AMD's Ryzen ROCm page: "the 6.14-1018 OEM kernel or newer" on Ubuntu 24.04
kernel_num="$(printf '%s' "$kernel" | cut -d- -f1,2 | tr '-' '.')"
if [ "$VERSION_ID" = "24.04" ]; then
  if version_ge "$kernel_num" "6.14.0.1018" && case "$kernel" in *-oem) true ;; *) false ;; esac; then
    ok "kernel" "$kernel (OEM, 6.14-1018 or newer)"
  elif version_ge "$kernel_num" "6.14.0.1018"; then
    note "kernel" "$kernel is new enough but is not the OEM kernel AMD's Ryzen page names: task 2"
  else
    missing "kernel" "$kernel is older than the 6.14-1018 OEM kernel AMD requires for ROCm: task 2"
  fi
else
  note "kernel" "$kernel on ${PRETTY_NAME:-this distribution}; AMD's Ryzen page documents Ubuntu 24.04 only"
fi
mem_kib="$(awk '/^MemTotal:/ {print $2}' /proc/meminfo)"
ok "memory (MemTotal)" "$(awk -v k="$mem_kib" 'BEGIN {printf "%.2f GiB", k / 1048576}')"
for group in render video; do
  if id -nG | tr ' ' '\n' | grep -qx "$group"; then
    ok "group $group" "$(id -un) is a member"
  else
    missing "group $group" "$(id -un) is not a member (log out and in after adding): task 3"
  fi
done

section "2. How much memory the GPU may address"
if have lspci; then
  ok "PCI display device" "$(lspci | awk 'tolower($0) ~ /vga|display/ && !done { sub(/^[^ ]+ /, ""); print; done = 1 }')"
fi
found_card="no"
for dev in /sys/class/drm/card*/device; do
  [ -r "$dev/mem_info_gtt_total" ] || continue
  found_card="yes"
  card="$(basename "$(dirname "$dev")")"
  vram="$(cat "$dev/mem_info_vram_total")"
  gtt="$(cat "$dev/mem_info_gtt_total")"
  ok "$card VRAM" "$(awk -v b="$vram" 'BEGIN {printf "%.2f GiB (the firmware carve-out)", b / 1073741824}')"
  ok "$card GTT" "$(awk -v b="$gtt" 'BEGIN {printf "%.2f GiB (system memory the GPU may map)", b / 1073741824}')"
  if [ "$vram" -gt 644245094 ]; then   # 0.6 GiB
    note "$card VRAM" "AMD recommends the minimum carve-out (0.5 GB) in the firmware settings, and a larger TTM limit"
  fi
done
[ "$found_card" = "yes" ] || missing "amdgpu" "no card exposes mem_info_gtt_total: the amdgpu driver is not loaded"
if [ -r /sys/module/ttm/parameters/pages_limit ]; then
  pages="$(cat /sys/module/ttm/parameters/pages_limit)"
  limit_line="$(awk -v p="$pages" -v s="$PAGE_BYTES" -v k="$mem_kib" \
    'BEGIN {g = p * s / 1073741824; printf "%d pages = %.2f GiB = %.0f%% of MemTotal", p, g, 100 * g / (k / 1048576)}')"
  percent="$(awk -v p="$pages" -v s="$PAGE_BYTES" -v k="$mem_kib" 'BEGIN {printf "%d", 100 * p * s / 1024 / k}')"
  if [ "$percent" -le 55 ]; then
    note "TTM pages_limit" "$limit_line: still about the kernel default of half; task 4 raises it"
  else
    ok "TTM pages_limit" "$limit_line"
  fi
else
  note "TTM pages_limit" "/sys/module/ttm/parameters/pages_limit not readable"
fi
if [ -f /etc/modprobe.d/ttm.conf ]; then
  ok "ttm.conf" "$(cat /etc/modprobe.d/ttm.conf)"
fi

section "3. The two GPU software paths"
if have vulkaninfo; then
  summary="$(vulkaninfo --summary 2>/dev/null || true)"
  names="$(printf '%s\n' "$summary" | awk -F'= ' '/deviceName/ {printf "%s%s", sep, $2; sep = "; "}')"
  drivers="$(printf '%s\n' "$summary" | awk -F'= ' '/driverName/ {printf "%s%s", sep, $2; sep = "; "}')"
  if [ -n "$names" ]; then
    ok "Vulkan devices" "$names"
    ok "Vulkan drivers" "${drivers:-not reported}"
  else
    missing "Vulkan devices" "vulkaninfo --summary lists no device: task 3"
  fi
else
  missing "vulkaninfo" "not found: task 3 installs mesa-vulkan-drivers and vulkan-tools"
fi
if have rocminfo; then
  gfx="$(rocminfo 2>/dev/null | awk '$1 == "Name:" && $2 ~ /^gfx/ && !done {print $2; done = 1}')"
  if [ -n "$gfx" ]; then
    ok "ROCm agent" "$gfx"
  else
    missing "ROCm agent" "rocminfo lists no gfx agent: groups, then a reboot (task 3)"
  fi
  if have amd-smi; then ok "amd-smi" "$(amd-smi version 2>/dev/null | first_match .)"; fi
else
  note "ROCm" "rocminfo not found: ROCm is not installed (task 3; PyTorch on the GPU needs it)"
fi
if have dpkg-query && dpkg-query -W -f='${db:Status-Status}\n' amdgpu-dkms 2>/dev/null | grep -qx installed; then
  missing "amdgpu-dkms" "installed, but AMD's Ryzen page requires the inbox driver: Troubleshooting"
fi

section "4. Containers"
if have docker; then
  ok "docker" "$(docker --version 2>/dev/null)"
  if docker info >/dev/null 2>&1; then
    ok "docker access" "'docker info' works for $(id -un)"
  else
    missing "docker access" "'docker info' failed for $(id -un): not in the docker group yet (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
elif have podman; then
  note "podman" "$(podman --version 2>/dev/null); Part 7's lab is written for Docker Engine with Compose"
else
  missing "docker" "neither docker nor podman found: task 6"
fi
if [ -e /dev/kfd ]; then
  ok "/dev/kfd" "present (ROCm opens it; Part 7 passes it into containers)"
else
  note "/dev/kfd" "absent: ROCm and ROCm containers cannot reach the GPU without it (Troubleshooting)"
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 ]