#!/bin/sh
# fillrate-task: nvidia_driver_update
# Apt-only NVIDIA driver branch update. fillrate-runner verifies the signed
# task and the one-time ticket, then starts this script. It refuses to run
# unless that runner set FILLRATE_TASK. Nothing here starts by itself.
#
# Parameters from the ticket (the runner exports them; they are not read from
# the host environment the runner cleared):
#   FILLRATE_BRANCH          approved major branch, for example 595
#   FILLRATE_RESOLVED_BUILD  exact build, or empty until apt-cache resolves it
#   FILLRATE_MODE            apply | dry-run | rollback | repair
#   FILLRATE_RESTORE_BRANCH  branch to put back (optional; repair and apply recovery)
#   FILLRATE_REPAIR_STRATEGY forward (default) | restore
#   FILLRATE_NOT_BEFORE      window start (unix seconds)
#   FILLRATE_EXPIRES         window end (unix seconds)
#   FILLRATE_REPAIR_AUTH_ID  repair-until-healthy grant, when this is a follow-on
#   FILLRATE_ATTEMPT         attempt number under that grant
#   FILLRATE_PARENT_TICKET   previous ticket id, when this attempt is paired
#   FILLRATE_VAST_WINDOW_ID  Vast maintenance window id that must stay open
#
# The installed flavor comes from dpkg, matched to the running driver branch:
#   open      nvidia-dkms-<N>-open, nvidia-driver-<N>-open, or
#             nvidia-kernel-source-<N>-open
#   standard  nvidia-dkms-<N> or nvidia-driver-<N>
# The target meta-package follows that flavor (nvidia-driver-<branch>-open
# or nvidia-driver-<branch>). apt-cache policy supplies the build when the
# ticket does not. A runfile install (nvidia-uninstall on PATH or in
# /usr/bin, or no flavor package for the running branch) aborts non-apt-install.
#
# dry-run records pre-flight state, runs an ordered list of read-only
# apt-get -s strategies, then checks unattended-upgrades, DKMS, Secure Boot,
# the Docker NVIDIA runtime, the Vast daemon, and free space. It prints the
# JSON report on stdout. It does not unhold, install, stop vastai, or reboot.
# rollback reinstalls the previous branch of the same flavor, recorded by an
# earlier apply in the state directory. The ticket branch stays the approved
# request branch. rollback and repair both tolerate a broken nvidia-smi.
# GPU presence then comes from lspci (NVIDIA VGA or 3D) and /proc/driver/nvidia.
# A non-empty dpkg --audit is copied into the report and is not an abort.
# The 'nothing running on the GPUs' check is skipped when nvidia-smi fails.
# Vast job containers (C.<id> or a vastai/ image) still abort rentals-present.
# repair is for a node left half-changed after a failed apply.
# Diagnostics in the JSON report on stdout (the stream fillrate-runner shows
# the fillrate user; term.log itself stays root-only) are the last 200 lines
# of apt term.log and history.log, DKMS make.log and dkms logs, dpkg -l
# nvidia, dpkg --audit, dkms status, and the crash file name when one is
# present. Every report also carries failed_step and the failing command's
# exit code when a step failed, plus apt-mark showhold, apt-cache policy for
# the target meta, dpkg -l for nvidia and libnvidia packages, and the last
# 120 lines of each $STATE/apt-*.txt, dpkg-configure.txt, and
# dkms-autoinstall.txt that exists. The ticket or
# FILLRATE_REPAIR_STRATEGY selects the path (forward or restore, default
# forward). forward removes auto-installed packages that own files the new
# debs also ship (libnvidia-egl-xcb1 and libnvidia-egl-xlib1 on a headless
# host), runs apt-get install -f, installs libnvidia-gl-<branch> and the
# target meta at the resolved build, runs dpkg --configure -a, and reboots
# only when DKMS shows that module installed for the running kernel and
# nvidia.ko is on disk under /lib/modules/$(uname -r)/updates/dkms (or
# updates/ or extra/). If that module is not on disk, forward stops and
# restore runs. On that success path it enables nvidia-persistenced
# from nvidia-compute-utils-<branch>, re-holds the new driver packages plus
# the baseline non-NVIDIA holds, installs the post-reboot unit, and reboots.
# The owner authorized that reboot; the maintenance window is in the report.
# The existing post phase then checks nvidia-smi, persistence, the container
# GPU, and vastai. If forward fails, or the strategy is restore, it puts the
# previous branch back from the network and does not reboot: FILLRATE_RESTORE_BRANCH,
# then saved previous_branch, then the ticket branch when that is not the
# half-installed branch, then the removed (rc) driver branch, then 580.
# Before any apt install or remove, repair and apply's restore unhold
# NVIDIA-family holds (nvidia_holds, or apt-mark showhold filtered to nvidia*
# when that fact is empty). nvidia-docker2 and nvidia-container* stay held.
# A package apt-mark hold turned from rc into hc is still a driver package.
# Holds are put back on success and from an EXIT trap on every failure path.
# apt-get update runs first. The restore install is
# apt-get install -y --no-install-recommends of the previous meta (Depends
# pull the userspace set; nvidia-prime is only a Recommends). libnvidia-egl-gbm1
# is not required on a headless host. libnvidia-egl-xcb1 and libnvidia-egl-xlib1
# are removed only when a downloaded deb ships the same path: apt-get -s does
# not report that overlap, but libnvidia-gl-580 does ship those json files and
# libnvidia-egl-xcb.so.1.0.5. DKMS autoinstall then runs for the running kernel,
# nvidia-persistenced is restarted, and vastai is reset-failed and started.
# It still purges unconfigured packages from the failed branch.
#
# Strategies, first success wins (apt-get -s exit 0, and the plan does not
# install nvidia-prime). That winner is the validated plan. There is no
# second simulation that tries to keep nvidia-persistenced or
# libnvidia-egl-gbm1 on the install side. Dry-run and apply both select
# against a copy of the dpkg status in which NVIDIA holds are marked
# install, so the solver sees the same problem apply will run after
# apt-mark unhold. Dry-run does not unhold anything. Other holds stay.
#   A  apt-get -s -o APT::Get::Allow-Change-Held-Packages=true
#      -o Dir::State::status=<unheld copy> install <target>
#   B  A with --no-install-recommends
#   C  A plus explicit removal of the previous branch meta and that branch's
#      installed driver packages (<old-meta>- and libnvidia-*/nvidia-* -OLD)
#   D  C with --no-install-recommends
# The report names the winner (sim_strategy) and lists every attempt's
# command, exit code, and captured output. Output is complete up to 200
# lines; longer output keeps the head and the tail. apt_install, apt_remove,
# and apt_upgrade are the winning strategy's sets only, de-duplicated.
# Losing strategies are not concatenated into those sets. If none succeed,
# abort apt-sim-failed with all four outputs. A solved plan that would
# install nvidia-prime is not a success; the next strategy is tried. If the
# winning plan installs nvidia-prime, or every solved plan does, abort
# apt-would-install-nvidia-prime.
#
# Removing the standalone nvidia-persistenced package is expected when the
# winning plan installs nvidia-compute-utils-<new branch>. That package
# replaces it (it ships /usr/bin/nvidia-persistenced and the systemd unit;
# apt-cache show and, when the package is installed, dpkg -L are checked).
# Removing libnvidia-egl-gbm1 is expected when the plan installs
# libnvidia-gl-<new branch> (the library is irrelevant on a headless host).
# Those pairs are recorded in replaced_by_new_branch. If the plan removes
# nvidia-persistenced and nvidia-compute-utils-<new> is not in the install
# or upgrade set, abort persistenced-not-replaced.
#
# Abort before package changes when Secure Boot is on and nvidia.ko has no
# signer, the driver is not an apt package, a GPU process or container is
# running, / or /boot is short on space, or the winning plan would remove
# nvidia-cuda-toolkit, the cuda metapackage, the cuda-drivers metapackage,
# a libcudnn* package, or a cuda-toolkit* package. Also abort when the plan
# would install or upgrade cuda-toolkit* (the host CUDA toolkit, for example
# cuda-toolkit-13-2, is not changed). Abort when the plan would upgrade a
# package outside the driver family, or remove docker, containerd,
# nvidia-container-toolkit, nvidia-docker2, libvirt, vast*, kernel or
# linux-image packages, cuda-* packages, or anything else that is not an
# old-branch driver package (nvidia-*, libnvidia-*, linux-modules-nvidia-*,
# xserver-xorg-video-nvidia* for the installed branch) or one of the
# replaced pairs (nvidia-persistenced, libnvidia-egl-gbm1), plus packages
# the pre-removal file-conflict check added. Those extra removals are
# accepted only from that check: the owner is auto-installed and is not on
# the never-remove list (cuda-*, nvidia-container*, nvidia-docker*, docker*,
# containerd*, vast*, kernel / linux-image). The winning transaction must
# install the target meta-package at the resolved build. Apply unholds and
# later reholds only NVIDIA packages on the success path. Holds on docker,
# containerd, libvirt, and xfsprogs are reported and left alone until a
# failed install is recovered. Before vastai is stopped or any hold is
# released, apply installs linux-headers for the running kernel when they
# are missing and that install would not remove a protected package, checks
# that nvidia-kernel-source for the new branch is available, downloads the
# solved plan, and compares each downloaded deb's files with dpkg -S. A file
# owned by an installed package that this plan does not already remove, and
# that the new package does not Replace, is a conflict. Auto-installed
# owners outside the never-remove list are added to the removal set and the
# safety checks run again with them included. Any other conflict aborts
# file-conflict with the conflict list in the report and changes nothing.
# A headers, download, or file-conflict failure aborts with nothing removed.
# After unhold, confirm_plan selects again and applies the same safety
# checks. It does not require the same strategy or the same install, remove,
# or upgrade sets. The file-conflict removals are attached again after that
# re-simulation. Apply installs the post-unhold plan that passed those
# checks (same packages, same --no-install-recommends choice, no
# held-package override, no extra keep packages). The report field
# plan_differs_from_pre_unhold is yes when that plan is not the pre-unhold
# one. A violation aborts with that check's abort_reason, re-holds, and
# restarts vastai. If the real install fails after that, apply restores the
# previous branch the same way repair's restore path does, including apt
# and DKMS diagnostics, and does not reboot. nvidia-prime is never
# installed. Listing and prices are not changed.
# Abort after the swap when the module does not match the running kernel.
# Abort after reboot when a GPU is missing or the network is still down at
# the ticket deadline. After reboot the post-check confirms nvidia-smi shows
# the new branch, requires the nvidia-persistenced systemd unit, runs
# systemctl enable --now nvidia-persistenced, and turns persistence mode on
# (nvidia-smi -pm 1 if the mode is still off). Then it runs the container GPU
# check. Persistence failure uses the same post-failure path as a missing GPU
# (rehold, start vastai, non-zero exit). The report lists the service state
# before and after, persistence_detail, and the winning install and remove sets.
#
# JSON on stdout is the report. It does not include keys, tokens, or the nonce.
# apt-sim.txt stays in the state directory for root. That directory is mode
# 700. The same JSON is copied to /var/lib/fillrate-ops (mode 755, file mode
# 644) so the fillrate user can read abort_reason, failed_step, and the apt
# transcript without opening the root-only state directory. Do not loosen
# the state directory. Every abort report includes those fields. When apt
# did not run, apt_transcript says so.
# A failed simulation is recorded and the dry run still runs the checks below
# (unattended-upgrades, DKMS, Secure Boot, Docker, Vast, free space), then
# exits non-zero with that abort_reason. The report includes the failing
# command and the captured apt text.

set -u

if [ "${FILLRATE_TASK:-}" != "nvidia_driver_update" ]; then
  printf '%s\n' "refusing: fillrate-runner did not launch this task" >&2
  exit 1
fi

BRANCH=${FILLRATE_BRANCH:-}
BUILD=${FILLRATE_RESOLVED_BUILD:-}
MODE=${FILLRATE_MODE:-apply}
EXPIRES=${FILLRATE_EXPIRES:-0}
PHASE=${FILLRATE_DRIVER_PHASE:-pre}
STATE=${FILLRATE_DRIVER_STATE:-/var/lib/fillrate-runner/nvidia-driver-update}
VAR_LOG=${FILLRATE_VAR_LOG:-/var/log}
VAR_LIB=${FILLRATE_VAR_LIB:-/var/lib}
VAR_CRASH=${FILLRATE_VAR_CRASH:-/var/crash}
# Used when a failed apply left no saved hold list. Machine 152178's set.
BASELINE_HOLDS="containerd.io docker-ce docker-ce-cli docker-ce-rootless-extras libvirt-dev nvidia-docker2 nvidia-driver-580-open xfsprogs"
ORIG_HOLDS=""
ORIG_NVIDIA_HOLDS=""
MIN_ROOT_KB=2097152
MIN_BOOT_KB=65536
PKG_SPEC=
PKG=
SIM_FAILED=0
SIM_REASON=
SIM_STRATEGY=
PLAN_PKGS=
PLAN_NO_RECOMMENDS=0
FELL_BACK=0
PRIME_HIT=0
CONFIRM_PLAN=0
STRAT_SEQ=0
UPDATE_FAILED=0
RUNTIME_ABORT=
HOLDS_RELEASED=0
HOLDS_RESTORED=0
ALLOW_EGL_CONFLICT_REMOVAL=0
SIM_APT_OPTS="-o APT::Get::Allow-Change-Held-Packages=true"

case "$MODE" in
  apply|dry-run|rollback|repair) ;;
  *)
    printf '%s\n' "refusing: bad mode" >&2
    exit 1
    ;;
esac

if [ "$MODE" = "repair" ]; then
  case "${FILLRATE_REPAIR_STRATEGY:-forward}" in
    forward|restore) ;;
    *)
      printf '%s\n' "refusing: bad repair_strategy" >&2
      exit 1
      ;;
  esac
fi

printf '%s\n' "$BRANCH" | grep -E -q '^[0-9]{3,4}$' || {
  printf '%s\n' "refusing: branch is required" >&2
  exit 1
}

mkdir -p "$STATE" || {
  printf '%s\n' "refusing: cannot write state" >&2
  exit 1
}
# Root-only. The report the caller can read is the JSON on stdout and the
# mode 644 copy under FILLRATE_OPS_REPORT_DIR.
chmod 700 "$STATE" 2>/dev/null || true

json_escape() {
  printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g' | tr '\n\r\t' '   '
}

json_escape_multiline() {
  printf '%s' "$1" | awk '
    BEGIN { ORS = "" }
    {
      gsub(/\r/, "")
      gsub(/\\/, "\\\\")
      gsub(/"/, "\\\"")
      gsub(/\t/, "\\t")
      if (NR > 1) printf "\\n"
      printf "%s", $0
    }
  '
}

file_tail() {
  path=$1
  if [ -f "$path" ]; then
    tail -n 60 "$path"
  fi
}

# Full capture through 200 lines. Longer apt output keeps the head and the tail
# so neither the first unmet-dependency lines nor the final error is dropped.
cap_text() {
  path=$1
  if [ ! -f "$path" ]; then
    return 0
  fi
  n=$(wc -l < "$path" | awk '{ print $1 }')
  if [ "$n" -le 200 ]; then
    cat "$path"
    return 0
  fi
  head -n 100 "$path"
  printf '%s\n' "... $((n - 200)) lines omitted ..."
  tail -n 100 "$path"
}

