# infra/install.sh: put this machine to work for Unitfabric (Plan U2). The site serves it; as a user who may sudo:
#   curl -fsSL https://unitfabric-site.vercel.app/install.sh | UNITFABRIC_SUPPLY_TOKEN=<uf_supply_...> sudo -E bash -s -- --control-url <url>
# The token is one of the org's supply tokens (the dashboard's Add a machine, or `unitfabric supply tokens create`). It
# enrols machines, each pending until the org's owner confirms it, and does nothing else; it works for a day, and the org
# revokes it at once (the Machines page, or `unitfabric supply tokens revoke`). The line above, as typed, stays in the
# shell's history for that day: leave UNITFABRIC_SUPPLY_TOKEN out of it and the installer asks for the token on the
# terminal instead, unechoed. Never an argument, which any user's `ps` could read; used once, to enrol; never on disk.
# What this does, idempotently (a second run changes only what is missing or out of date, and restarts the agent only then):
#   1. checks the machine (--check stops here, changing nothing: a report, then exit 0, or 1 when a check failed);
#   2. installs Podman 4.4 or newer and Python 3.12 or newer from the distribution, on Ubuntu 24.04, Debian 13, Fedora 42
#      and newer, RHEL, Rocky Linux and AlmaLinux 9 and 10, openSUSE Leap 15.6 and Tumbleweed, Arch, and their
#      derivatives. Podman 4.4: the CDI device support the agent gives a worker its GPU by (`--device nvidia.com/gpu=N`).
#      Not served, with the reason: Ubuntu 22.04 (Podman 3.4, without CDI: use 24.04) and Debian 12 (Podman 4.3 and
#      Python 3.11: use 13). On Arch it never upgrades the system: --check says when `pacman -Syu` is due first, and the
#      install adds only its own packages (--needed);
#   3. with an NVIDIA driver loaded (never installed here: a missing driver is said, and the machine runs CPU templates),
#      the NVIDIA Container Toolkit and its CDI spec (--mode=wsl on WSL2, whose one device is `all`); on Linux, CRIU where
#      the distribution packages a 4.x with its CUDA plugin, and NVIDIA's cuda-checkpoint, which together let the agent
#      checkpoint a GPU worker (never on WSL2: such a machine restarts units instead);
#   4. the agent's package from the site, checked against the sha256 the site serves beside it (or --agent-sha256), in
#      /opt/grid/agent/.venv with pip and setuptools pinned and its dependencies at the versions agent/constraints.txt pins;
#      /etc/grid/agent.toml (mode 600; network "none": a supplier's machine runs jobs whose inputs are uploads or inline,
#      with no network; userns "auto": each worker in a user namespace of its own, from the ranges /etc/subuid and
#      /etc/subgid give "containers"), and its systemd unit, grid-agent.service;
#   5. enrols with the token (the rate given with --rate, or the control plane's suggestion for its GPU class), benchmarks
#      the templates it can run (again only when the agent changed or none were reported yet), and starts the agent.
# On WSL2 (/proc/version names Microsoft) systemd must run: a run without it turns it on in /etc/wsl.conf and exits 3, and
# `wsl --terminate <the distribution>` in Windows then a second run finish (infra/install.ps1 turns it on first, in the
# distribution of its own it installs into). Exit status: 0 done; 1 a step failed; 2 usage, not root, or a machine this
# cannot serve; 3 restart the WSL distribution and run again. The whole install is main(), called on the last line: a
# download cut short defines what it got and runs nothing.
set -euo pipefail