emit_strategies() {
  printf '    "strategies": ['
  if [ -f "$STATE/strategies/order" ]; then
    first=1
    while IFS= read -r n; do
      [ -n "$n" ] || continue
      id=$(cat "$STATE/strategies/$n.id" 2>/dev/null || true)
      rc=$(cat "$STATE/strategies/$n.rc" 2>/dev/null || true)
      cmd=$(cat "$STATE/strategies/$n.cmd" 2>/dev/null || true)
      text=$(cap_text "$STATE/strategies/$n.out")
      case "$rc" in
        ''|*[!0-9]*) rc=1 ;;
      esac
      if [ "$first" = "0" ]; then
        printf ','
      fi
      first=0
      printf '\n      {"id": "%s", "command": "%s", "exit_code": %s, "output": "%s"}' \
        "$(json_escape "$id")" "$(json_escape "$cmd")" "$rc" "$(json_escape_multiline "$text")"
    done < "$STATE/strategies/order"
  fi
  printf '\n    ]\n'
}

write_fact() {
  key=$1
  val=$2
  tmp="$STATE/facts.new"
  if [ -f "$STATE/facts" ]; then
    grep -v "^${key}=" "$STATE/facts" > "$tmp" || true
  else
    : > "$tmp"
  fi
  printf '%s=%s\n' "$key" "$val" >> "$tmp"
  mv "$tmp" "$STATE/facts"
}

read_fact() {
  key=$1
  if [ ! -f "$STATE/facts" ]; then
    return 1
  fi
  awk -F= -v k="$key" 'index($0, k "=") == 1 { sub(/^[^=]*=/, ""); print; found=1; exit } END { exit !found }' "$STATE/facts"
}

write_fact repair_auth_id "${FILLRATE_REPAIR_AUTH_ID:-}"
write_fact attempt "${FILLRATE_ATTEMPT:-}"
write_fact parent_ticket "${FILLRATE_PARENT_TICKET:-}"
write_fact vast_window_id "${FILLRATE_VAST_WINDOW_ID:-}"
write_fact maint_window "${FILLRATE_NOT_BEFORE:-0}-${EXPIRES}"

emit_report() {
  capture_diagnostics
  current=$(apt-mark showhold 2>/dev/null | tr '\n' ' ' || true)
  write_fact holds_after "$current"
  result=$(read_fact result || true)
  reason=$(read_fact abort_reason || true)
  failed_step=$(read_fact failed_step || true)
  failed_rc=$(read_fact failed_exit_code || true)
  forward_reason=$(read_fact forward_abort_reason || true)
  forward_step=$(read_fact forward_failed_step || true)
  forward_rc=$(read_fact forward_exit_code || true)
  resolved=$(read_fact resolved_build || true)
  prev_b=$(read_fact previous_branch || true)
  prev_v=$(read_fact previous_build || true)
  gpu_before=$(read_fact gpu_count || true)
  gpu_after=$(read_fact gpu_count_after || true)
  driver_before=$(read_fact driver_before || true)
  driver_after=$(read_fact driver_after || true)
  cuda_before=$(read_fact cuda_before || true)
  holds=$(read_fact holds || true)
  holds_after=$(read_fact holds_after || true)
  space_root=$(read_fact space_root_kb || true)
  space_boot=$(read_fact space_boot_kb || true)
  secure=$(read_fact secure_boot || true)
  method=$(read_fact install_method || true)
  flavor=$(read_fact driver_flavor || true)
  target_pkg=$(read_fact target_package || true)
  apt_install=$(read_fact apt_install || true)
  apt_remove=$(read_fact apt_remove || true)
  apt_upgrade=$(read_fact apt_upgrade || true)
  sim_strategy=$(read_fact sim_strategy || true)
  pre_unhold_strategy=$(read_fact pre_unhold_strategy || true)
  plan_differs=$(read_fact plan_differs_from_pre_unhold || true)
  cuda_impact=$(read_fact cuda_impact || true)
  dkms=$(read_fact dkms || true)
  xid_before=$(read_fact xid_count || true)
  xid_after=$(read_fact xid_count_after || true)
  ecc_before=$(read_fact ecc_uncorrected || true)
  ecc_after=$(read_fact ecc_uncorrected_after || true)
  persistence=$(read_fact persistence || true)
  persistenced=$(read_fact persistenced || true)
  persistence_after=$(read_fact persistence_after || true)
  persistenced_after=$(read_fact persistenced_after || true)
  replaced=$(read_fact replaced_by_new_branch || true)
  provider_files=$(read_fact persistenced_provider_files || true)
  persistence_detail=$(read_fact persistence_detail || true)
  fabric=$(read_fact fabric_manager || true)
  docker_rt=$(read_fact docker_runtime || true)
  vast=$(read_fact vast_daemon || true)
  unattended=$(read_fact unattended_upgrades || true)
  packages=$(read_fact packages || true)
  mig=$(read_fact mig || true)
  temps=$(read_fact temps || true)
  apt_command=$(read_fact apt_command || true)
  apt_update=$(read_fact apt_update || true)
  apt_sim=$(read_fact apt_sim || true)
  held_meta=$(read_fact held_meta || true)
  held_conflicts=$(read_fact held_meta_conflicts || true)
  version_fallback=$(read_fact version_fallback || true)
  recovery=$(read_fact recovery || true)
  restore_branch=$(read_fact restore_branch || true)
  reboot_required=$(read_fact reboot_required || true)
  repair_strategy=$(read_fact repair_strategy || true)
  file_conflicts=$(read_fact file_conflicts || true)
  maint_window=$(read_fact maint_window || true)
  repair_auth_id=$(read_fact repair_auth_id || true)
  repair_attempt=$(read_fact attempt || true)
  parent_ticket=$(read_fact parent_ticket || true)
  vast_window_id=$(read_fact vast_window_id || true)
  module_ko=$(read_fact module_ko || true)
  diag_term=$(cat "$STATE/diag-term.log" 2>/dev/null || true)
  diag_history=$(cat "$STATE/diag-history.log" 2>/dev/null || true)
  diag_make=$(cat "$STATE/diag-dkms-make.log" 2>/dev/null || true)
  diag_dkms=$(cat "$STATE/diag-dkms.log" 2>/dev/null || true)
  diag_dpkg=$(cat "$STATE/diag-dpkg-nvidia.txt" 2>/dev/null || true)
  diag_audit=$(cat "$STATE/diag-dpkg-audit.txt" 2>/dev/null || true)
  diag_status=$(cat "$STATE/diag-dkms-status.txt" 2>/dev/null || true)
  diag_crash=$(cat "$STATE/diag-crash.txt" 2>/dev/null || true)
  diag_apt_logs=$(cat "$STATE/diag-apt-state.txt" 2>/dev/null || true)
  apt_transcript=$diag_apt_logs
  if [ -z "$(printf '%s' "$apt_transcript" | tr -d '[:space:]')" ]; then
    apt_transcript=$diag_term
  fi
  if [ -z "$(printf '%s' "$apt_transcript" | tr -d '[:space:]')" ]; then
    apt_transcript="(no apt transcript yet)"
  fi
  diag_showhold=$(cat "$STATE/diag-showhold.txt" 2>/dev/null || true)
  diag_dpkg_live=$(cat "$STATE/diag-dpkg-live.txt" 2>/dev/null || true)
  diag_policy_target=$(cat "$STATE/diag-policy-target.txt" 2>/dev/null || true)
  policy_text=$(file_tail "$STATE/apt-policy.txt")
  holds_text=$(file_tail "$STATE/apt-holds.txt")
  sim_text=$(cap_text "$STATE/apt-sim.txt")
  update_text=$(file_tail "$STATE/apt-update.txt")
  update_err=
  if [ "$apt_update" = "failed" ]; then
    update_err=$update_text
  fi
  error_text=
  if [ -n "$reason" ]; then
    if [ "$reason" = "apt-update-failed" ] && [ -n "$update_text" ]; then
      error_text=$update_text
    elif [ -n "$sim_text" ]; then
      error_text=$sim_text
    else
      error_text=$update_err
    fi
  fi
  report="$STATE/report.json"
  {
    printf '%s\n' '{'
    printf '  "task": "nvidia_driver_update",\n'
    printf '  "mode": "%s",\n' "$(json_escape "$MODE")"
    printf '  "phase": "%s",\n' "$(json_escape "$PHASE")"
    printf '  "machine_id": "%s",\n' "$(json_escape "${FILLRATE_MACHINE_ID:-}")"
    printf '  "branch": "%s",\n' "$(json_escape "$BRANCH")"
    printf '  "flavor": "%s",\n' "$(json_escape "$flavor")"
    printf '  "target_package": "%s",\n' "$(json_escape "$target_pkg")"
    printf '  "resolved_build": "%s",\n' "$(json_escape "$resolved")"
    printf '  "result": "%s",\n' "$(json_escape "$result")"
    printf '  "abort_reason": "%s",\n' "$(json_escape "$reason")"
    printf '  "failed_step": "%s",\n' "$(json_escape "$failed_step")"
    printf '  "failed_exit_code": "%s",\n' "$(json_escape "$failed_rc")"
    printf '  "apt_transcript": "%s",\n' "$(json_escape_multiline "$apt_transcript")"
    printf '  "repair_auth_id": "%s",\n' "$(json_escape "$repair_auth_id")"
    printf '  "attempt": "%s",\n' "$(json_escape "$repair_attempt")"
    printf '  "parent_ticket": "%s",\n' "$(json_escape "$parent_ticket")"
    printf '  "vast_window_id": "%s",\n' "$(json_escape "$vast_window_id")"
    printf '  "forward_abort_reason": "%s",\n' "$(json_escape "$forward_reason")"
    printf '  "forward_failed_step": "%s",\n' "$(json_escape "$forward_step")"
    printf '  "forward_exit_code": "%s",\n' "$(json_escape "$forward_rc")"
    printf '  "recovery": "%s",\n' "$(json_escape "$recovery")"
    printf '  "restore_branch": "%s",\n' "$(json_escape "$restore_branch")"
    printf '  "reboot_required": "%s",\n' "$(json_escape "$reboot_required")"
    printf '  "repair_strategy": "%s",\n' "$(json_escape "$repair_strategy")"
    printf '  "file_conflicts": "%s",\n' "$(json_escape "$file_conflicts")"
    printf '  "maint_window": "%s",\n' "$(json_escape "$maint_window")"
    printf '  "module_ko": "%s",\n' "$(json_escape "$module_ko")"
    printf '  "sim_strategy": "%s",\n' "$(json_escape "$sim_strategy")"
    printf '  "pre_unhold_strategy": "%s",\n' "$(json_escape "$pre_unhold_strategy")"
    printf '  "plan_differs_from_pre_unhold": "%s",\n' "$(json_escape "$plan_differs")"
    printf '  "apt_install": "%s",\n' "$(json_escape "$apt_install")"
    printf '  "apt_remove": "%s",\n' "$(json_escape "$apt_remove")"
    printf '  "replaced_by_new_branch": "%s",\n' "$(json_escape "$replaced")"
    printf '  "persistenced_provider_files": "%s",\n' "$(json_escape "$provider_files")"
    printf '  "persistence_detail": "%s",\n' "$(json_escape "$persistence_detail")"
    printf '  "apt_upgrade": "%s",\n' "$(json_escape "$apt_upgrade")"
    printf '  "cuda_impact": "%s",\n' "$(json_escape "$cuda_impact")"
    printf '  "before": {\n'
    printf '    "driver": "%s",\n' "$(json_escape "$driver_before")"
    printf '    "cuda": "%s",\n' "$(json_escape "$cuda_before")"
    printf '    "gpu_count": "%s",\n' "$(json_escape "$gpu_before")"
    printf '    "previous_branch": "%s",\n' "$(json_escape "$prev_b")"
    printf '    "previous_build": "%s",\n' "$(json_escape "$prev_v")"
    printf '    "holds": "%s",\n' "$(json_escape "$holds")"
    printf '    "packages": "%s",\n' "$(json_escape "$packages")"
    printf '    "space_root_kb": "%s",\n' "$(json_escape "$space_root")"
    printf '    "space_boot_kb": "%s",\n' "$(json_escape "$space_boot")"
    printf '    "secure_boot": "%s",\n' "$(json_escape "$secure")"
    printf '    "install_method": "%s",\n' "$(json_escape "$method")"
    printf '    "dkms": "%s",\n' "$(json_escape "$dkms")"
    printf '    "xid_count": "%s",\n' "$(json_escape "$xid_before")"
    printf '    "ecc_uncorrected": "%s",\n' "$(json_escape "$ecc_before")"
    printf '    "persistence": "%s",\n' "$(json_escape "$persistence")"
    printf '    "persistenced": "%s",\n' "$(json_escape "$persistenced")"
    printf '    "fabric_manager": "%s",\n' "$(json_escape "$fabric")"
    printf '    "docker_runtime": "%s",\n' "$(json_escape "$docker_rt")"
    printf '    "vast_daemon": "%s",\n' "$(json_escape "$vast")"
    printf '    "unattended_upgrades": "%s",\n' "$(json_escape "$unattended")"
    printf '    "mig": "%s",\n' "$(json_escape "$mig")"
    printf '    "health": "%s"\n' "$(json_escape "$temps")"
    printf '  },\n'
    printf '  "after": {\n'
    printf '    "driver": "%s",\n' "$(json_escape "$driver_after")"
    printf '    "gpu_count": "%s",\n' "$(json_escape "$gpu_after")"
    printf '    "xid_count": "%s",\n' "$(json_escape "$xid_after")"
    printf '    "ecc_uncorrected": "%s",\n' "$(json_escape "$ecc_after")"
    printf '    "persistence": "%s",\n' "$(json_escape "$persistence_after")"
    printf '    "persistenced": "%s",\n' "$(json_escape "$persistenced_after")"
    printf '    "holds": "%s"\n' "$(json_escape "$holds_after")"
    printf '  },\n'
    printf '  "apt": {\n'
    printf '    "command": "%s",\n' "$(json_escape "$apt_command")"
    printf '    "update": "%s",\n' "$(json_escape "$apt_update")"
    printf '    "update_error": "%s",\n' "$(json_escape_multiline "$update_err")"
    printf '    "sim": "%s",\n' "$(json_escape "$apt_sim")"
    printf '    "error": "%s",\n' "$(json_escape_multiline "$error_text")"
    printf '    "sim_output": "%s",\n' "$(json_escape_multiline "$sim_text")"
    printf '    "policy": "%s",\n' "$(json_escape_multiline "$policy_text")"
    printf '    "holds": "%s",\n' "$(json_escape_multiline "$holds_text")"
    printf '    "held_meta": "%s",\n' "$(json_escape "$held_meta")"
    printf '    "held_meta_conflicts": "%s",\n' "$(json_escape "$held_conflicts")"
    printf '    "version_fallback": "%s",\n' "$(json_escape "$version_fallback")"
    emit_strategies
    printf '  },\n'
    printf '  "diagnostics": {\n'
    printf '    "term_log": "%s",\n' "$(json_escape_multiline "$diag_term")"
    printf '    "history_log": "%s",\n' "$(json_escape_multiline "$diag_history")"
    printf '    "dkms_make_log": "%s",\n' "$(json_escape_multiline "$diag_make")"
    printf '    "dkms_log": "%s",\n' "$(json_escape_multiline "$diag_dkms")"
    printf '    "dpkg_nvidia": "%s",\n' "$(json_escape_multiline "$diag_dpkg")"
    printf '    "dpkg_audit": "%s",\n' "$(json_escape_multiline "$diag_audit")"
    printf '    "dkms_status": "%s",\n' "$(json_escape_multiline "$diag_status")"
    printf '    "crash_file": "%s",\n' "$(json_escape_multiline "$diag_crash")"
    printf '    "failed_step": "%s",\n' "$(json_escape "$failed_step")"
    printf '    "failed_exit_code": "%s",\n' "$(json_escape "$failed_rc")"
    printf '    "apt_showhold": "%s",\n' "$(json_escape_multiline "$diag_showhold")"
    printf '    "dpkg_nvidia_packages": "%s",\n' "$(json_escape_multiline "$diag_dpkg_live")"
    printf '    "apt_policy": "%s",\n' "$(json_escape_multiline "$diag_policy_target")"
    printf '    "apt_state_logs": "%s"\n' "$(json_escape_multiline "$diag_apt_logs")"
    printf '  }\n'
    printf '%s\n' '}'
  } | tee "$report"
  publish_ops_report
}

# Copy the JSON report where the fillrate user can read it. The state
# directory stays mode 700. A copy failure must not hide the stdout report.
publish_ops_report() {
  src="$STATE/report.json"
  [ -f "$src" ] || return 0
  dest_dir=${FILLRATE_OPS_REPORT_DIR:-/var/lib/fillrate-ops}
  mkdir -p "$dest_dir" 2>/dev/null || return 0
  chmod 755 "$dest_dir" 2>/dev/null || true
  dest="$dest_dir/nvidia-driver-update-report.json"
  cp "$src" "$dest" 2>/dev/null || return 0
  chmod 644 "$dest" 2>/dev/null || true
  attempt_n=$(read_fact attempt || true)
  case "$attempt_n" in
    ''|*[!0-9]*) ;;
    *)
      numbered="$dest_dir/nvidia-driver-update-attempt-$attempt_n.json"
      cp "$src" "$numbered" 2>/dev/null || return 0
      chmod 644 "$numbered" 2>/dev/null || true
      ;;
  esac
  return 0
}

abort() {
  reason=$1
  step=${2:-}
  write_fact result abort
  write_fact abort_reason "$reason"
  if [ -n "$step" ]; then
    write_fact failed_step "$step"
  else
    existing=$(read_fact failed_step || true)
    if [ -z "$existing" ]; then
      write_fact failed_step "$reason"
    fi
  fi
  capture_apt_context || true
  emit_report
  exit 2
}

# Record which apt step failed. The report prints this even when a later
# restore overwrites abort_reason; callers copy it to forward_* first.
note_step_fail() {
  reason=$1
  step=$2
  rc=$3
  write_fact abort_reason "$reason"
  write_fact failed_step "$step"
  write_fact failed_exit_code "$rc"
}

disk_kb() {
  mount=$1
  df -Pk "$mount" 2>/dev/null | awk 'NR==2 { print $4 }'
}

count_xid() {
  if command -v journalctl >/dev/null 2>&1; then
    journalctl -k -b --no-pager 2>/dev/null | grep -c 'Xid' || true
  else
    printf '%s\n' "0"
  fi
}

sum_ecc() {
  nvidia-smi --query-gpu=ecc.errors.uncorrected.volatile.total --format=csv,noheader 2>/dev/null \
    | awk -F, '{ gsub(/ /, "", $1); if ($1 ~ /^[0-9]+$/) s += $1 } END { print s+0 }'
}

package_list() {
  dpkg -l 2>/dev/null | awk '
    ($1 == "ii" || $1 == "hi") && ($2 ~ /^(nvidia-|libnvidia-|linux-modules-nvidia-|cuda-|cuda$|cuda-drivers$|libcudnn|nvidia-cuda-toolkit)/) {
      pkg = $2
      sub(/:.*/, "", pkg)
      printf "%s=%s ", pkg, $3
    }
  '
}

has_pkg() {
  name=$1
  dpkg -l 2>/dev/null | awk -v n="$name" '
    ($1 == "ii" || $1 == "hi") {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg == n) found = 1
    }
    END { exit !found }
  '
}

# open if a -open flavor package for this branch is installed; otherwise
# standard if nvidia-dkms-<N> or nvidia-driver-<N> is installed.
classify_flavor() {
  branch=$1
  dpkg -l 2>/dev/null | awk -v br="$branch" '
    ($1 == "ii" || $1 == "hi") {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg ~ ("^nvidia-(dkms|driver|kernel-source)-" br "-open$")) open = 1
      else if (pkg ~ ("^nvidia-(dkms|driver)-" br "$")) standard = 1
    }
    END {
      if (open) print "open"
      else if (standard) print "standard"
    }
  '
}

# Repair sees unpacked (iU/iF), removed (rc), and held config-files (hc).
# apt-mark hold on an rc package rewrites the desired state to hc.
classify_flavor_broad() {
  dpkg -l 2>/dev/null | awk '
    $1 ~ /^(ii|hi|iU|iF|iH|rc|hc)$/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg ~ /^nvidia-(dkms|driver|kernel-source)-[0-9]+-open$/) open = 1
      else if (pkg ~ /^nvidia-(dkms|driver)-[0-9]+$/) standard = 1
    }
    END {
      if (open) print "open"
      else if (standard) print "standard"
    }
  '
}

meta_for() {
  ver=$1
  if [ "$(read_fact driver_flavor || true)" = "open" ]; then
    printf '%s\n' "nvidia-driver-${ver}-open"
  else
    printf '%s\n' "nvidia-driver-${ver}"
  fi
}

nvidia_only_holds() {
  printf '%s\n' "$1" | tr ' ' '\n' | grep -E 'nvidia' | tr '\n' ' ' || true
}

protected_removal() {
  sim_file=$1
  awk '
    $1 == "Remv" {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg == "nvidia-cuda-toolkit" || pkg == "cuda" || pkg == "cuda-drivers" || pkg ~ /^libcudnn/ || pkg ~ /^cuda-toolkit($|-)/) print pkg
    }
  ' "$sim_file" | awk 'NR==1 { print; exit }'
}

cuda_toolkit_install() {
  sim_file=$1
  awk '
    $1 == "Inst" {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg ~ /^cuda-toolkit($|-)/) { print pkg; exit }
    }
  ' "$sim_file"
}

# One strategy's package list. Repeated names collapse; order stays.
dedupe_list() {
  printf '%s\n' "$1" | tr ' ' '\n' | awk 'NF && !seen[$0]++ { printf "%s ", $0 }'
}

list_has() {
  hay=$1
  needle=$2
  case " $hay " in
    *" $needle "*) return 0 ;;
  esac
  return 1
}

record_sim() {
  sim_file=$1
  # A fresh install is "Inst pkg (new ...)". An upgrade is "Inst pkg [old] (new ...)".
  # The "[amd64]" inside the new-version parentheses is not an upgrade.
  # sim_file is one strategy. Do not pass the concatenated attempt log.
  inst=$(awk '
    $1 == "Inst" && $0 !~ /\[[^]]+\][[:space:]]*\(/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      printf "%s ", pkg
    }
  ' "$sim_file")
  upg=$(awk '
    $1 == "Inst" && $0 ~ /\[[^]]+\][[:space:]]*\(/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      printf "%s ", pkg
    }
  ' "$sim_file")
  remv=$(awk '
    $1 == "Remv" {
      pkg = $2
      sub(/:.*/, "", pkg)
      printf "%s ", pkg
    }
  ' "$sim_file")
  impact=$(awk '
    $1 == "Inst" || $1 == "Remv" {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg ~ /cuda/ || pkg ~ /^libcudnn/) printf "%s:%s ", $1, pkg
    }
  ' "$sim_file")
  inst=$(dedupe_list "$inst")
  upg=$(dedupe_list "$upg")
  remv=$(dedupe_list "$remv")
  write_fact apt_install "$inst"
  write_fact apt_remove "$remv"
  write_fact apt_upgrade "$upg"
  if [ -z "$impact" ]; then
    impact="none"
  else
    impact=$(dedupe_list "$impact")
  fi
  write_fact cuda_impact "$impact"
}

rehold() {
  holds=$(read_fact nvidia_holds || true)
  if [ -z "$holds" ]; then
    return 0
  fi
  # shellcheck disable=SC2086
  apt-mark hold $holds >/dev/null 2>&1 || return 1
}

# Driver holds to release before repair/restore apt changes.
# nvidia-docker2 and the container stack stay held. When preflight recorded
# nothing (a held rc package can be hc, and an older fact can be empty),
# apt-mark showhold is the source. Live holds are always unioned in so a
# stale fact cannot hide nvidia-driver-580-open.
nvidia_hold_names() {
  fact=$(read_fact nvidia_holds || true)
  live=$(apt-mark showhold 2>/dev/null | grep -E 'nvidia' | tr '\n' ' ' || true)
  if [ -z "$fact" ]; then
    fact=$live
  else
    fact="$fact $live"
  fi
  printf '%s\n' "$fact" | tr ' ' '\n' | awk '
    NF && !seen[$0]++ {
      if ($0 ~ /^nvidia-docker/) next
      if ($0 ~ /^nvidia-container/) next
      if ($0 !~ /nvidia/) next
      print
    }
  ' | tr '\n' ' ' | awk '{$1=$1; print}'
}

release_nvidia_holds() {
  names=$(nvidia_hold_names)
  write_fact nvidia_unheld "$names"
  HOLDS_RELEASED=1
  HOLDS_RESTORED=0
  if [ -z "$names" ]; then
    return 0
  fi
  # shellcheck disable=SC2086
  apt-mark unhold $names >/dev/null 2>&1 || return 1
  return 0
}

exit_rehold() {
  [ "${HOLDS_RELEASED:-0}" = "1" ] || return 0
  [ "${HOLDS_RESTORED:-0}" = "1" ] && return 0
  HOLDS_RESTORED=1
  br=$(read_fact restore_branch || true)
  if [ -z "$br" ]; then
    br=$(choose_restore_branch || true)
  fi
  if [ -z "$br" ]; then
    br=580
  fi
  rehold_restore "$br" || true
}

start_vast() {
  if command -v systemctl >/dev/null 2>&1; then
    systemctl start vastai.service >/dev/null 2>&1 || return 1
  fi
  return 0
}

fail_after_change() {
  reason=$1
  rehold || true
  start_vast || true
  abort "$reason"
}

kernel_module_matches() {
  kernel=$(uname -r 2>/dev/null || true)
  [ -n "$kernel" ] || return 1
  status=$(dkms status 2>/dev/null || true)
  if printf '%s\n' "$status" | grep -F 'nvidia' | grep -F "$kernel" >/dev/null 2>&1; then
    return 0
  fi
  # A dkms status that names nvidia for some other kernel is a mismatch.
  # Fall back to the loaded module only when dkms has nothing to say.
  if printf '%s\n' "$status" | grep -F 'nvidia' >/dev/null 2>&1; then
    return 1
  fi
  vermagic=$(modinfo -F vermagic nvidia 2>/dev/null || true)
  printf '%s\n' "$vermagic" | grep -F "$kernel" >/dev/null 2>&1
}

install_post_unit() {
  cat > "$STATE/post.service" << EOF
[Unit]
Description=Fillrate NVIDIA driver update post-checks
ConditionPathExists=$STATE/pending-post
After=network-online.target

[Service]
Type=oneshot
Environment=FILLRATE_TASK=nvidia_driver_update
Environment=FILLRATE_DRIVER_PHASE=post
Environment=FILLRATE_BRANCH=$BRANCH
Environment=FILLRATE_RESOLVED_BUILD=$BUILD
Environment=FILLRATE_MODE=$MODE
Environment=FILLRATE_EXPIRES=$EXPIRES
Environment=FILLRATE_MACHINE_ID=${FILLRATE_MACHINE_ID:-}
Environment=FILLRATE_DRIVER_STATE=$STATE
ExecStart=/bin/sh $STATE/task.sh

[Install]
WantedBy=multi-user.target
EOF
  cp "$0" "$STATE/task.sh"
  chmod 700 "$STATE/task.sh"
  if [ "${FILLRATE_SKIP_BOOT_UNIT:-}" = "1" ]; then
    return 0
  fi
  if [ ! -d /etc/systemd/system ]; then
    return 1
  fi
  cp "$STATE/post.service" /etc/systemd/system/fillrate-nvidia-driver-update-post.service
  systemctl daemon-reload
  systemctl enable fillrate-nvidia-driver-update-post.service
}

remove_post_unit() {
  rm -f "$STATE/pending-post"
  if [ "${FILLRATE_SKIP_BOOT_UNIT:-}" = "1" ]; then
    return 0
  fi
  if command -v systemctl >/dev/null 2>&1; then
    systemctl disable --now fillrate-nvidia-driver-update-post.service >/dev/null 2>&1 || true
    rm -f /etc/systemd/system/fillrate-nvidia-driver-update-post.service
    systemctl daemon-reload >/dev/null 2>&1 || true
  fi
}

cuda_version() {
  nvidia-smi 2>/dev/null | awk '
    /CUDA Version/ {
      line = $0
      sub(/.*CUDA Version[[:space:]]*:[[:space:]]*/, "", line)
      sub(/[^0-9.].*/, "", line)
      print line
      exit
    }
  '
}

smi_version_ok() {
  printf '%s\n' "$1" | grep -E -q '^[0-9]+\.[0-9]+'
}

lspci_nvidia_count() {
  if ! command -v lspci >/dev/null 2>&1; then
    printf '%s\n' 0
    return 0
  fi
  lspci -nn 2>/dev/null | grep -E -i 'NVIDIA' | grep -E -i 'VGA compatible controller|3D controller' | wc -l | awk '{ print $1 }'
}

proc_nvidia_loaded() {
  proc=${FILLRATE_PROC_NVIDIA:-/proc/driver/nvidia}
  [ -d "$proc" ] || [ -e "$proc/version" ] || [ -d "$proc/gpus" ]
}

# Vast rentals are C.<contract> containers or vastai/ images. Other
# containers (exporters, sidecars) are not jobs. Exit 0 when a job is running.
vast_job_running() {
  if ! command -v docker >/dev/null 2>&1; then
    return 1
  fi
  listing=$(docker ps --format '{{.Names}} {{.Image}}' 2>/dev/null) || return 0
  printf '%s\n' "$listing" | awk '
    NF {
      name = $1
      image = $2
      if (name ~ /^C\.[0-9]+/) found = 1
      if (image ~ /(^|\/)vastai\//) found = 1
    }
    END { exit !found }
  '
}

refuse_vast_jobs() {
  if vast_job_running; then
    abort "rentals-present"
  fi
}

# nvidia.ko must exist before a reboot. The loaded 580 module can still be
# in memory after its file was deleted; rebooting then leaves no driver.
module_ko_present() {
  kernel=$(uname -r 2>/dev/null || true)
  [ -n "$kernel" ] || return 1
  root=${FILLRATE_MODULES_ROOT:-/lib/modules}
  for dir in \
    "$root/$kernel/updates/dkms" \
    "$root/$kernel/updates" \
    "$root/$kernel/extra"
  do
    for name in nvidia.ko nvidia.ko.xz nvidia.ko.zst; do
      if [ -f "$dir/$name" ]; then
        write_fact module_ko "$dir/$name"
        return 0
      fi
    done
  done
  write_fact module_ko ""
  return 1
}

preflight() {
  driver=""
  if command -v nvidia-smi >/dev/null 2>&1; then
    driver=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | awk 'NR==1 { gsub(/ /, "", $1); print $1 }')
  fi
  if ! smi_version_ok "$driver"; then
    driver=""
  fi
  if [ -z "$driver" ]; then
    if [ "$MODE" != "repair" ] && [ "$MODE" != "rollback" ]; then
      if ! command -v nvidia-smi >/dev/null 2>&1; then
        abort "nvidia-smi-missing"
      fi
      abort "nvidia-smi-failed"
    fi
    # Broken userspace (NVML Driver/library mismatch) is expected on a
    # half-installed node. dpkg --audit may list iU packages; that is recorded
    # later and does not abort.
    gpus=$(lspci_nvidia_count)
    case "$gpus" in
      ''|*[!0-9]*) gpus=0 ;;
    esac
    write_fact gpu_count "$gpus"
    write_fact nvidia_smi unavailable
    if ! proc_nvidia_loaded || [ "$gpus" -lt 1 ]; then
      abort "no-gpu"
    fi
    refuse_vast_jobs
  else
    write_fact driver_before "$driver"
    write_fact nvidia_smi ok
    cuda=$(cuda_version)
    write_fact cuda_before "${cuda:-unknown}"
    gpus=$(nvidia-smi -L 2>/dev/null | grep -c '^GPU ' || true)
    write_fact gpu_count "$gpus"
    [ "$gpus" -gt 0 ] || abort "no-gpu"
    if [ "$MODE" = "repair" ] || [ "$MODE" = "rollback" ]; then
      refuse_vast_jobs
    else
      apps=$(nvidia-smi --query-compute-apps=pid,process_name --format=csv,noheader 2>/dev/null | grep -E '[0-9]' || true)
      if [ -n "$apps" ]; then
        abort "gpu-processes"
      fi
      if command -v docker >/dev/null 2>&1; then
        running=$(docker ps -q 2>/dev/null || true)
        if [ -n "$running" ]; then
          abort "rentals-present"
        fi
      fi
    fi
    health=$(nvidia-smi --query-gpu=index,temperature.gpu,ecc.errors.uncorrected.volatile.total,power.limit,clocks.sm,clocks.mem,fan.speed,persistence_mode --format=csv,noheader 2>/dev/null | tr '\n' ';' || true)
    write_fact temps "$health"
    persistence=$(printf '%s\n' "$health" | awk -F, '{ gsub(/[ ;]/, "", $8); if ($8 != "") print $8; exit }')
    write_fact persistence "${persistence:-unknown}"
    write_fact ecc_uncorrected "$(sum_ecc)"
    write_fact xid_count "$(count_xid)"
    mig=$(nvidia-smi mig -lgi 2>/dev/null | tr '\n' ' ' || true)
    write_fact mig "${mig:-none}"
  fi
  if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet nvidia-persistenced; then
    write_fact persistenced active
  else
    write_fact persistenced inactive
  fi
  if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet nvidia-fabricmanager; then
    write_fact fabric_manager active
  else
    write_fact fabric_manager inactive
  fi
  packages=$(package_list)
  write_fact packages "$packages"
  if [ -n "$driver" ]; then
    prev=$(printf '%s\n' "$driver" | awk -F. 'NR==1 { print $1 }')
    write_fact previous_branch "$prev"
    write_fact previous_build "$driver"
  fi
  if command -v nvidia-uninstall >/dev/null 2>&1 || [ -x /usr/bin/nvidia-uninstall ]; then
    write_fact install_method non-apt
    write_fact driver_flavor ""
    abort "non-apt-install"
  fi
  flavor=""
  if [ -n "${driver:-}" ]; then
    flavor=$(classify_flavor "$(read_fact previous_branch || true)")
  fi
  if [ -z "$flavor" ] && { [ "$MODE" = "repair" ] || [ "$MODE" = "rollback" ]; }; then
    flavor=$(classify_flavor_broad)
  fi
  if [ -z "$flavor" ]; then
    write_fact install_method non-apt
    write_fact driver_flavor ""
    abort "non-apt-install"
  fi
  write_fact install_method apt
  write_fact driver_flavor "$flavor"
  holds=$(apt-mark showhold 2>/dev/null | tr '\n' ' ' || true)
  write_fact holds "$holds"
  write_fact nvidia_holds "$(nvidia_only_holds "$holds")"
}