SITE=${UNITFABRIC_SITE_URL:-https://unitfabric-site.vercel.app}    # another when the install line names one, or a test's copy
SITE=${SITE%/}
AGENT_URL=$SITE/grid-agent.tar.gz                 # the agent package, which the site serves beside this script (Plan U4)
CUDA_CKPT_URL=https://raw.githubusercontent.com/NVIDIA/cuda-checkpoint/00d5cce84c628088d6caa203fc4af40c1538b6f7/bin/x86_64_Linux/cuda-checkpoint
CUDA_CKPT_SHA256=707fa7f54136824d6c1d6dd724b9b1717610f831033c00d06da474de363a06db    # as infra/provision_host.sh pins it
MIN_PODMAN=4.4
MIN_PYTHON=3.12
PIP=pip==26.2.1                                  # the venv's installer and the agent's build backend, pinned as its
SETUPTOOLS=setuptools==84.0.0                    # dependencies are: what every machine builds with is what was tested
VENV=/opt/grid/agent/.venv
PACKAGE_SUM=/opt/grid/agent/package.sha256       # the agent's files as last installed: an unchanged package is not reinstalled
STATE=/var/lib/grid-agent/node.json
BENCHED=/var/lib/grid-agent/bench.json           # what the last install's benchmarks reported

usage() {
  cat >&2 <<'USAGE'
usage: UNITFABRIC_SUPPLY_TOKEN=TOKEN install.sh --control-url URL [--name NAME] [--rate USD] [--cpu] [--gpu-index N]
                  [--agent-url URL [--agent-sha256 HEX] | --agent-dir DIR] [--no-bench] [--no-start] [--reenrol]
       install.sh --check [--control-url URL]
  UNITFABRIC_SUPPLY_TOKEN  the org's supply token, used once to enrol; left out, asked for on the terminal
  UNITFABRIC_SITE_URL      where the agent's package is (default: the site; a control plane served from another names it
                           in its install lines), https; a test serves a copy on localhost
  --agent-sha256 the package's sha256 (default: the one the site serves beside it, in grid-agent.tar.gz.sha256)
  --control-url  the control plane (or UF_CONTROL_URL), https
  --name         the machine's name (default: its hostname)
  --rate         what it asks, USD an hour (default: the control plane's suggestion for its GPU class)
  --cpu          CPU templates only, even with a GPU
  --agent-dir    install the agent from this directory (a checkout's agent/) instead of the site's package
  --no-bench     enrol without benchmarking (`grid-agent bench` later; no work comes until it has run)
  --no-start     install and enrol, but leave grid-agent.service stopped (a container without systemd)
  --reenrol      forget this machine's enrolment and enrol it again, as a new machine
  --check        report on the machine and change nothing
USAGE
  exit 2
}

main() {
TOKEN=${UNITFABRIC_SUPPLY_TOKEN:-}; CONTROL_URL=${UF_CONTROL_URL:-}; NAME=""; RATE=""; FORCE_CPU=""; GPU_INDEX=0
unset UNITFABRIC_SUPPLY_TOKEN                     # no process this starts inherits it: enrol reads it once from its stdin
AGENT_DIR=""; AGENT_SHA256=""; NO_BENCH=""; NO_START=""; REENROL=""; CHECK=""
while [ $# -gt 0 ]; do
  case "$1" in
    --token|--token=*) echo "install.sh: the token is never an argument (any user's ps reads a command line): set UNITFABRIC_SUPPLY_TOKEN, or leave it out to be asked" >&2; exit 2 ;;
    --control-url) CONTROL_URL=${2:-}; shift 2 ;;
    --name) NAME=${2:-}; shift 2 ;;
    --rate) RATE=${2:-}; shift 2 ;;
    --cpu) FORCE_CPU=1; shift ;;
    --gpu-index) GPU_INDEX=${2:-}; shift 2 ;;
    --agent-url) AGENT_URL=${2:-}; shift 2 ;;
    --agent-sha256) AGENT_SHA256=${2:-}; shift 2 ;;
    --agent-dir) AGENT_DIR=${2:-}; shift 2 ;;
    --no-bench) NO_BENCH=1; shift ;;
    --no-start) NO_START=1; shift ;;
    --reenrol) REENROL=1; shift ;;
    --check) CHECK=1; shift ;;
    -h|--help) usage ;;
    *) echo "install.sh: unknown argument $1" >&2; usage ;;
  esac
done
CONTROL_URL=${CONTROL_URL%/}
if [ -z "$CHECK" ]; then
  [ -n "$CONTROL_URL" ] || { echo "install.sh: --control-url is required" >&2; exit 2; }
  if [ -z "$TOKEN" ] && { : < /dev/tty; } 2>/dev/null; then      # the script is curl's stdin: the terminal is /dev/tty
    printf 'Supply token (uf_supply_...): ' > /dev/tty
    IFS= read -rs TOKEN < /dev/tty || true
    printf '\n' > /dev/tty
  fi
  [[ $TOKEN =~ ^uf_supply_[0-9a-f]{40}$ ]] || { echo "install.sh: UNITFABRIC_SUPPLY_TOKEN must be a supply token, uf_supply_ and 40 hex digits" >&2; exit 2; }