runtime_fail() {
  reason=$1
  if [ "${abort_on_fail:-1}" = "1" ]; then
    abort "$reason"
  fi
  if [ -z "$RUNTIME_ABORT" ]; then
    RUNTIME_ABORT=$reason
  fi
}

preflight_late() {
  abort_on_fail=${1:-1}
  RUNTIME_ABORT=
  if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet unattended-upgrades; then
    write_fact unattended_upgrades active
  else
    write_fact unattended_upgrades inactive
  fi
  dkms_status=$(dkms status 2>/dev/null | tr '\n' ';' || true)
  write_fact dkms "${dkms_status:-missing}"
  if ! command -v dkms >/dev/null 2>&1; then
    runtime_fail "dkms-missing"
  fi
  sb=$(mokutil --sb-state 2>/dev/null || true)
  if printf '%s\n' "$sb" | grep -i -q 'SecureBoot enabled'; then
    signer=$(modinfo -F signer nvidia 2>/dev/null || true)
    if [ -z "$signer" ]; then
      write_fact secure_boot "enabled-unsigned"
      runtime_fail "secure-boot-unsigned"
    else
      write_fact secure_boot "enabled-signed"
    fi
  else
    write_fact secure_boot "${sb:-disabled}"
  fi
  if command -v docker >/dev/null 2>&1; then
    info=$(docker info 2>/dev/null || true)
    if ! printf '%s\n' "$info" | grep -i -q nvidia; then
      write_fact docker_runtime missing
      runtime_fail "docker-nvidia-runtime"
    else
      write_fact docker_runtime present
    fi
  else
    write_fact docker_runtime missing
    runtime_fail "docker-missing"
  fi
  if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet vastai.service; then
    write_fact vast_daemon active
  else
    write_fact vast_daemon inactive
    runtime_fail "vast-daemon-down"
  fi
  root_kb=$(disk_kb /)
  boot_kb=$(disk_kb /boot || true)
  write_fact space_root_kb "${root_kb:-0}"
  write_fact space_boot_kb "${boot_kb:-0}"
  if [ -z "$root_kb" ] || [ "$root_kb" -lt "$MIN_ROOT_KB" ]; then
    runtime_fail "low-space-root"
  elif [ -n "$boot_kb" ] && [ "$boot_kb" -lt "$MIN_BOOT_KB" ]; then
    runtime_fail "low-space-boot"
  fi
}

hold_has() {
  name=$1
  list=$2
  printf '%s\n' "$list" | sed 's/:.*//' | grep -F -x "$name" >/dev/null 2>&1
}

conflict_section_names() {
  show=$1
  name=$2
  printf '%s\n' "$show" | awk -v n="$name" '
    /^Conflicts:/ || /^Breaks:/ || /^Replaces:/ { on = 1 }
    /^[A-Za-z]/ && $0 !~ /^(Conflicts|Breaks|Replaces):/ { on = 0 }
    on && index($0, n) { found = 1 }
    END { exit !found }
  '
}

record_hold_conflict() {
  spec=$1
  old=$(read_fact previous_branch || true)
  flavor=$(read_fact driver_flavor || true)
  hold_list=$(apt-mark showhold 2>/dev/null || true)
  printf '%s\n' "$hold_list" > "$STATE/apt-holds.txt"
  show=$(apt-cache show "$spec" 2>/dev/null || true)
  printf '%s\n' "$show" > "$STATE/apt-show.txt"
  meta=
  conflicts=no
  if [ -n "$old" ]; then
    if [ "$flavor" = "open" ]; then
      names="nvidia-driver-${old}-open nvidia-driver-${old}"
    else
      names="nvidia-driver-${old} nvidia-driver-${old}-open"
    fi
    for name in $names; do
      if hold_has "$name" "$hold_list"; then
        meta=$name
        break
      fi
    done
  fi
  if [ -n "$meta" ] && [ "$meta" != "$spec" ]; then
    conflicts=yes
  fi
  if hold_has "$spec" "$hold_list"; then
    if [ -z "$meta" ]; then
      meta=$spec
    fi
    conflicts=yes
  fi
  if [ -n "$meta" ] && [ -n "$show" ] && conflict_section_names "$show" "$meta"; then
    conflicts=yes
  fi
  write_fact held_meta "$meta"
  write_fact held_meta_conflicts "$conflicts"
}

capture_apt_context() {
  spec=${PKG_SPEC:-}
  if [ -z "$spec" ]; then
    spec=$(meta_for "$BRANCH")
  fi
  if [ ! -s "$STATE/apt-policy.txt" ]; then
    apt-cache policy "$spec" > "$STATE/apt-policy.txt" 2>&1 || true
  fi
  if [ ! -s "$STATE/apt-holds.txt" ]; then
    apt-mark showhold > "$STATE/apt-holds.txt" 2>&1 || true
  fi
  if ! read_fact held_meta_conflicts >/dev/null 2>&1; then
    record_hold_conflict "$spec" || true
  fi
}

version_listed() {
  policy=$1
  ver=$2
  printf '%s\n' "$policy" | awk -v v="$ver" '
    $1 == v { found = 1 }
    $1 == "***" && $2 == v { found = 1 }
    $1 == "Candidate:" && $2 == v { found = 1 }
    END { exit !found }
  '
}

sim_version_miss() {
  grep -E -q "was not found|has no installation candidate|Couldn't find any package|Unable to locate package" "$STATE/apt-sim.txt"
}

refresh_policy() {
  spec=$1
  apt-cache policy "$spec" > "$STATE/apt-policy.txt" 2>&1 || true
  awk '/Candidate:/ { print $2; exit }' "$STATE/apt-policy.txt" 2>/dev/null || true
}

run_apt_update() {
  printf '%s\n' "+ apt-get update" > "$STATE/apt-update.txt"
  if apt-get update >> "$STATE/apt-update.txt" 2>&1; then
    write_fact apt_update ok
    UPDATE_FAILED=0
    return 0
  fi
  write_fact apt_update failed
  write_fact apt_command "apt-get update"
  UPDATE_FAILED=1
  if [ "$MODE" != "dry-run" ]; then
    abort "apt-update-failed"
  fi
  return 1
}

mark_sim_failed() {
  reason=$1
  sim_state=${2:-failed}
  write_fact apt_sim "$sim_state"
  SIM_FAILED=1
  SIM_REASON=$reason
  # Apply aborts here. A confirming re-sim returns so the caller can rehold.
  if [ "$MODE" != "dry-run" ] && [ "$CONFIRM_PLAN" != "1" ]; then
    abort "$reason"
  fi
  return 1
}

# 0 means this package must not be removed. Driver-stack packages may go.
# docker, containerd, libvirt, and the NVIDIA container packages may not.
removal_forbidden() {
  pkg=$1
  case "$pkg" in
    docker|docker-*|containerd|containerd-*|libvirt|libvirt*|nvidia-container-toolkit|nvidia-container-toolkit-*|nvidia-container-runtime|nvidia-container-runtime-*|nvidia-docker|nvidia-docker2)
      return 0
      ;;
    nvidia-*|libnvidia-*|xserver-xorg-video-nvidia*|linux-modules-nvidia-*|cuda-compat|cuda-compat-*)
      return 1
      ;;
    *)
      return 0
      ;;
  esac
}

first_disallowed_removal() {
  sim_file=$1
  awk '
    $1 == "Remv" {
      pkg = $2
      sub(/:.*/, "", pkg)
      print pkg
    }
  ' "$sim_file" | while IFS= read -r pkg; do
    [ -n "$pkg" ] || continue
    if removal_forbidden "$pkg"; then
      printf '%s\n' "$pkg"
      exit 0
    fi
  done
}

# 0 when the package name has the branch as its own hyphen-separated piece.
name_has_branch() {
  pkg=$1
  branch=$2
  [ -n "$branch" ] || return 1
  printf '%s\n' "$pkg" | awk -F- -v b="$branch" '{
    for (i = 1; i <= NF; i++) if ($i == b) found = 1
    exit !found
  }'
}

# A removal that survived the allowlist still has to be an old-branch driver
# package or one of the replaced-by-new-branch pairs. cuda-* is not removed.
removal_expected() {
  pkg=$1
  old=$2
  case "$pkg" in
    nvidia-persistenced|libnvidia-egl-gbm1) return 0 ;;
  esac
  # Only the pre-removal file-conflict check may add these. A simulated
  # removal of the same name, without that fact, is still unexpected.
  extras=$(read_fact file_conflict_removals || true)
  case " $extras " in
    *" $pkg "*) return 0 ;;
  esac
  case "$pkg" in
    nvidia-*|libnvidia-*|xserver-xorg-video-nvidia*|linux-modules-nvidia-*) ;;
    *) return 1 ;;
  esac
  name_has_branch "$pkg" "$old"
}

first_unexpected_removal() {
  sim_file=$1
  old=$(installed_branch)
  [ -n "$old" ] || return 0
  awk '
    $1 == "Remv" {
      pkg = $2
      sub(/:.*/, "", pkg)
      print pkg
    }
  ' "$sim_file" | while IFS= read -r pkg; do
    [ -n "$pkg" ] || continue
    if removal_expected "$pkg" "$old"; then
      continue
    fi
    printf '%s\n' "$pkg"
    exit 0
  done
}

# Upgrades stay inside the driver family. cuda-*, docker*, containerd,
# vast*, and kernel/linux-image upgrades are not part of this swap.
first_disallowed_upgrade() {
  sim_file=$1
  awk '
    $1 == "Inst" && $0 ~ /\[[^]]+\][[:space:]]*\(/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      print pkg
    }
  ' "$sim_file" | while IFS= read -r pkg; do
    [ -n "$pkg" ] || continue
    case "$pkg" in
      nvidia-*|libnvidia-*|xserver-xorg-video-nvidia*|linux-modules-nvidia-*) continue ;;
    esac
    printf '%s\n' "$pkg"
    exit 0
  done
}

# The sim must install the target meta at the resolved build. The executed
# transaction also pins that version; this catches a solver that named a
# different build.
plan_installs_target() {
  sim_file=$1
  spec_name=$2
  build=$3
  [ -n "$spec_name" ] || return 1
  [ -n "$build" ] || return 1
  awk -v spec="$spec_name" -v build="$build" '
    $1 == "Inst" {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg != spec) next
      found = 1
      ver = ""
      if (match($0, /\([^)]*\)/)) {
        ver = substr($0, RSTART + 1, RLENGTH - 2)
        sub(/ .*/, "", ver)
      }
      if (ver != build) bad = 1
    }
    END { if (!found || bad) exit 1 }
  ' "$sim_file"
}

installed_branch() {
  driver=$(read_fact driver_before || true)
  printf '%s\n' "$driver" | awk -F. 'NR==1 && $1 ~ /^[0-9]+$/ { print $1 }'
}

branch_of_meta() {
  printf '%s\n' "$1" | awk '
    {
      n = $0
      sub(/^nvidia-driver-/, "", n)
      sub(/-open$/, "", n)
      if (n ~ /^[0-9][0-9][0-9][0-9]?$/) print n
    }
  '
}

# Previous-branch meta plus installed driver packages for that branch.
# The meta is included even when dpkg only shows it as held, not unpacked.
branch_removals() {
  old=$1
  target=$2
  [ -n "$old" ] || return 0
  [ -n "$target" ] || return 0
  [ "$old" != "$target" ] || return 0
  flavor=$(read_fact driver_flavor || true)
  if [ "$flavor" = "open" ]; then
    meta="nvidia-driver-${old}-open"
  else
    meta="nvidia-driver-${old}"
  fi
  extras=$(dpkg -l 2>/dev/null | awk -v old="$old" -v meta="$meta" '
    ($1 == "ii" || $1 == "hi") {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg == meta) next
      suffix = "-" old
      idx = index(pkg, suffix)
      if (idx < 1) next
      rest = substr(pkg, idx + length(suffix))
      if (rest != "" && substr(rest, 1, 1) != "-") next
      base = substr(pkg, 1, idx - 1)
      ok = 0
      if (base ~ /^(nvidia-driver|nvidia-dkms|nvidia-kernel-common|nvidia-kernel-source|nvidia-utils|nvidia-compute-utils|nvidia-headless|nvidia-headless-no-dkms|nvidia-firmware|xserver-xorg-video-nvidia|linux-modules-nvidia)$/) ok = 1
      if (base ~ /^libnvidia-/) ok = 1
      if (ok) print pkg
    }
  ' | sort -u)
  out="${meta}-"
  for extra in $extras; do
    [ "$extra" = "$meta" ] && continue
    out="$out ${extra}-"
  done
  printf '%s\n' "$out"
}

reset_strategies() {
  rm -rf "$STATE/strategies"
  mkdir -p "$STATE/strategies"
  : > "$STATE/strategies/order"
  STRAT_SEQ=0
}

record_attempt() {
  attempt_id=$1
  attempt_rc=$2
  attempt_cmd=$3
  attempt_out=$4
  STRAT_SEQ=$((STRAT_SEQ + 1))
  n=$(printf '%03d' "$STRAT_SEQ")
  printf '%s\n' "$n" >> "$STATE/strategies/order"
  printf '%s\n' "$attempt_id" > "$STATE/strategies/$n.id"
  printf '%s\n' "$attempt_rc" > "$STATE/strategies/$n.rc"
  printf '%s\n' "$attempt_cmd" > "$STATE/strategies/$n.cmd"
  cp "$attempt_out" "$STATE/strategies/$n.out"
}

sim_installs_prime() {
  awk '
    $1 == "Inst" {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg == "nvidia-prime") found = 1
    }
    END { exit !found }
  ' "$1"
}

dpkg_status_path() {
  if [ -n "${FILLRATE_DPKG_STATUS:-}" ]; then
    printf '%s\n' "$FILLRATE_DPKG_STATUS"
    return 0
  fi
  printf '%s\n' /var/lib/dpkg/status
}

# Copy dpkg status and mark NVIDIA holds as install. apt-get -s then solves
# the problem apply will solve after apt-mark unhold. Docker, containerd,
# libvirt, and xfsprogs holds stay. The live status is not changed, so a
# dry-run does not unhold anything.
prepare_unheld_status() {
  src=$(dpkg_status_path)
  dest=$STATE/dpkg-status-unheld
  rm -f "$dest"
  if [ ! -f "$src" ]; then
    return 1
  fi
  holds=$(read_fact nvidia_holds || true)
  list=$STATE/nvidia-holds.list
  : > "$list"
  for name in $holds; do
    [ -n "$name" ] || continue
    printf '%s\n' "${name%%:*}" >> "$list"
  done
  if [ ! -s "$list" ]; then
    return 1
  fi
  if ! awk '
    NR == FNR {
      if ($0 != "") hold[$0] = 1
      next
    }
    /^Package: / { pkg = $2 }
    /^Status: hold / && pkg != "" && (pkg in hold) {
      sub(/^Status: hold /, "Status: install ")
    }
    /^$/ { pkg = "" }
    { print }
  ' "$list" "$src" > "$dest"; then
    rm -f "$dest"
    return 1
  fi
  return 0
}

# Same solver inputs for dry-run and for apply, before and after the unhold.
refresh_sim_opts() {
  SIM_APT_OPTS="-o APT::Get::Allow-Change-Held-Packages=true"
  if prepare_unheld_status; then
    SIM_APT_OPTS="$SIM_APT_OPTS -o Dir::State::status=$STATE/dpkg-status-unheld"
    return 0
  fi
  # No status file to edit. Ignore holds so a held NVIDIA package can still
  # leave the transaction. confirm_plan runs after the real unhold, where
  # this is a no-op for packages that are already not held.
  if [ ! -f "$(dpkg_status_path)" ]; then
    SIM_APT_OPTS="$SIM_APT_OPTS -o APT::Ignore-Hold=true"
  fi
}

# Read-only. Exit 0 only when apt accepts the plan and it does not install nvidia-prime.
run_strategy() {
  strat_id=$1
  strat_extra=$2
  strat_removals=$3
  strat_pkg=$4
  strat_cmd="apt-get -s $SIM_APT_OPTS install"
  if [ -n "$strat_extra" ]; then
    strat_cmd="$strat_cmd $strat_extra"
  fi
  strat_cmd="$strat_cmd $strat_pkg"
  if [ -n "$strat_removals" ]; then
    strat_cmd="$strat_cmd $strat_removals"
  fi
  strat_out="$STATE/strategy-attempt.txt"
  printf '%s\n' "+ ${strat_cmd}" > "$strat_out"
  # shellcheck disable=SC2086
  apt-get -s $SIM_APT_OPTS install $strat_extra $strat_pkg $strat_removals >> "$strat_out" 2>&1
  strat_rc=$?
  record_attempt "$strat_id" "$strat_rc" "$strat_cmd" "$strat_out"
  if [ "$strat_rc" -ne 0 ]; then
    return 1
  fi
  if sim_installs_prime "$strat_out"; then
    PRIME_HIT=1
    return 1
  fi
  SIM_STRATEGY=$strat_id
  if [ -n "$strat_removals" ]; then
    PLAN_PKGS="$strat_pkg $strat_removals"
  else
    PLAN_PKGS=$strat_pkg
  fi
  if [ "$strat_extra" = "--no-install-recommends" ]; then
    PLAN_NO_RECOMMENDS=1
  else
    PLAN_NO_RECOMMENDS=0
  fi
  cp "$strat_out" "$STATE/apt-sim.txt"
  cp "$strat_out" "$STATE/winner-sim.txt"
  write_fact sim_strategy "$strat_id"
  write_fact apt_command "$strat_cmd"
  return 0
}

# apt-cache show always runs. dpkg -L is decisive only when the package is
# already installed and its file list can be read. A missing list before the
# swap is not a failure: the new compute-utils package is not installed yet.
provider_ships_persistenced() {
  pkg=$1
  apt-cache show "$pkg" > "$STATE/provider-show.txt" 2>&1 || true
  if dpkg -L "$pkg" > "$STATE/provider-files.txt" 2>/dev/null; then
    if grep -qx '/usr/bin/nvidia-persistenced' "$STATE/provider-files.txt" \
      && grep -E -q '(^|/)nvidia-persistenced\.service$' "$STATE/provider-files.txt"; then
      write_fact persistenced_provider_files "present"
      return 0
    fi
    write_fact persistenced_provider_files "absent"
    return 1
  fi
  write_fact persistenced_provider_files "unverified"
  return 0
}

# Standalone nvidia-persistenced and libnvidia-egl-gbm1 leave with the old
# stack. The new branch packages replace them. Record the pair. Abort when
# persistenced is removed and nvidia-compute-utils-<new> is not installed
# by this plan.
accept_expected_replacements() {
  inst=$(read_fact apt_install || true)
  upg=$(read_fact apt_upgrade || true)
  remv=$(read_fact apt_remove || true)
  new_branch=$(branch_of_meta "${spec:-}")
  if [ -z "$new_branch" ]; then
    new_branch=$BRANCH
  fi
  provider="nvidia-compute-utils-${new_branch}"
  glpkg="libnvidia-gl-${new_branch}"
  replaced=""
  if list_has "$remv" nvidia-persistenced; then
    if ! list_has "$inst" "$provider" && ! list_has "$upg" "$provider"; then
      mark_sim_failed "persistenced-not-replaced" "would-remove"
      return 1
    fi
    if ! provider_ships_persistenced "$provider"; then
      mark_sim_failed "persistenced-provider-incomplete" "would-remove"
      return 1
    fi
    replaced="nvidia-persistenced=${provider}"
  fi
  if list_has "$remv" libnvidia-egl-gbm1; then
    if list_has "$inst" "$glpkg" || list_has "$upg" "$glpkg"; then
      if [ -n "$replaced" ]; then
        replaced="$replaced libnvidia-egl-gbm1=${glpkg}"
      else
        replaced="libnvidia-egl-gbm1=${glpkg}"
      fi
    fi
  fi
  write_fact replaced_by_new_branch "$replaced"
  return 0
}

maybe_fallback() {
  fb_spec=$1
  [ "${FELL_BACK:-0}" = "1" ] && return 1
  pinned=$BUILD
  candidate=$(read_fact apt_candidate || true)
  [ -n "$pinned" ] || return 1
  [ -n "$candidate" ] || return 1
  [ "$candidate" != "(none)" ] || return 1
  [ "$candidate" != "$pinned" ] || return 1
  policy=$(cat "$STATE/apt-policy.txt" 2>/dev/null || true)
  cp "$STATE/strategy-attempt.txt" "$STATE/apt-sim.txt"
  if version_listed "$policy" "$pinned" && ! sim_version_miss; then
    return 1
  fi
  write_fact version_fallback "pinned ${pinned} unavailable; using ${candidate}"
  cp "$STATE/strategy-attempt.txt" "$STATE/apt-sim-pinned.txt"
  FELL_BACK=1
  BUILD=$candidate
  write_fact resolved_build "$candidate"
  PKG="${fb_spec}=${candidate}"
  return 0
}

try_strategies() {
  ts_pkg=$1
  ts_removals=$2
  ts_spec=$3
  if run_strategy A "" "" "$ts_pkg"; then
    return 0
  fi
  if maybe_fallback "$ts_spec"; then
    ts_pkg=$PKG
    if run_strategy A "" "" "$ts_pkg"; then
      return 0
    fi
  fi
  if run_strategy B "--no-install-recommends" "" "$ts_pkg"; then
    return 0
  fi
  if run_strategy C "" "$ts_removals" "$ts_pkg"; then
    return 0
  fi
  if run_strategy D "--no-install-recommends" "$ts_removals" "$ts_pkg"; then
    return 0
  fi
  return 1
}

combine_attempts() {
  : > "$STATE/apt-sim.txt"
  [ -f "$STATE/strategies/order" ] || return 0
  while IFS= read -r n; do
    [ -n "$n" ] || continue
    id=$(cat "$STATE/strategies/$n.id" 2>/dev/null || true)
    rc=$(cat "$STATE/strategies/$n.rc" 2>/dev/null || true)
    {
      printf '%s\n' "---- strategy ${id} exit ${rc} ----"
      cat "$STATE/strategies/$n.out"
    } >> "$STATE/apt-sim.txt"
  done < "$STATE/strategies/order"
  last=$(awk 'END { print }' "$STATE/strategies/order")
  if [ -n "$last" ] && [ -f "$STATE/strategies/$last.cmd" ]; then
    write_fact apt_command "$(cat "$STATE/strategies/$last.cmd")"
  fi
}

judge_plan() {
  # Winner only. The concatenated attempt log is for the error text.
  plan_file="$STATE/winner-sim.txt"
  if [ ! -f "$plan_file" ]; then
    plan_file="$STATE/apt-sim.txt"
  fi
  record_sim "$plan_file"
  bad=$(protected_removal "$plan_file" || true)
  if [ -n "$bad" ]; then
    mark_sim_failed "apt-would-remove-${bad}" "would-remove"
    return 1
  fi
  bad=$(first_disallowed_removal "$plan_file" || true)
  if [ -n "$bad" ]; then
    mark_sim_failed "apt-would-remove-${bad}" "would-remove"
    return 1
  fi
  bad=$(first_unexpected_removal "$plan_file" || true)
  if [ -n "$bad" ]; then
    mark_sim_failed "apt-would-remove-${bad}" "would-remove"
    return 1
  fi
  touched=$(cuda_toolkit_install "$plan_file" || true)
  if [ -n "$touched" ]; then
    mark_sim_failed "apt-would-touch-${touched}" "would-touch"
    return 1
  fi
  bad=$(first_disallowed_upgrade "$plan_file" || true)
  if [ -n "$bad" ]; then
    mark_sim_failed "apt-would-upgrade-${bad}" "would-upgrade"
    return 1
  fi
  inst_list=$(read_fact apt_install || true)
  up_list=$(read_fact apt_upgrade || true)
  case " ${inst_list} ${up_list} " in
    *" nvidia-prime "*)
      mark_sim_failed "apt-would-install-nvidia-prime" "would-install"
      return 1
      ;;
  esac
  if ! plan_installs_target "$plan_file" "$spec" "$BUILD"; then
    mark_sim_failed "target-not-in-plan" "would-install"
    return 1
  fi
  if ! accept_expected_replacements; then
    return 1
  fi
  if [ "$UPDATE_FAILED" = "1" ]; then
    write_fact apt_command "apt-get update"
    mark_sim_failed "apt-update-failed" ok
    return 1
  fi
  write_fact apt_sim ok
  return 0
}

resolve_build() {
  spec=$1
  candidate=$(refresh_policy "$spec")
  write_fact apt_candidate "${candidate:-}"
  if [ -n "$BUILD" ]; then
    write_fact resolved_build "$BUILD"
    PKG="${spec}=${BUILD}"
    return 0
  fi
  if [ -z "$candidate" ] || [ "$candidate" = "(none)" ]; then
    write_fact apt_sim failed
    write_fact apt_command "apt-cache policy ${spec}"
    mark_sim_failed "unresolved-build"
    return 1
  fi
  BUILD=$candidate
  write_fact resolved_build "$candidate"
  PKG="${spec}=${candidate}"
  return 0
}

simulate() {
  pkg=$1
  spec=$2
  record_hold_conflict "$spec"
  if [ "$SIM_FAILED" = "1" ]; then
    return 1
  fi
  refresh_sim_opts
  reset_strategies
  SIM_STRATEGY=
  PRIME_HIT=0
  FELL_BACK=0
  PLAN_PKGS=
  PLAN_NO_RECOMMENDS=0
  rm -f "$STATE/winner-sim.txt"
  write_fact sim_strategy ""
  write_fact replaced_by_new_branch ""
  write_fact persistenced_provider_files ""
  write_fact apt_install ""
  write_fact apt_remove ""
  write_fact apt_upgrade ""
  old_branch=$(installed_branch)
  target_branch=$(branch_of_meta "$spec")
  removals=$(branch_removals "$old_branch" "$target_branch")
  if try_strategies "$pkg" "$removals" "$spec"; then
    judge_plan
    return $?
  fi
  # No winner. Keep every attempt's text for the error, but do not mix
  # their install and remove sets into the report.
  combine_attempts
  write_fact apt_install ""
  write_fact apt_remove ""
  write_fact apt_upgrade ""
  write_fact sim_strategy ""
  if [ "$PRIME_HIT" = "1" ]; then
    mark_sim_failed "apt-would-install-nvidia-prime" "would-install"
    return 1
  fi
  mark_sim_failed "apt-sim-failed"
  return 1
}

# Re-simulate after the NVIDIA unhold. The post-unhold plan is installed
# when it passes the same safety checks as the first validation. A different
# strategy or package set is recorded; it is not by itself a failure.
confirm_plan() {
  saved_strategy=$(read_fact sim_strategy || true)
  saved_install=$(read_fact apt_install || true)
  saved_remove=$(read_fact apt_remove || true)
  saved_upgrade=$(read_fact apt_upgrade || true)
  write_fact pre_unhold_strategy "$saved_strategy"
  CONFIRM_PLAN=1
  SIM_FAILED=0
  SIM_REASON=
  if ! simulate "$PKG" "$spec"; then
    CONFIRM_PLAN=0
    fail_after_change "${SIM_REASON:-apt-sim-failed}"
  fi
  CONFIRM_PLAN=0
  now_strategy=$(read_fact sim_strategy || true)
  now_install=$(read_fact apt_install || true)
  now_remove=$(read_fact apt_remove || true)
  now_upgrade=$(read_fact apt_upgrade || true)
  if [ "$now_strategy" != "$saved_strategy" ] || [ "$now_install" != "$saved_install" ] || [ "$now_remove" != "$saved_remove" ] || [ "$now_upgrade" != "$saved_upgrade" ]; then
    write_fact plan_differs_from_pre_unhold yes
  else
    write_fact plan_differs_from_pre_unhold no
  fi
}