fi
url_ok() { [[ $1 =~ ^https://[A-Za-z0-9.-]+(:[0-9]{1,5})?$ || $1 =~ ^http://(localhost|127\.0\.0\.1)(:[0-9]{1,5})?$ ]]; }
if [ -n "$CONTROL_URL" ] && ! url_ok "$CONTROL_URL"; then
  echo "install.sh: --control-url must be https://host[:port] (plain http only for localhost)" >&2; exit 2
fi
url_ok "$SITE" || { echo "install.sh: UNITFABRIC_SITE_URL must be https://host[:port] (plain http only for localhost)" >&2; exit 2; }
rate_ok() {          # USD an hour, above 0 and at most 1000, four decimals at most; in bash alone (a minimal system has no awk)
  [[ $1 =~ ^([0-9]{1,4})(\.([0-9]{1,4}))?$ ]] || return 1
  local whole=$((10#${BASH_REMATCH[1]})) part=$((10#${BASH_REMATCH[3]:-0}))
  { [ "$whole" -lt 1000 ] && [ $((whole + part)) -gt 0 ]; } || { [ "$whole" = 1000 ] && [ "$part" = 0 ]; }
}
if [ -n "$RATE" ] && ! rate_ok "$RATE"; then
  echo "install.sh: --rate is USD an hour, above 0 and at most 1000, with at most four decimals" >&2; exit 2
fi
[[ $GPU_INDEX =~ ^[0-9]{1,2}$ ]] || { echo "install.sh: --gpu-index is a number" >&2; exit 2; }
[ -z "$AGENT_SHA256" ] || [[ $AGENT_SHA256 =~ ^[0-9a-f]{64}$ ]] || { echo "install.sh: --agent-sha256 is 64 hex digits" >&2; exit 2; }
if [ -z "$NAME" ]; then NAME=$(uname -n); fi
NAME=$(printf '%s' "$NAME" | tr -cd '[:alnum:]._ -' | cut -c1-64)
[ -n "$NAME" ] || NAME=machine

FAILED=0
ok() { printf 'ok    %-12s %s\n' "$1" "$2"; }
warn() { printf 'warn  %-12s %s\n' "$1" "$2"; }
fail() { printf 'FAIL  %-12s %s\n' "$1" "$2"; FAILED=1; }
die() { echo "install.sh: $*" >&2; exit 1; }
version_ge() { [ "$(printf '%s\n%s\n' "$2" "$1" | sort -V | head -1)" = "$2" ]; }   # version_ge HAVE NEED
same() { [ -f "$2" ] && [ "$(sha256sum < "$1")" = "$(sha256sum < "$2")" ]; }   # same NEW OLD: cmp is not on every system

# --- What this machine is -------------------------------------------------------------------------------------------
# shellcheck source=/dev/null
OS_ID=$(. /etc/os-release 2>/dev/null && echo "${ID:-}"); OS_LIKE=$(. /etc/os-release 2>/dev/null && echo "${ID_LIKE:-}")
OS_NAME=$(. /etc/os-release 2>/dev/null && echo "${PRETTY_NAME:-unknown}"); OS_VERSION=$(. /etc/os-release 2>/dev/null && echo "${VERSION_ID:-}")
FAMILY=""
for id in $OS_ID $OS_LIKE; do
  case "$id" in
    debian|ubuntu) FAMILY=apt ;;
    fedora|rhel|centos|rocky|almalinux) FAMILY=dnf ;;
    arch) FAMILY=pacman ;;
    opensuse-tumbleweed|opensuse-leap|opensuse|suse) FAMILY=zypper ;;
  esac
  [ -z "$FAMILY" ] || break
done
PLATFORM=linux
grep -qi microsoft /proc/version 2>/dev/null && PLATFORM=wsl2
SYSTEMD=""
[ -d /run/systemd/system ] && SYSTEMD=1
NVSMI=$(command -v nvidia-smi || true)
[ -n "$NVSMI" ] || { [ -x /usr/lib/wsl/lib/nvidia-smi ] && NVSMI=/usr/lib/wsl/lib/nvidia-smi; } || true
GPU_NAME=""
if [ -n "$NVSMI" ] && [ -z "$FORCE_CPU" ]; then
  GPU_NAME=$("$NVSMI" --id="$GPU_INDEX" --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || true)
fi
MODE=cpu; [ -z "$GPU_NAME" ] || MODE=gpu

python_ok() {        # the first interpreter of MIN_PYTHON or newer with venv, on stdout
  local py
  for py in python3.13 python3.12 python3; do
    if command -v "$py" >/dev/null 2>&1 && "$py" -c 'import sys, venv, ensurepip; sys.exit(sys.version_info < (3, 12))' 2>/dev/null; then
      command -v "$py"; return 0
    fi
  done
  return 1
}
podman_version() { podman version --format '{{.Client.Version}}' 2>/dev/null || podman --version 2>/dev/null | { read -r _ _ v; echo "$v"; }; }
criu_version() { criu --version 2>/dev/null | while read -r k v _; do [ "$k" = Version: ] && echo "$v" && break; done; }
release_unserved() {   # why this release of a distribution is not served (nothing: it is, or a derivative the versions decide)
  local major=${OS_VERSION%%.*}
  case "$OS_ID" in
    ubuntu) version_ge "$OS_VERSION" 24.04 || echo "Ubuntu $OS_VERSION's Podman is 3.4, without the CDI support the GPU is given by: use Ubuntu 24.04" ;;
    debian) [ -z "$major" ] || version_ge "$major" 13 || echo "Debian $OS_VERSION's Podman is 4.3 and its Python 3.11, under the agent's $MIN_PODMAN and $MIN_PYTHON: use Debian 13" ;;
    fedora) version_ge "$OS_VERSION" 42 || echo "Fedora $OS_VERSION is past its end of life: use Fedora 42 or newer" ;;
    rhel|rocky|almalinux|centos) version_ge "$major" 9 || echo "$OS_NAME is too old: use release 9 or 10" ;;
    opensuse-leap) version_ge "$OS_VERSION" 15.6 || echo "openSUSE Leap $OS_VERSION is past its end of life: use 15.6, or Tumbleweed" ;;
  esac
}
pacman_due() {         # Arch: why a system upgrade must come first (nothing: none is due)
  [ "$FAMILY" = pacman ] || return 0
  ls /var/lib/pacman/sync/*.db >/dev/null 2>&1 || { echo "no package databases yet: run pacman -Syu first"; return 0; }
  local due; due=$(pacman -Qu 2>/dev/null | wc -l)
  [ "$due" = 0 ] || echo "$due packages have upgrades pending: run pacman -Syu first (this installer never upgrades the system, and a partial upgrade breaks Arch)"
}
criu_ok() {          # CRIU 4 or newer with its CUDA plugin: what the agent's checkpoint needs beside cuda-checkpoint
  local v d
  v=$(criu_version)
  [ -n "$v" ] && version_ge "$v" 4.0 || return 1
  for d in /usr/local/lib/criu /usr/lib/criu /usr/lib64/criu /usr/libexec/criu; do [ -f "$d/cuda_plugin.so" ] && return 0; done
  return 1
}

preflight() {
  echo "Unitfabric: checking $(uname -n) ($OS_NAME, $(uname -m), $PLATFORM)"
  if [ "$(id -u)" -eq 0 ]; then ok root "running as root"; elif [ -n "$CHECK" ]; then warn root "not root: the install itself needs sudo"; else fail root "run as root: sudo bash"; fi
  if [ "$(uname -m)" = x86_64 ]; then ok arch "x86_64"; else fail arch "$(uname -m): the templates' images are built for x86_64 only"; fi
  local why; why=$(release_unserved)
  if [ -z "$FAMILY" ]; then fail distro "$OS_NAME: not a distribution this installs on (apt, dnf, pacman or zypper)"
  elif [ -n "$why" ]; then fail distro "$why"
  else ok distro "$OS_NAME ($FAMILY)"; fi
  why=$(pacman_due)
  if [ -n "$why" ]; then fail upgrade "$why"; elif [ "$FAMILY" = pacman ]; then ok upgrade "the system is up to date: the install adds only its own packages"; fi
  if [ -n "$SYSTEMD" ]; then ok systemd "systemd runs the agent's service"
  elif [ "$PLATFORM" = wsl2 ]; then warn systemd "not running in this WSL2 distribution: the install turns it on, then the distribution restarts (wsl --terminate)"
  elif [ -n "$NO_START" ]; then warn systemd "not running: with --no-start, start the agent yourself (grid-agent run)"
  else fail systemd "not running: the agent runs as a systemd service"; fi
  local py; py=$(python_ok || true)
  if [ -n "$py" ]; then ok python "$("$py" -V 2>&1) at $py"; else warn python "no Python $MIN_PYTHON or newer with venv yet: the install adds the distribution's"; fi
  local pv; pv=$(podman_version || true)
  if [ -z "$pv" ]; then warn podman "not installed yet: the install adds the distribution's ($MIN_PODMAN or newer)"
  elif version_ge "$pv" "$MIN_PODMAN"; then ok podman "$pv"; else warn podman "$pv: the install upgrades it to $MIN_PODMAN or newer, or stops"; fi
  if [ "$MODE" = gpu ]; then
    ok gpu "$GPU_NAME, driver $("$NVSMI" --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1)"
    if command -v nvidia-ctk >/dev/null 2>&1; then ok toolkit "$(nvidia-ctk --version 2>/dev/null | head -1)"; else warn toolkit "the NVIDIA Container Toolkit is not installed yet: the install adds it"; fi
    if [ "$PLATFORM" = wsl2 ]; then ok checkpoint "off on WSL2: a lost lease's units restart, and only jobs whose units are short come here"
    elif criu_ok && command -v cuda-checkpoint >/dev/null 2>&1; then ok checkpoint "CRIU $(criu_version) with its CUDA plugin, and cuda-checkpoint"
    else warn checkpoint "no CRIU 4 with its CUDA plugin yet: the install adds the distribution's where it has one; without, units restart"; fi
  elif [ -n "$FORCE_CPU" ]; then ok gpu "CPU templates only (--cpu)"
  else warn gpu "no NVIDIA driver loaded: this machine runs CPU templates only. For GPU work install your distribution's NVIDIA driver (550 or newer; on WSL2, the Windows driver) and run this again"; fi
  local free; free=$(df -Pk /var/lib 2>/dev/null | { read -r _; read -r _ _ _ avail _; echo $(( ${avail:-0} / 1048576 )); })
  if [ "${free:-0}" -ge 100 ]; then ok disk "${free} GB free for images and work"
  elif [ "${free:-0}" -ge 30 ]; then warn disk "${free} GB free: a template image is up to 15 GB; 100 GB leaves room for all of them"
  else fail disk "${free:-0} GB free: at least 30 GB for the agent's images and work"; fi
  local mem k v; while read -r k v _; do [ "$k" = MemTotal: ] && mem=$((v / 1048576)) && break; done < /proc/meminfo
  if [ "$mem" -ge 8 ]; then ok memory "${mem} GB, $(nproc) CPUs"; else warn memory "${mem} GB: templates ask for 4 to 16 GB; this machine takes the smaller ones"; fi
  if [ -z "$CONTROL_URL" ]; then :
  elif ! command -v curl >/dev/null 2>&1; then warn control "no curl yet to reach $CONTROL_URL with: the install adds it"
  elif curl -fsS -m 10 -o /dev/null "$CONTROL_URL/health"; then ok control "$CONTROL_URL answers"
  else fail control "$CONTROL_URL/health does not answer from here"; fi
  if command -v curl >/dev/null 2>&1; then
    local code; code=$(curl -s -m 10 -o /dev/null -w '%{http_code}' https://ghcr.io/v2/ || true)
    if [ "$code" = 401 ] || [ "$code" = 200 ]; then ok registry "ghcr.io answers (template images come from it)"; else warn registry "ghcr.io did not answer ($code): template images come from it"; fi
  fi
}

preflight
if [ -n "$CHECK" ]; then
  echo "check done: $([ "$FAILED" = 0 ] && echo 'this machine can be installed' || echo 'fix the FAIL lines first'); nothing was changed"
  exit "$FAILED"
fi
[ "$(id -u)" -eq 0 ] || { echo "install.sh: run as root (sudo)" >&2; exit 2; }
[ "$FAILED" = 0 ] || { echo "install.sh: fix the FAIL lines above first" >&2; exit 2; }

# --- WSL2: systemd first -------------------------------------------------------------------------------------------
if [ "$PLATFORM" = wsl2 ] && [ -z "$SYSTEMD" ] && [ -z "$NO_START" ]; then
  if ! grep -Eq '^[[:space:]]*systemd[[:space:]]*=[[:space:]]*true' /etc/wsl.conf 2>/dev/null; then
    if grep -q '^\[boot\]' /etc/wsl.conf 2>/dev/null; then
      sed -i '/^\[boot\]/a systemd=true' /etc/wsl.conf
    else
      printf '\n[boot]\nsystemd=true\n' >> /etc/wsl.conf
    fi
  fi
  echo "systemd is now on in /etc/wsl.conf. Restart this distribution from Windows (PowerShell: wsl --terminate ${WSL_DISTRO_NAME:-<its name>}), then run this again." >&2
  exit 3
fi

CHANGED=""                                  # set by every step that changed something: the agent is restarted only then
WORK=$(mktemp -d)
trap 'rm -rf "$WORK"' EXIT

# --- Packages ------------------------------------------------------------------------------------------------------
echo "== packages ($FAMILY)"
export DEBIAN_FRONTEND=noninteractive
EXTRA=()                                    # curl only where there is none: Fedora's and RHEL's curl-minimal would conflict
command -v curl >/dev/null 2>&1 || EXTRA+=(curl)
command -v tar >/dev/null 2>&1 || EXTRA+=(tar)        # a minimal system (a container image) may lack them
command -v gzip >/dev/null 2>&1 || EXTRA+=(gzip)
command -v find >/dev/null 2>&1 || EXTRA+=(findutils)  # find and xargs, which sum the agent's package (Tumbleweed's image lacks them)
case "$FAMILY" in
  apt)
    apt-get update -qq
    apt-get install -y -qq --no-install-recommends ca-certificates podman python3 python3-venv "${EXTRA[@]}" >/dev/null ;;
  dnf)
    dnf install -y -q ca-certificates podman python3 "${EXTRA[@]}" >/dev/null ;;
  pacman)                                   # never -Sy or -Syu here: the check found the system up to date
    pacman -S --needed --noconfirm ca-certificates podman python "${EXTRA[@]}" >/dev/null ;;
  zypper)
    zypper --non-interactive --quiet refresh
    zypper --non-interactive --quiet install ca-certificates podman python3 "${EXTRA[@]}" >/dev/null ;;
esac
PY=$(python_ok || true)
if [ -z "$PY" ]; then                       # an older python3 by default: the distribution's newer one, where it has it
  case "$FAMILY" in
    apt) apt-get install -y -qq python3.12 python3.12-venv >/dev/null 2>&1 || true ;;
    dnf) dnf install -y -q python3.12 python3.12-pip >/dev/null 2>&1 || true ;;      # RHEL 9's python3 is 3.9
    zypper) zypper --non-interactive --quiet install python312 python312-pip >/dev/null 2>&1 || true ;;   # Leap 15.6's is 3.6
  esac
  PY=$(python_ok || true)
fi
[ -n "$PY" ] || die "no Python $MIN_PYTHON or newer: $OS_NAME's is $(python3 -V 2>&1); upgrade the distribution"
PV=$(podman_version || true)
[ -n "$PV" ] && version_ge "$PV" "$MIN_PODMAN" || die "Podman ${PV:-missing}: $MIN_PODMAN or newer is needed, and $OS_NAME has no newer one"
echo "Python $("$PY" -c 'import platform; print(platform.python_version())'), Podman $PV"
# Every worker runs in a user namespace of its own (podman --userns=auto, agent.toml's userns): rootful podman takes their
# uids from the ranges /etc/subuid and /etc/subgid give the user "containers", which need not exist.
for f in /etc/subuid /etc/subgid; do
  grep -q '^containers:' "$f" 2>/dev/null || { echo "containers:2147483647:2147483648" >> "$f"; CHANGED=1; }
done

# --- NVIDIA: the toolkit, its CDI spec, and checkpointing where it can work -----------------------------------------
GPU_DEVICE=""
CAPABLE=false
if [ "$MODE" = gpu ]; then
  echo "== NVIDIA Container Toolkit"
  if ! command -v nvidia-ctk >/dev/null 2>&1; then
    case "$FAMILY" in
      apt)
        command -v gpg >/dev/null 2>&1 || apt-get install -y -qq gnupg >/dev/null
        curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor --yes -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
        curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
          | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
          > /etc/apt/sources.list.d/nvidia-container-toolkit.list
        apt-get update -qq && apt-get install -y -qq nvidia-container-toolkit >/dev/null ;;
      dnf)
        curl -fsSL https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo -o /etc/yum.repos.d/nvidia-container-toolkit.repo
        dnf install -y -q nvidia-container-toolkit >/dev/null ;;
      pacman)
        pacman -S --needed --noconfirm nvidia-container-toolkit >/dev/null ;;
      zypper)                               # NVIDIA's own key, imported as its repository asks, and no other repository's
        rpm --import https://nvidia.github.io/libnvidia-container/gpgkey
        zypper --non-interactive ar -f https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo >/dev/null 2>&1 || true
        zypper --non-interactive --quiet install nvidia-container-toolkit >/dev/null ;;
    esac
  fi
  CDI_MODE=""; [ "$PLATFORM" = wsl2 ] && CDI_MODE="--mode=wsl" && GPU_DEVICE=all
  mkdir -p /etc/cdi /etc/systemd/system/grid-agent.service.d
  # shellcheck disable=SC2086  # CDI_MODE is one word or none
  nvidia-ctk cdi generate $CDI_MODE --output="$WORK/nvidia.yaml" >/dev/null
  same "$WORK/nvidia.yaml" /etc/cdi/nvidia.yaml || { install -m 644 "$WORK/nvidia.yaml" /etc/cdi/nvidia.yaml; CHANGED=1; }
  cat > "$WORK/10-cdi.conf" <<UNIT
# Written by infra/install.sh: the CDI spec names the devices of the driver loaded now, so it is made again at every start
# (a driver upgrade, a new card); "-": an agent that cannot make it still starts, and its benchmarks say what is wrong.
[Service]
ExecStartPre=-$(command -v nvidia-ctk) cdi generate $CDI_MODE --output=/etc/cdi/nvidia.yaml
UNIT
  same "$WORK/10-cdi.conf" /etc/systemd/system/grid-agent.service.d/10-cdi.conf \
    || { install -m 644 "$WORK/10-cdi.conf" /etc/systemd/system/grid-agent.service.d/10-cdi.conf; CHANGED=1; }
  if [ "$(cat /sys/fs/selinux/enforce 2>/dev/null)" = 1 ] && ! getsebool container_use_devices 2>/dev/null | grep -q ' on$'; then
    setsebool -P container_use_devices 1    # SELinux labels the workers (they are not unconfined): this lets them reach the GPU
    CHANGED=1
  fi
  if [ "$PLATFORM" = linux ]; then
    echo "== checkpointing (CRIU and cuda-checkpoint)"
    # CRIU only with its CUDA plugin, without which it cannot checkpoint a GPU worker: Fedora packages the plugin as
    # criu-cuda-plugin, openSUSE as criu-plugin-cuda, Arch inside criu. Debian's and RHEL 10's CRIU have none, Ubuntu has
    # no CRIU: such a machine restarts units (a dnf without the plugin installs neither package).
    case "$FAMILY" in
      dnf) dnf install -y -q criu criu-cuda-plugin >/dev/null 2>&1 || true ;;
      pacman) pacman -S --needed --noconfirm criu >/dev/null 2>&1 || true ;;
      zypper) zypper --non-interactive --quiet install criu criu-plugin-cuda >/dev/null 2>&1 || true ;;
    esac
    if criu_ok; then
      if ! echo "$CUDA_CKPT_SHA256  /usr/local/bin/cuda-checkpoint" | sha256sum --check --status 2>/dev/null; then
        tmp=$(mktemp)
        curl -fsSLo "$tmp" "$CUDA_CKPT_URL"
        echo "$CUDA_CKPT_SHA256  $tmp" | sha256sum --check --status || { rm -f "$tmp"; die "cuda-checkpoint does not have its pinned sha256"; }
        install -m 755 "$tmp" /usr/local/bin/cuda-checkpoint && rm -f "$tmp"
      fi
    else
      echo "no CRIU 4 with its CUDA plugin from $OS_NAME: this machine restarts units instead of checkpointing them"
    fi
  fi
fi

# --- The agent -----------------------------------------------------------------------------------------------------
echo "== agent"
SRC=$WORK/agent
mkdir "$SRC"
if [ -n "$AGENT_DIR" ]; then                # copied: pip builds in the source tree, which may be read-only
  tar -C "$AGENT_DIR" --exclude=.venv --exclude=__pycache__ --exclude='*.egg-info' --exclude=.pytest_cache --exclude=build -cf - . \
    | tar -C "$SRC" -xf -
else
  curl -fsSLo "$WORK/agent.tar.gz" "$AGENT_URL"
  if [ -z "$AGENT_SHA256" ]; then           # the sum the site serves beside the package, in sha256sum's format (the hash first)
    if curl -fsSLo "$WORK/agent.tar.gz.sha256" "$AGENT_URL.sha256" 2>/dev/null; then
      read -r AGENT_SHA256 _ < "$WORK/agent.tar.gz.sha256" || true
      [[ $AGENT_SHA256 =~ ^[0-9a-f]{64}$ ]] || die "$AGENT_URL.sha256 is not a sha256sum line"
    else
      echo "warn: $AGENT_URL.sha256 is not served: the package is not checked (--agent-sha256 HEX checks it)"
    fi
  fi
  [ -z "$AGENT_SHA256" ] || echo "$AGENT_SHA256  $WORK/agent.tar.gz" | sha256sum --check --status \
    || die "$AGENT_URL is not the package its sha256 ($AGENT_SHA256) names: refused"
  tar -xzf "$WORK/agent.tar.gz" -C "$SRC" --strip-components=1
fi
[ -f "$SRC/pyproject.toml" ] && [ -f "$SRC/constraints.txt" ] && [ -f "$SRC/systemd/grid-agent-supply.service" ] || die "$SRC is not the agent's package"
mkdir -p /opt/grid/agent /etc/grid /srv/grid/work
SUM=$(cd "$SRC" && find . -type f -print0 | LC_ALL=C sort -z | xargs -0 sha256sum | sha256sum | cut -d' ' -f1)
if [ -x "$VENV/bin/grid-agent" ] && [ "$(cat "$PACKAGE_SUM" 2>/dev/null)" = "$SUM" ] \
    && "$VENV/bin/python" -c 'import sys; sys.exit(sys.version_info < (3, 12))' 2>/dev/null; then
  echo "the agent is this package already: not installed again"
else
  [ -x "$VENV/bin/python" ] && "$VENV/bin/python" -c 'import sys; sys.exit(sys.version_info < (3, 12))' 2>/dev/null || "$PY" -m venv --clear "$VENV"
  "$VENV/bin/pip" install -q "$PIP" "$SETUPTOOLS" >/dev/null
  "$VENV/bin/pip" install -q --no-build-isolation --force-reinstall --no-deps "$SRC" >/dev/null   # built by the pinned setuptools
  "$VENV/bin/pip" install -q --no-build-isolation -c "$SRC/constraints.txt" "$SRC" >/dev/null     # the dependencies the suites ran
  echo "$SUM" > "$PACKAGE_SUM"
  CHANGED=1
fi
CAPABLE=$("$VENV/bin/python" -c "from grid_agent import hw; print(str(hw.checkpoint_capable('$MODE', '$PLATFORM')).lower())")

config() {           # agent.toml; with an argument, the copy `grid-agent enrol` reads once from its stdin
  printf 'control_url = "%s"\nname = "%s"\nmode = "%s"\ngpu_index = %s\n' "$CONTROL_URL" "$NAME" "$MODE" "$GPU_INDEX"
  [ -z "$GPU_DEVICE" ] || printf 'gpu_device = "%s"\n' "$GPU_DEVICE"
  printf 'network = "none"\nuserns = "auto"\nstate_dir = "/var/lib/grid-agent"\nwork_dir = "/srv/grid/work"\nheartbeat_s = 10\npoll_wait_s = 20\n'
  if [ $# -gt 0 ]; then
    printf 'supply_token = "%s"\n' "$1"
    [ -z "$RATE" ] || printf 'rate_usd_per_hour = "%s"\n' "$RATE"
  fi
}
config | install -m 600 /dev/stdin /etc/grid/agent.toml.new
"$VENV/bin/python" -c 'import sys; from grid_agent import config; config.load(sys.argv[1])' /etc/grid/agent.toml.new \
  || { rm -f /etc/grid/agent.toml.new; die "the agent refuses the configuration written for it"; }
if same /etc/grid/agent.toml.new /etc/grid/agent.toml; then rm -f /etc/grid/agent.toml.new
else mv /etc/grid/agent.toml.new /etc/grid/agent.toml; CHANGED=1; fi
if ! same "$SRC/systemd/grid-agent-supply.service" /etc/systemd/system/grid-agent.service; then
  install -D -m 644 "$SRC/systemd/grid-agent-supply.service" /etc/systemd/system/grid-agent.service   # -D: a container may have no systemd
  CHANGED=1
fi
[ -z "$SYSTEMD" ] || [ -z "$CHANGED" ] || systemctl daemon-reload

# --- Enrol, benchmark, start ---------------------------------------------------------------------------------------
if [ -n "$REENROL" ] && [ -f "$STATE" ]; then
  [ -z "$SYSTEMD" ] || systemctl stop grid-agent 2>/dev/null || true
  mv "$STATE" "$STATE.replaced-$(date +%s)"
  echo "the machine's old enrolment stays on the Machines page until it is removed there"
fi
if [ -f "$STATE" ]; then
  echo "== already enrolled as $("$VENV/bin/python" -c 'import json, sys; print(json.load(open(sys.argv[1]))["node_id"])' "$STATE") (--reenrol to enrol again)"
else
  echo "== enrol"
  config "$TOKEN" | "$VENV/bin/grid-agent" --config /dev/stdin enrol > "$WORK/enrolled.json" \
    || die "the control plane did not enrol this machine: the line above says why"
  "$VENV/bin/python" -c 'import json, sys; e = json.load(open(sys.argv[1])); print("enrolled as", e["node_id"], "at", e.get("rate_usd_per_hour"), "USD an hour")' "$WORK/enrolled.json"
  CHANGED=1
fi
if [ -z "$NO_BENCH" ] && { [ -n "$CHANGED" ] || [ ! -s "$BENCHED" ]; }; then
  echo "== benchmarks (each template image is pulled once: this can take a while)"
  [ -z "$SYSTEMD" ] || systemctl stop grid-agent 2>/dev/null || true
  "$VENV/bin/grid-agent" bench > "$WORK/bench.json" || true      # templates it could not run are said above; the rest count
  "$VENV/bin/python" -c 'import json, sys; r = json.load(open(sys.argv[1])); sys.exit(0 if r else 1)' "$WORK/bench.json" \
    || die "no template could be benchmarked on this machine: it would be offered no work (the lines above say why)"
  install -m 600 "$WORK/bench.json" "$BENCHED"
  echo "benchmarked: $(cat "$BENCHED")"
  CHANGED=1
elif [ -z "$NO_BENCH" ]; then
  echo "== benchmarks as the last install reported them: $(cat "$BENCHED") (grid-agent bench runs them again)"
fi
if [ -n "$NO_START" ] || [ -z "$SYSTEMD" ]; then
  echo "installed; not started (--no-start): run /opt/grid/agent/.venv/bin/grid-agent run, or systemctl enable --now grid-agent"
elif [ -z "$CHANGED" ] && systemctl is-active --quiet grid-agent; then
  echo "grid-agent is running, and nothing it runs on changed: left as it was"
else
  systemctl enable grid-agent >/dev/null 2>&1
  systemctl restart grid-agent              # a running agent hands its lease back and starts again on the code just installed
  sleep 15
  systemctl is-active --quiet grid-agent && [ "$(systemctl show -p NRestarts --value grid-agent)" = 0 ] \
    || die "the agent did not stay up: journalctl -u grid-agent -n 50"
  echo "grid-agent is running"
fi
[ -n "$CHANGED" ] || echo "nothing needed changing: this machine is installed as this run would install it"
echo "done: $NAME, $MODE mode on $PLATFORM, checkpointing $([ "$CAPABLE" = true ] && echo on || echo 'off (units restart)'); manage it on the dashboard's Machines page or with unitfabric machines"
}

main "$@"