# The real transaction is the post-unhold plan, without -s and without the
# held-package simulation flag. NVIDIA holds were already cleared.
real_install() {
  [ -n "$PLAN_PKGS" ] || fail_after_change "apt-install-failed"
  for token in $PLAN_PKGS; do
    name=${token%%=*}
    case "$name" in
      nvidia-prime)
        fail_after_change "apt-would-install-nvidia-prime"
        ;;
    esac
  done
  if [ "$PLAN_NO_RECOMMENDS" = "1" ]; then
    # shellcheck disable=SC2086
    apt-get install -y --no-install-recommends $PLAN_PKGS > "$STATE/apt-install.txt" 2>&1 || recover_failed_install
  else
    # shellcheck disable=SC2086
    apt-get install -y $PLAN_PKGS > "$STATE/apt-install.txt" 2>&1 || recover_failed_install
  fi
  if ! kernel_module_matches; then
    fail_after_change "dkms-kernel-mismatch"
  fi
}

post_abort() {
  start_vast || true
  rehold || true
  abort "$1"
}

persistence_mode() {
  nvidia-smi --query-gpu=index,temperature.gpu,ecc.errors.uncorrected.volatile.total,power.limit,clocks.sm,clocks.mem,fan.speed,persistence_mode --format=csv,noheader 2>/dev/null \
    | awk -F, 'NR==1 { gsub(/ /, "", $8); print $8; exit }'
}

# After the new branch is up: the persistenced unit must exist, enable --now
# must leave it active, and persistence mode must be on. nvidia-smi -pm 1 is
# the fallback when the mode is still off. Failure here is the same post
# path as a missing GPU (rehold, start vastai, non-zero report).
post_ensure_persistenced() {
  provider="nvidia-compute-utils-${BRANCH}"
  if ! provider_ships_persistenced "$provider"; then
    write_fact persistence_detail "${provider} is installed but does not ship /usr/bin/nvidia-persistenced and nvidia-persistenced.service"
    post_abort "persistence-not-restored"
  fi
  unit_ok=0
  if [ -f /lib/systemd/system/nvidia-persistenced.service ] \
    || [ -f /usr/lib/systemd/system/nvidia-persistenced.service ] \
    || [ -f /etc/systemd/system/nvidia-persistenced.service ]; then
    unit_ok=1
  fi
  if command -v systemctl >/dev/null 2>&1 && systemctl cat nvidia-persistenced >/dev/null 2>&1; then
    unit_ok=1
  fi
  if [ "$unit_ok" != "1" ]; then
    write_fact persistence_detail "nvidia-persistenced systemd unit missing"
    write_fact persistenced_after inactive
    write_fact persistence_after "$(persistence_mode || true)"
    post_abort "persistence-not-restored"
  fi
  if command -v systemctl >/dev/null 2>&1; then
    systemctl enable --now nvidia-persistenced >/dev/null 2>&1 || true
  fi
  mode_now=$(persistence_mode || true)
  if [ "$mode_now" != "Enabled" ]; then
    nvidia-smi -pm 1 >/dev/null 2>&1 || true
    mode_now=$(persistence_mode || true)
  fi
  write_fact persistence_after "${mode_now:-unknown}"
  if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet nvidia-persistenced; then
    write_fact persistenced_after active
  else
    write_fact persistenced_after inactive
  fi
  if [ "$mode_now" != "Enabled" ] || [ "$(read_fact persistenced_after || true)" != "active" ]; then
    write_fact persistence_detail "persistence mode=${mode_now:-unknown} persistenced=$(read_fact persistenced_after || true) after systemctl enable --now nvidia-persistenced and nvidia-smi -pm 1"
    post_abort "persistence-not-restored"
  fi
  write_fact persistence_detail "ok"
}

post_checks() {
  if [ ! -f "$STATE/pending-post" ]; then
    abort "missing-post-marker"
  fi
  now=$(date +%s)
  end=$EXPIRES
  if [ -n "${FILLRATE_NET_WAIT_SEC:-}" ]; then
    cap=$((now + FILLRATE_NET_WAIT_SEC))
    if [ "$cap" -lt "$end" ]; then
      end=$cap
    fi
  fi
  online=0
  while [ "$(date +%s)" -le "$end" ]; do
    if ping -c 1 -W 1 1.1.1.1 >/dev/null 2>&1; then
      online=1
      break
    fi
    if [ "$(date +%s)" -ge "$end" ]; then
      break
    fi
    sleep 1
  done
  if [ "$online" != "1" ]; then
    abort "no-network"
  fi
  driver=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | awk 'NR==1 { gsub(/ /, "", $1); print $1 }')
  write_fact driver_after "$driver"
  gpus=$(nvidia-smi -L 2>/dev/null | grep -c '^GPU ' || true)
  write_fact gpu_count_after "$gpus"
  before_gpus=$(read_fact gpu_count || echo 0)
  if [ "$gpus" -lt 1 ] || [ "$gpus" -lt "$before_gpus" ]; then
    post_abort "missing-gpu"
  fi
  resolved=$(read_fact resolved_build || true)
  if [ -n "$resolved" ] && [ "$driver" != "$resolved" ]; then
    post_abort "driver-version-mismatch"
  fi
  case "$driver" in
    "$BRANCH".*) ;;
    *) post_abort "driver-branch-mismatch" ;;
  esac
  xid_after=$(count_xid)
  write_fact xid_count_after "$xid_after"
  xid_before=$(read_fact xid_count || echo 0)
  if [ "$xid_after" -gt "$xid_before" ]; then
    post_abort "new-xid"
  fi
  ecc_after=$(sum_ecc)
  write_fact ecc_uncorrected_after "$ecc_after"
  ecc_before=$(read_fact ecc_uncorrected || echo 0)
  if [ "$ecc_after" -gt "$ecc_before" ]; then
    post_abort "ecc-increased"
  fi
  if [ "$(read_fact fabric_manager || true)" = "active" ]; then
    systemctl start nvidia-fabricmanager >/dev/null 2>&1 || post_abort "fabric-manager-down"
    systemctl is-active --quiet nvidia-fabricmanager || post_abort "fabric-manager-down"
    if ! printf '%s\n' "$(package_list)" | grep -F -q 'fabricmanager'; then
      post_abort "fabric-manager-mismatch"
    fi
  fi
  post_ensure_persistenced
  if ! docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi >/dev/null 2>&1; then
    post_abort "docker-gpu-test-failed"
  fi
  if ! has_pkg nvidia-cuda-toolkit && printf '%s\n' "$(read_fact packages || true)" | grep -F -q 'nvidia-cuda-toolkit='; then
    post_abort "cuda-toolkit-removed"
  fi
  if ! has_pkg cuda-drivers && printf '%s\n' "$(read_fact packages || true)" | grep -E -q '(^| )cuda-drivers='; then
    post_abort "cuda-drivers-removed"
  fi
  old_packages=$(read_fact packages || true)
  for token in $old_packages; do
    pkg=${token%%=*}
    case "$pkg" in
      cuda-toolkit|cuda-toolkit-*)
        has_pkg "$pkg" || post_abort "cuda-toolkit-removed"
        ;;
    esac
  done
  rehold || post_abort "holds-not-restored"
  start_vast || post_abort "vast-daemon-not-started"
  systemctl is-active --quiet vastai.service || post_abort "vast-daemon-not-started"
  write_fact result ok
  write_fact abort_reason ""
  emit_report
  remove_post_unit
  exit 0
}

last_lines() {
  src=$1
  dest=$2
  if [ -f "$src" ]; then
    tail -n 200 "$src" > "$dest" 2>/dev/null || : > "$dest"
  else
    : > "$dest"
  fi
}

# Copy root-only apt and DKMS logs into the state dir so the JSON report
# can carry them. Each source is capped at the last 200 lines.
capture_diagnostics() {
  last_lines "$VAR_LOG/apt/term.log" "$STATE/diag-term.log"
  last_lines "$VAR_LOG/apt/history.log" "$STATE/diag-history.log"
  raw="$STATE/diag-dkms-make.raw"
  : > "$raw"
  for f in "$VAR_LIB"/dkms/nvidia*/*/build/make.log; do
    [ -f "$f" ] || continue
    printf '%s\n' "----- $f -----" >> "$raw"
    tail -n 200 "$f" >> "$raw"
  done
  tail -n 200 "$raw" > "$STATE/diag-dkms-make.log" 2>/dev/null || : > "$STATE/diag-dkms-make.log"
  raw="$STATE/diag-dkms.raw"
  : > "$raw"
  for f in "$VAR_LOG"/dkms*; do
    [ -f "$f" ] || continue
    printf '%s\n' "----- $f -----" >> "$raw"
    tail -n 200 "$f" >> "$raw"
  done
  tail -n 200 "$raw" > "$STATE/diag-dkms.log" 2>/dev/null || : > "$STATE/diag-dkms.log"
  dpkg -l 2>/dev/null | grep -i nvidia | tail -n 200 > "$STATE/diag-dpkg-nvidia.txt" || : > "$STATE/diag-dpkg-nvidia.txt"
  dpkg --audit > "$STATE/diag-dpkg-audit.raw" 2>&1 || true
  tail -n 200 "$STATE/diag-dpkg-audit.raw" > "$STATE/diag-dpkg-audit.txt" 2>/dev/null || : > "$STATE/diag-dpkg-audit.txt"
  dkms status > "$STATE/diag-dkms-status.raw" 2>&1 || true
  tail -n 200 "$STATE/diag-dkms-status.raw" > "$STATE/diag-dkms-status.txt" 2>/dev/null || : > "$STATE/diag-dkms-status.txt"
  : > "$STATE/diag-crash-names.txt"
  for dir in "$VAR_CRASH" "$VAR_LIB/crash"; do
    [ -d "$dir" ] || continue
    for f in "$dir"/*; do
      [ -f "$f" ] || continue
      basename "$f" >> "$STATE/diag-crash-names.txt"
    done
  done
  if [ -s "$STATE/diag-crash-names.txt" ]; then
    sort -u "$STATE/diag-crash-names.txt" | head -n 20 > "$STATE/diag-crash.txt"
  else
    : > "$STATE/diag-crash.txt"
  fi
  capture_apt_state_logs
}

# Last 120 lines of every apt/dpkg/dkms transcript in the state directory,
# plus the live hold list, nvidia package table, and target policy. These
# files are root-only; the JSON report is what the operator can read.
capture_apt_state_logs() {
  out="$STATE/diag-apt-state.txt"
  : > "$out"
  for f in "$STATE"/apt-*.txt "$STATE"/dpkg-configure.txt "$STATE"/dkms-autoinstall.txt "$STATE"/dkms-autoinstall-restore.txt; do
    [ -f "$f" ] || continue
    printf '%s\n' "===== $(basename "$f") =====" >> "$out"
    tail -n 120 "$f" >> "$out"
    printf '\n' >> "$out"
  done
  apt-mark showhold > "$STATE/diag-showhold.txt" 2>&1 || true
  dpkg -l 2>/dev/null | grep -E 'nvidia|libnvidia' | tail -n 120 > "$STATE/diag-dpkg-live.txt" || true
  meta=$(read_fact target_package || true)
  meta=${meta%%=*}
  if [ -z "$meta" ]; then
    meta=$(meta_for "$BRANCH")
  fi
  if [ -n "$meta" ]; then
    apt-cache policy "$meta" 2>&1 | tail -n 120 > "$STATE/diag-policy-target.txt" || true
  else
    : > "$STATE/diag-policy-target.txt"
  fi
}

half_installed_branch() {
  dpkg -l 2>/dev/null | awk '
    $1 ~ /^(iU|iF|iH)$/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg !~ /^nvidia-/) next
      n = pkg
      sub(/-open$/, "", n)
      if (match(n, /[0-9]{3,4}$/)) { print substr(n, RSTART, RLENGTH); exit }
    }
  '
}

rc_driver_branch() {
  dpkg -l 2>/dev/null | awk '
    $1 == "rc" {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg !~ /^nvidia-(dkms|driver)-/) next
      n = pkg
      sub(/-open$/, "", n)
      if (match(n, /[0-9]{3,4}$/)) { print substr(n, RSTART, RLENGTH); exit }
    }
  '
}

# Restore target: explicit restore_branch, saved facts, the ticket branch
# when it is not the half-installed new branch, an rc driver package, then 580.
choose_restore_branch() {
  if printf '%s\n' "${FILLRATE_RESTORE_BRANCH:-}" | grep -E -q '^[0-9]{3,4}$'; then
    printf '%s\n' "$FILLRATE_RESTORE_BRANCH"
    return 0
  fi
  saved=$(read_fact previous_branch || true)
  if printf '%s\n' "$saved" | grep -E -q '^[0-9]{3,4}$'; then
    printf '%s\n' "$saved"
    return 0
  fi
  broken=$(half_installed_branch || true)
  finished=$(read_fact forward_branch || true)
  if printf '%s\n' "$BRANCH" | grep -E -q '^[0-9]{3,4}$'; then
    if [ "$BRANCH" != "$finished" ]; then
      if [ -z "$broken" ] || [ "$BRANCH" != "$broken" ]; then
        printf '%s\n' "$BRANCH"
        return 0
      fi
    fi
  fi
  rc=$(rc_driver_branch || true)
  if printf '%s\n' "$rc" | grep -E -q '^[0-9]{3,4}$'; then
    printf '%s\n' "$rc"
    return 0
  fi
  printf '%s\n' "580"
}

version_matches_branch() {
  ver=$1
  br=$2
  printf '%s\n' "$ver" | awk -F'[.-]' -v b="$br" 'NR==1 && $1 == b { found = 1 } END { exit !found }'
}

choose_restore_spec() {
  br=$1
  meta=$(meta_for "$br")
  policy=$(apt-cache policy "$meta" 2>/dev/null || true)
  saved_v=$(read_fact previous_build || true)
  deb=$(dpkg -l 2>/dev/null | awk -v meta="$meta" -v br="$br" '
    $1 ~ /^(ii|hi|rc|hc|iU)$/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg == meta || pkg == ("nvidia-dkms-" br) || pkg == ("nvidia-dkms-" br "-open")) {
        print $3
        exit
      }
    }
  ')
  build=""
  for cand in $saved_v $deb; do
    [ -n "$cand" ] || continue
    if version_matches_branch "$cand" "$br" && version_listed "$policy" "$cand"; then
      build=$cand
      break
    fi
  done
  if [ -z "$build" ]; then
    cand=$(printf '%s\n' "$policy" | awk '/Candidate:/ { print $2; exit }')
    if [ -n "$cand" ] && [ "$cand" != "(none)" ] && version_matches_branch "$cand" "$br" && version_listed "$policy" "$cand"; then
      build=$cand
    fi
  fi
  if [ -n "$build" ]; then
    printf '%s\n' "${meta}=${build}"
  else
    printf '%s\n' "$meta"
  fi
}

unconfigured_nvidia_pkgs() {
  br=$1
  dpkg -l 2>/dev/null | awk -v br="$br" '
    $1 ~ /^(iU|iF|iH)$/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg !~ /^(nvidia-|libnvidia-|xserver-xorg-video-nvidia|linux-modules-nvidia)/) next
      n = pkg
      sub(/-open$/, "", n)
      major = ""
      if (match(n, /[0-9]{3,4}$/)) major = substr(n, RSTART, RLENGTH)
      if (major == br) next
      print pkg
    }
  '
}

dkms_built_for() {
  br=$1
  kernel=$(uname -r 2>/dev/null || true)
  [ -n "$kernel" ] || return 1
  dkms status 2>/dev/null | awk -v b="$br" -v k="$kernel" '
    index($0, "nvidia") && index($0, k) && index($0, b ".") && $0 ~ /installed/ { found = 1 }
    END { exit !found }
  '
}

loaded_module_version() {
  modinfo -F version nvidia 2>/dev/null | awk 'NR==1 { gsub(/ /, "", $0); print; exit }'
}

userspace_driver_version() {
  nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | awk 'NR==1 { gsub(/ /, "", $1); print; exit }'
}

versions_match() {
  loaded=$(loaded_module_version || true)
  user=$(userspace_driver_version || true)
  [ -n "$loaded" ] && [ -n "$user" ] && [ "$loaded" = "$user" ]
}

rehold_restore() {
  br=$1
  flavor=$(read_fact driver_flavor || true)
  if [ "$flavor" = "open" ]; then
    meta="nvidia-driver-${br}-open"
  else
    meta="nvidia-driver-${br}"
  fi
  if [ -n "$ORIG_HOLDS" ]; then
    list=$ORIG_HOLDS
  elif [ -n "$ORIG_NVIDIA_HOLDS" ]; then
    list="$ORIG_NVIDIA_HOLDS $BASELINE_HOLDS"
  elif [ "$MODE" = "repair" ]; then
    list="$BASELINE_HOLDS $meta"
  else
    list=$(read_fact holds || true)
    if [ -z "$list" ]; then
      list=$(read_fact nvidia_holds || true)
    fi
    if [ -z "$list" ]; then
      list="$BASELINE_HOLDS $meta"
    fi
  fi
  case " $list " in
    *" $meta "*) ;;
    *) list="$list $meta" ;;
  esac
  # shellcheck disable=SC2086
  apt-mark hold $list >/dev/null 2>&1 || return 1
  HOLDS_RESTORED=1
  return 0
}

restart_services_strict() {
  if command -v systemctl >/dev/null 2>&1; then
    systemctl reset-failed nvidia-persistenced >/dev/null 2>&1 || true
    systemctl reset-failed vastai.service >/dev/null 2>&1 || true
    systemctl enable --now nvidia-persistenced >/dev/null 2>&1 || true
    systemctl restart nvidia-persistenced >/dev/null 2>&1 || true
    systemctl start vastai.service >/dev/null 2>&1 || {
      note_step_fail "vast-daemon-not-started" "vastai" "1"
      return 1
    }
  fi
  nvidia-smi -pm 1 >/dev/null 2>&1 || true
  if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet nvidia-persistenced; then
    write_fact persistenced_after active
  else
    write_fact persistenced_after inactive
    write_fact persistence_detail "nvidia-persistenced inactive after systemctl enable --now"
    note_step_fail "persistence-not-restored" "persistenced" "1"
    return 1
  fi
  mode_now=$(persistence_mode || true)
  write_fact persistence_after "${mode_now:-unknown}"
  if command -v systemctl >/dev/null 2>&1 && ! systemctl is-active --quiet vastai.service; then
    note_step_fail "vast-daemon-not-started" "vastai" "1"
    return 1
  fi
  write_fact vast_daemon active
  return 0
}

verify_restored() {
  br=$1
  driver=$(userspace_driver_version || true)
  if [ -z "$driver" ]; then
    note_step_fail "nvidia-smi-failed" "nvidia-smi" "1"
    return 1
  fi
  write_fact driver_after "$driver"
  gpus=$(nvidia-smi -L 2>/dev/null | grep -c '^GPU ' || true)
  write_fact gpu_count_after "$gpus"
  before=$(read_fact gpu_count || true)
  case "$before" in
    ''|*[!0-9]*) before=0 ;;
  esac
  if [ "$gpus" -lt 1 ]; then
    note_step_fail "missing-gpu" "nvidia-smi" "1"
    return 1
  fi
  if [ "$before" -gt 0 ] && [ "$gpus" -lt "$before" ]; then
    note_step_fail "missing-gpu" "nvidia-smi" "1"
    return 1
  fi
  case "$driver" in
    "$br".*) ;;
    *)
      note_step_fail "driver-branch-mismatch" "nvidia-smi" "1"
      return 1
      ;;
  esac
  if ! docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi >/dev/null 2>&1; then
    note_step_fail "docker-gpu-test-failed" "docker-gpu" "1"
    return 1
  fi
  return 0
}

refresh_apt_lists() {
  if [ -s "$STATE/apt-update.txt" ]; then
    printf '\n%s\n' "+ apt-get update" >> "$STATE/apt-update.txt"
  else
    printf '%s\n' "+ apt-get update" > "$STATE/apt-update.txt"
  fi
  if apt-get update >> "$STATE/apt-update.txt" 2>&1; then
    write_fact apt_update ok
    return 0
  fi
  write_fact apt_update failed
  return 1
}

remove_recorded_conflicts() {
  log=$1
  extras=$(read_fact file_conflict_removals || true)
  [ -n "$extras" ] || return 0
  # shellcheck disable=SC2086
  apt-get remove -y $extras > "$log" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "file-conflict-remove-failed" "conflict-remove" "$rc"
    return 1
  fi
  return 0
}

dkms_autoinstall_running() {
  log=$1
  kernel=$(uname -r 2>/dev/null || true)
  if [ -n "$kernel" ]; then
    if [ -s "$log" ]; then
      printf '\n===== dkms autoinstall -k %s =====\n' "$kernel" >> "$log"
      dkms autoinstall -k "$kernel" >> "$log" 2>&1 || true
    else
      dkms autoinstall -k "$kernel" > "$log" 2>&1 || true
    fi
  else
    if [ -s "$log" ]; then
      printf '\n===== dkms autoinstall =====\n' >> "$log"
      dkms autoinstall >> "$log" 2>&1 || true
    else
      dkms autoinstall > "$log" 2>&1 || true
    fi
  fi
}

ensure_persistenced_pkg() {
  br=$1
  provider="nvidia-compute-utils-${br}"
  if has_pkg "$provider" || has_pkg nvidia-persistenced; then
    return 0
  fi
  apt-get install -y --no-install-recommends "$provider" >> "$STATE/apt-persist.txt" 2>&1
  rc=$?
  if [ "$rc" -eq 0 ]; then
    return 0
  fi
  apt-get install -y --no-install-recommends nvidia-persistenced >> "$STATE/apt-persist.txt" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "persistenced-not-restored" "persistenced" "$rc"
    return 1
  fi
  return 0
}

restore_previous_branch() {
  br=$(choose_restore_branch)
  write_fact restore_branch "$br"
  write_fact recovery restore
  spec=$(choose_restore_spec "$br")
  write_fact target_package "$spec"
  refresh_apt_lists || true
  if ! release_nvidia_holds; then
    note_step_fail "unhold-failed" "unhold" "1"
    return 1
  fi
  pkgs=$(unconfigured_nvidia_pkgs "$br")
  if [ -n "$pkgs" ]; then
    # shellcheck disable=SC2086
    apt-get remove --purge -y $pkgs > "$STATE/apt-purge.txt" 2>&1
    rc=$?
    if [ "$rc" -ne 0 ]; then
      note_step_fail "purge-failed" "restore-purge" "$rc"
      return 1
    fi
  fi
  # Download first so file conflicts apt-get -s cannot see (the EGL json
  # files inside libnvidia-gl) can be removed before the real install.
  # A download failure is recorded and does not skip the install: the cache
  # may already hold the debs.
  arch=$STATE/restore-archives
  rm -rf "$arch"
  mkdir -p "$arch"
  # shellcheck disable=SC2086
  apt-get install -y --download-only --no-install-recommends -o Dir::Cache::archives="$arch" $spec > "$STATE/apt-restore-download.txt" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "apt-download-failed" "download-only" "$rc"
  fi
  ALLOW_EGL_CONFLICT_REMOVAL=1
  saved_plan=$PLAN_PKGS
  PLAN_PKGS=$spec
  if ! collect_file_conflicts "$arch"; then
    note_step_fail "file-conflict" "file-conflict" "1"
  fi
  PLAN_PKGS=$saved_plan
  ALLOW_EGL_CONFLICT_REMOVAL=0
  # Remove the owners that are safe even if some other conflict was blocked.
  # The install below is still attempted so a file-conflict note is not the
  # last word when the debs can be unpacked.
  if ! remove_recorded_conflicts "$STATE/apt-restore-conflict-remove.txt"; then
    return 1
  fi
  # shellcheck disable=SC2086
  apt-get install -y --no-install-recommends $spec > "$STATE/apt-restore.txt" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "restore-install-failed" "restore-install" "$rc"
    return 1
  fi
  ensure_persistenced_pkg "$br" || return 1
  if ! dkms_built_for "$br" || ! module_ko_present; then
    dkms_autoinstall_running "$STATE/dkms-autoinstall.txt"
  fi
  if ! dkms_built_for "$br" || ! module_ko_present; then
    write_fact reboot_required no
    note_step_fail "dkms-kernel-mismatch" "dkms" "1"
    rehold_restore "$br" || true
    return 1
  fi
  rehold_restore "$br" || {
    note_step_fail "holds-not-restored" "rehold" "1"
    return 1
  }
  restart_services_strict || return 1
  if ! versions_match; then
    write_fact reboot_required yes
    note_step_fail "reboot-required" "versions" "1"
    return 1
  fi
  write_fact reboot_required no
  verify_restored "$br" || return 1
  return 0
}

finish_restored() {
  write_fact result repaired
  write_fact abort_reason ""
  write_fact failed_step ""
  write_fact failed_exit_code ""
  write_fact reboot_required no
  emit_report
  exit 0
}

recover_failed_install() {
  if [ "$MODE" = "apply" ]; then
    trap exit_rehold EXIT
  fi
  capture_diagnostics
  if [ "$MODE" != "apply" ]; then
    fail_after_change "apt-install-failed"
  fi
  ORIG_HOLDS=$(read_fact holds || true)
  ORIG_NVIDIA_HOLDS=$(read_fact nvidia_holds || true)
  if restore_previous_branch; then
    finish_restored
  fi
  exit_rehold
  start_vast || true
  reason=$(read_fact abort_reason || true)
  [ -n "$reason" ] || reason="restore-failed"
  abort "$reason"
}

run_repair() {
  trap exit_rehold EXIT
  strategy=${FILLRATE_REPAIR_STRATEGY:-forward}
  write_fact repair_strategy "$strategy"
  write_fact maint_window "${FILLRATE_NOT_BEFORE:-0}-${EXPIRES}"
  ORIG_HOLDS=$(read_fact holds || true)
  ORIG_NVIDIA_HOLDS=$(read_fact nvidia_holds || true)
  capture_diagnostics
  preflight
  if [ -n "$ORIG_HOLDS" ]; then
    write_fact holds "$ORIG_HOLDS"
  fi
  if [ -n "$ORIG_NVIDIA_HOLDS" ]; then
    write_fact nvidia_holds "$ORIG_NVIDIA_HOLDS"
  fi
  if [ "$strategy" = "forward" ]; then
    if ! repair_forward; then
      write_fact forward_abort_reason "$(read_fact abort_reason || true)"
      write_fact forward_failed_step "$(read_fact failed_step || true)"
      write_fact forward_exit_code "$(read_fact failed_exit_code || true)"
    fi
  fi
  if restore_previous_branch; then
    finish_restored
  fi
  exit_rehold
  reason=$(read_fact abort_reason || true)
  [ -n "$reason" ] || reason="restore-failed"
  abort "$reason"
}

pkg_available() {
  name=$1
  cand=$(apt-cache policy "$name" 2>/dev/null | awk '/Candidate:/ { print $2; exit }')
  [ -n "$cand" ] && [ "$cand" != "(none)" ]
}

kernel_source_name() {
  flavor=$(read_fact driver_flavor || true)
  if [ "$flavor" = "open" ]; then
    printf '%s\n' "nvidia-kernel-source-${BRANCH}-open"
  else
    printf '%s\n' "nvidia-kernel-source-${BRANCH}"
  fi
}

download_plan() {
  [ -n "${PLAN_PKGS:-}" ] || return 1
  dest=$STATE/archives
  mkdir -p "$dest"
  if [ "$PLAN_NO_RECOMMENDS" = "1" ]; then
    # shellcheck disable=SC2086
    apt-get install -y --download-only --no-install-recommends -o Dir::Cache::archives="$dest" $PLAN_PKGS > "$STATE/apt-download.txt" 2>&1
  else
    # shellcheck disable=SC2086
    apt-get install -y --download-only -o Dir::Cache::archives="$dest" $PLAN_PKGS > "$STATE/apt-download.txt" 2>&1
  fi
}

# cuda-*, container stacks, vast, and kernel images are never added as
# file-conflict removals. removal_forbidden covers the rest of the
# non-driver set (docker, containerd, libvirt, anything that is not nvidia).
conflict_protected() {
  pkg=$1
  case "$pkg" in
    cuda|cuda-*|libcudnn*|nvidia-cuda-toolkit|nvidia-cuda-toolkit-*|nvidia-container*|nvidia-docker*|docker|docker-*|containerd|containerd-*|vast|vast*|linux-image*|linux-headers*|linux-modules*|linux-kernel*)
      return 0
      ;;
  esac
  if removal_forbidden "$pkg"; then
    return 0
  fi
  return 1
}

plan_removes() {
  pkg=$1
  case " ${PLAN_PKGS:-} " in
    *" ${pkg}- "*) return 0 ;;
  esac
  remv=$(read_fact apt_remove || true)
  case " $remv " in
    *" $pkg "*) return 0 ;;
  esac
  return 1
}

replaces_has() {
  list=$1
  owner=$2
  printf '%s\n' "$list" | tr ',' '\n' | awk -v o="$owner" '
    {
      gsub(/^[ \t]+|[ \t]+$/, "", $0)
      name = $1
      sub(/:.*/, "", name)
      sub(/\(.*$/, "", name)
      gsub(/[ \t]+/, "", name)
      if (name == o) found = 1
    }
    END { exit !found }
  '
}

is_auto_pkg() {
  pkg=$1
  list=$2
  printf '%s\n' "$list" | awk -v p="$pkg" '
    {
      n = $1
      sub(/:.*/, "", n)
      if (n == p) found = 1
    }
    END { exit !found }
  '
}

# Auto-installed owners are removable. During repair/restore, the EGL platform
# packages are removable too when a downloaded deb ships their files. apt-get -s
# does not see that overlap: libnvidia-gl-580 Conflicts only libnvidia-egl-gbm1
# and libnvidia-gl, but it also ships 20_nvidia_xcb.json, 20_nvidia_xlib.json,
# and libnvidia-egl-xcb.so.1.0.5.
conflict_owner_allowed() {
  owner=$1
  auto_list=$2
  if conflict_protected "$owner"; then
    return 1
  fi
  if is_auto_pkg "$owner" "$auto_list"; then
    return 0
  fi
  if [ "${ALLOW_EGL_CONFLICT_REMOVAL:-0}" = "1" ]; then
    case "$owner" in
      libnvidia-egl-xcb1|libnvidia-egl-xlib1) return 0 ;;
    esac
  fi
  return 1
}

deb_file_paths() {
  awk '
    {
      mode = $1
      i = index($0, " ./")
      if (i < 1) next
      path = substr($0, i + 2)
      if (substr(mode, 1, 1) == "d") next
      if (path ~ /\/$/) next
      if (path == "" || path == "/") next
      print path
    }
  ' "$1"
}

record_deb_conflicts() {
  deb=$1
  newpkg=$(dpkg-deb -f "$deb" Package 2>/dev/null | awk 'NR==1 { sub(/:.*/, "", $1); print; exit }')
  [ -n "$newpkg" ] || return 0
  replaces=$(dpkg-deb -f "$deb" Replaces 2>/dev/null | tr '\n' ',' || true)
  dpkg-deb -c "$deb" > "$STATE/deb-contents.txt" 2>/dev/null || return 0
  deb_file_paths "$STATE/deb-contents.txt" > "$STATE/deb-files.txt"
  while IFS= read -r path; do
    [ -n "$path" ] || continue
    dpkg -S "$path" > "$STATE/dpkg-s.txt" 2>/dev/null || continue
    while IFS= read -r line; do
      [ -n "$line" ] || continue
      prefix=${line%"$path"}
      prefix=$(printf '%s' "$prefix" | sed 's/[[:space:]]*$//; s/:$//')
      [ -n "$prefix" ] || continue
      printf '%s\n' "$prefix" | tr ',' '\n' | while IFS= read -r raw; do
        owner=$(printf '%s' "$raw" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//; s/:.*//')
        [ -n "$owner" ] || continue
        [ "$owner" = "$newpkg" ] && continue
        if plan_removes "$owner"; then
          continue
        fi
        if replaces_has "$replaces" "$owner"; then
          continue
        fi
        printf '%s %s %s\n' "$owner" "$path" "$newpkg" >> "$STATE/file-conflicts.txt"
      done
    done < "$STATE/dpkg-s.txt"
  done < "$STATE/deb-files.txt"
}

scan_debs() {
  arch=$1
  [ -d "$arch" ] || return 0
  for deb in "$arch"/*.deb; do
    [ -f "$deb" ] || continue
    record_deb_conflicts "$deb"
  done
}

# Download directory first. The default apt cache is the fallback when
# download-only left the debs in the cache that was already populated.
conflict_scan_dirs() {
  arch=$1
  : > "$STATE/file-conflicts.txt"
  scan_debs "$arch"
  if [ ! -s "$STATE/file-conflicts.txt" ]; then
    have_deb=0
    for deb in "$arch"/*.deb; do
      [ -f "$deb" ] || continue
      have_deb=1
      break
    done
    if [ "$have_deb" = "0" ] && [ -d /var/cache/apt/archives ]; then
      scan_debs /var/cache/apt/archives
    fi
  fi
}

# Writes file_conflicts. On success, file_conflict_removals lists owners to
# remove. Returns 1 when a conflict cannot be removed safely.
collect_file_conflicts() {
  arch=$1
  conflict_scan_dirs "$arch"
  text=$(awk '{ printf "%s %s %s; ", $1, $2, $3 }' "$STATE/file-conflicts.txt")
  write_fact file_conflicts "$text"
  owners=$(awk 'NF && !seen[$1]++ { print $1 }' "$STATE/file-conflicts.txt")
  if [ -z "$owners" ]; then
    write_fact file_conflict_removals ""
    write_fact file_conflict_blocked ""
    return 0
  fi
  auto_list=$(apt-mark showauto 2>/dev/null || true)
  allowed=""
  blocked=""
  for owner in $owners; do
    if conflict_owner_allowed "$owner" "$auto_list"; then
      allowed="$allowed $owner"
    else
      blocked="$blocked $owner"
    fi
  done
  allowed=$(dedupe_list "$allowed")
  blocked=$(dedupe_list "$blocked")
  write_fact file_conflict_blocked "$blocked"
  # Keep the safe removals even when something else is blocked. Apply aborts
  # on that return and does not use them. Repair's restore still removes the
  # EGL packages that actually overlap the debs.
  write_fact file_conflict_removals "$allowed"
  if [ -n "$blocked" ]; then
    return 1
  fi
  return 0
}

attach_conflict_removals() {
  extras=$(read_fact file_conflict_removals || true)
  [ -n "$extras" ] || return 0
  for pkg in $extras; do
    case " $PLAN_PKGS " in
      *" ${pkg}- "*) ;;
      *)
        if [ -n "$PLAN_PKGS" ]; then
          PLAN_PKGS="$PLAN_PKGS ${pkg}-"
        else
          PLAN_PKGS="${pkg}-"
        fi
        ;;
    esac
  done
  remv=$(read_fact apt_remove || true)
  for pkg in $extras; do
    case " $remv " in
      *" $pkg "*) ;;
      *) remv="$remv $pkg" ;;
    esac
  done
  write_fact apt_remove "$(dedupe_list "$remv")"
}

# Before any unhold or removal: scan the downloaded debs. Safe auto-installed
# owners are added and the plan is judged again. Anything else aborts.
resolve_downloaded_conflicts() {
  if ! collect_file_conflicts "$STATE/archives"; then
    abort "file-conflict"
  fi
  extras=$(read_fact file_conflict_removals || true)
  [ -n "$extras" ] || return 0
  if [ ! -f "$STATE/winner-sim.txt" ]; then
    abort "apt-sim-failed"
  fi
  for pkg in $extras; do
    printf 'Remv %s\n' "$pkg" >> "$STATE/winner-sim.txt"
  done
  judge_plan || return 1
  attach_conflict_removals
}

installed_pkg_version() {
  name=$1
  dpkg -l 2>/dev/null | awk -v name="$name" '
    $1 ~ /^(ii|hi|iU|iF|iH)$/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      if (pkg == name) { print $3; exit }
    }
  '
}

dkms_built_version() {
  ver=$1
  kernel=$(uname -r 2>/dev/null || true)
  [ -n "$kernel" ] || return 1
  [ -n "$ver" ] || return 1
  dkms status 2>/dev/null | awk -v v="$ver" -v k="$kernel" '
    index($0, "nvidia/" v) && index($0, k) && $0 ~ /installed/ { found = 1 }
    END { exit !found }
  '
}

rehold_forward() {
  meta=$1
  pkgs=$(dpkg -l 2>/dev/null | awk -v br="$BRANCH" '
    $1 ~ /^(ii|hi|iU|iF)$/ {
      pkg = $2
      sub(/:.*/, "", pkg)
      n = pkg
      sub(/-open$/, "", n)
      if (n ~ ("-" br "$")) print pkg
    }
  ' | sort -u)
  base="containerd.io docker-ce docker-ce-cli docker-ce-rootless-extras libvirt-dev nvidia-docker2 xfsprogs"
  list=$meta
  for p in $pkgs; do
    [ "$p" = "$meta" ] && continue
    list="$list $p"
  done
  list="$list $base"
  # shellcheck disable=SC2086
  apt-mark hold $list >/dev/null 2>&1 || return 1
  HOLDS_RESTORED=1
  return 0
}

# Finish the half-installed new branch and reboot so the new module loads.
# Returns 1 when that cannot be done; the caller restores without rebooting.
repair_forward() {
  write_fact recovery forward
  write_fact forward_branch "$BRANCH"
  flavor=$(read_fact driver_flavor || true)
  if [ "$flavor" = "open" ]; then
    meta="nvidia-driver-${BRANCH}-open"
  else
    meta="nvidia-driver-${BRANCH}"
  fi
  full=$(installed_pkg_version "$meta")
  if [ -z "$full" ]; then
    full=$(apt-cache policy "$meta" 2>/dev/null | awk '/Candidate:/ { print $2; exit }')
    if [ "$full" = "(none)" ]; then
      full=""
    fi
  fi
  upstream=""
  if [ -n "${BUILD:-}" ]; then
    upstream=$BUILD
  elif [ -n "$full" ]; then
    upstream=$(printf '%s\n' "$full" | awk -F- 'NR==1 { print $1 }')
  fi
  if [ -n "$upstream" ]; then
    write_fact resolved_build "$upstream"
    BUILD=$upstream
  fi
  pin=""
  if [ -n "$full" ] && [ -n "$upstream" ]; then
    case "$full" in
      "$upstream"|"$upstream"-*) pin=$full ;;
    esac
  fi
  if [ -z "$pin" ] && [ -n "$upstream" ]; then
    pin=$upstream
  fi
  if [ -n "$pin" ]; then
    spec="${meta}=${pin}"
    gl="libnvidia-gl-${BRANCH}=${pin}"
  else
    spec=$meta
    gl="libnvidia-gl-${BRANCH}"
  fi
  write_fact target_package "$spec"
  refresh_apt_lists || true
  if ! release_nvidia_holds; then
    note_step_fail "unhold-failed" "unhold" "1"
    return 1
  fi
  arch=$STATE/archives
  rm -rf "$arch"
  mkdir -p "$arch"
  apt-get install -y --download-only --no-install-recommends -o Dir::Cache::archives="$arch" "$spec" "$gl" > "$STATE/apt-download.txt" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "apt-download-failed" "download-only" "$rc"
    return 1
  fi
  saved_plan=$PLAN_PKGS
  PLAN_PKGS="$spec $gl"
  ALLOW_EGL_CONFLICT_REMOVAL=1
  if ! collect_file_conflicts "$arch"; then
    PLAN_PKGS=$saved_plan
    ALLOW_EGL_CONFLICT_REMOVAL=0
    note_step_fail "file-conflict" "file-conflict" "1"
    return 1
  fi
  PLAN_PKGS=$saved_plan
  ALLOW_EGL_CONFLICT_REMOVAL=0
  if ! remove_recorded_conflicts "$STATE/apt-conflict-remove.txt"; then
    return 1
  fi
  apt-get install -y --no-install-recommends -f > "$STATE/apt-install-f.txt" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "apt-install-failed" "install -f" "$rc"
    return 1
  fi
  apt-get install -y --no-install-recommends "$gl" "$spec" > "$STATE/apt-install.txt" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "apt-install-failed" "install" "$rc"
    return 1
  fi
  dpkg --configure -a > "$STATE/dpkg-configure.txt" 2>&1
  rc=$?
  if [ "$rc" -ne 0 ]; then
    note_step_fail "dpkg-configure-failed" "configure" "$rc"
    return 1
  fi
  capture_diagnostics
  if [ -n "$(half_installed_branch || true)" ]; then
    note_step_fail "dpkg-configure-failed" "configure" "1"
    return 1
  fi
  ver=$upstream
  [ -n "$ver" ] || ver=$BRANCH
  if ! dkms_built_version "$ver" || ! module_ko_present; then
    dkms_autoinstall_running "$STATE/dkms-autoinstall.txt"
  fi
  if ! dkms_built_version "$ver"; then
    note_step_fail "dkms-kernel-mismatch" "dkms" "1"
    return 1
  fi
  if ! module_ko_present; then
    note_step_fail "module-not-on-disk" "dkms" "1"
    return 1
  fi
  provider="nvidia-compute-utils-${BRANCH}"
  if ! has_pkg "$provider"; then
    apt-get install -y --no-install-recommends "$provider" >> "$STATE/apt-persist.txt" 2>&1
    rc=$?
    if [ "$rc" -ne 0 ]; then
      note_step_fail "persistenced-not-restored" "persistenced" "$rc"
      return 1
    fi
  fi
  if command -v systemctl >/dev/null 2>&1; then
    systemctl enable --now nvidia-persistenced >/dev/null 2>&1 || {
      note_step_fail "persistenced-not-restored" "persistenced" "1"
      return 1
    }
    systemctl is-enabled --quiet nvidia-persistenced >/dev/null 2>&1 || {
      note_step_fail "persistenced-not-restored" "persistenced" "1"
      return 1
    }
  fi
  rehold_forward "$meta" || {
    note_step_fail "holds-not-restored" "rehold" "1"
    return 1
  }
  write_fact reboot_required yes
  write_fact maint_window "${FILLRATE_NOT_BEFORE:-0}-${EXPIRES}"
  touch "$STATE/pending-post"
  install_post_unit || {
    note_step_fail "post-unit-failed" "post-unit" "1"
    return 1
  }
  write_fact result reboot-scheduled
  write_fact abort_reason ""
  write_fact failed_step ""
  write_fact failed_exit_code ""
  emit_report
  reboot || {
    write_fact result abort
    note_step_fail "reboot-failed" "reboot" "1"
    emit_report
    exit 2
  }
  exit 0
}

pre_removal_check() {
  k=$(uname -r 2>/dev/null || true)
  [ -n "$k" ] || abort "kernel-unknown"
  hdr="linux-headers-${k}"
  if ! has_pkg "$hdr"; then
    if ! pkg_available "$hdr"; then
      abort "linux-headers-missing"
    fi
    apt-get -s install "$hdr" > "$STATE/headers-sim.txt" 2>&1 || true
    bad=$(first_disallowed_removal "$STATE/headers-sim.txt" || true)
    if [ -n "$bad" ]; then
      abort "linux-headers-conflict"
    fi
    apt-get install -y "$hdr" > "$STATE/headers-install.txt" 2>&1 || abort "linux-headers-missing"
  fi
  src=$(kernel_source_name)
  if ! has_pkg "$src" && ! pkg_available "$src"; then
    abort "kernel-source-unavailable"
  fi
  download_plan || abort "apt-download-failed"
  resolve_downloaded_conflicts
}

if [ "$PHASE" = "post" ]; then
  post_checks
fi

if [ "$MODE" = "repair" ]; then
  run_repair
fi

saved_prev=$(read_fact previous_branch || true)
saved_prev_build=$(read_fact previous_build || true)
saved_target=$(read_fact target_branch || true)

preflight

if [ "$MODE" = "rollback" ]; then
  if [ -z "$saved_prev" ] || [ "$saved_target" != "$BRANCH" ]; then
    abort "no-rollback-record"
  fi
  write_fact previous_branch "$saved_prev"
  write_fact previous_build "$saved_prev_build"
  spec=$(meta_for "$saved_prev")
  BUILD=$saved_prev_build
else
  spec=$(meta_for "$BRANCH")
fi

write_fact target_package "$spec"
PKG_SPEC=$spec
PKG=
run_apt_update || true
resolve_build "$spec" || true
simulate "$PKG" "$spec" || true

if [ "$MODE" = "dry-run" ]; then
  if [ "$SIM_FAILED" = "1" ]; then
    # Disk, DKMS, Secure Boot, Docker, and Vast are still recorded.
    preflight_late 0
    abort "$SIM_REASON"
  fi
  preflight_late 1
  write_fact result dry-run
  write_fact abort_reason ""
  emit_report
  exit 0
fi

if [ "$SIM_FAILED" = "1" ]; then
  abort "${SIM_REASON:-apt-sim-failed}"
fi
preflight_late 1

if [ "$MODE" = "apply" ]; then
  pre_removal_check
fi

write_fact target_branch "$BRANCH"
if command -v systemctl >/dev/null 2>&1; then
  systemctl stop vastai.service >/dev/null 2>&1 || abort "vast-daemon-stop-failed"
fi
holds=$(read_fact nvidia_holds || true)
if [ -n "$holds" ]; then
  # shellcheck disable=SC2086
  apt-mark unhold $holds >/dev/null 2>&1 || fail_after_change "unhold-failed"
fi
confirm_plan
if [ "$MODE" = "apply" ]; then
  attach_conflict_removals
  download_plan || fail_after_change "apt-download-failed"
fi
real_install
touch "$STATE/pending-post"
install_post_unit || fail_after_change "post-unit-failed"
write_fact result reboot-scheduled
write_fact abort_reason ""
emit_report
reboot || fail_after_change "reboot-failed"
exit 0