#!/bin/bash
#
# Acceptance test harness: boots an Omarchy ISO in a headless QEMU VM and
# drives the real interactive install flow the way a user would — reading
# each screen through QMP screendumps + OCR, and answering with virtual
# keystrokes. Nothing in the product is bypassed or stubbed: the configurator
# wizard, install dashboard, reboot prompt, and SDDM login are all exercised.
#
# After first boot it logs into a console TTY (typed, like everything else)
# to authorize an SSH key, snapshots the installed disk, then boots a
# throwaway overlay and runs the in-guest acceptance suite (test/acceptance
# in the omarchy repo) over SSH.
#
# Usage: omarchy-iso-test [release/omarchy.iso] [options]
#
#   --encrypt          Drive the encrypted install flow (default: unencrypted)
#   --provision              Drive the deferred-provisioning install flow: no user during install, the
#                      first boot runs omarchy-provision-owner and creates it there
#   --reuse-base       Skip the install phase if a base image already exists
#   --install-only     Stop after producing the installed base image
#   --sync-omarchy DIR Push DIR's test/acceptance into the guest before running
#                      (defaults to OMARCHY_PATH when the suite is present)
#   --sync-all DIR     Push all of DIR (bin/config/shell/test) into the guest
#   --port PORT        Host port forwarded to guest SSH (default 2222)
#   --memory MB        VM memory (default 8192)
#   --timeout SECS     Install timeout (default 2400)
#   --keep-running     Leave the acceptance VM running after the suite finishes
#   --no-preview       Do not open the collected screenshots in imv

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_DIR="$(dirname "$SCRIPT_DIR")"

ISO=""
SSH_PORT=2222
MEMORY=8192
INSTALL_TIMEOUT=2400
BOOT_TIMEOUT=600
ENCRYPT=false
PROVISION=false
REUSE_BASE=false
INSTALL_ONLY=false
KEEP_RUNNING=false
NO_PREVIEW=false
SYNC_DIR=""
SYNC_ALL=false

GUEST_USER="omarchy"
GUEST_PASSWORD="omarchy"
GUEST_HOSTNAME="omarchy-test"

while (($#)); do
  case "$1" in
    --encrypt) ENCRYPT=true ;;
    --provision) PROVISION=true ;;
    --reuse-base) REUSE_BASE=true ;;
    --install-only) INSTALL_ONLY=true ;;
    --keep-running) KEEP_RUNNING=true ;;
    --no-preview) NO_PREVIEW=true ;;
    --sync-omarchy) SYNC_DIR="$2"; shift ;;
    --sync-all) SYNC_DIR="$2"; SYNC_ALL=true; shift ;;
    --port) SSH_PORT="$2"; shift ;;
    --memory) MEMORY="$2"; shift ;;
    --timeout) INSTALL_TIMEOUT="$2"; shift ;;
    -*) echo "Unknown option: $1" >&2; exit 1 ;;
    *) ISO="$1" ;;
  esac
  shift
done

if [[ -z $SYNC_DIR && -f ${OMARCHY_PATH:-}/test/acceptance ]]; then
  SYNC_DIR="$OMARCHY_PATH"
fi

if [[ -z $ISO && -t 0 ]]; then
  ISO=$(ls -t "$REPO_DIR"/release/*.iso 2>/dev/null | gum choose --header "Select ISO to test")
fi

if [[ ! -f ${ISO:-} ]]; then
  echo "No ISO to test. Pass one explicitly or put it in release/." >&2
  echo "Usage: omarchy-iso-test [release/omarchy.iso] [options]" >&2
  exit 1
fi

omarchy-pkg-add qemu-full edk2-ovmf socat imagemagick tesseract tesseract-data-eng

OVMF_CODE="/usr/share/edk2/x64/OVMF_CODE.4m.fd"
OVMF_VARS_TEMPLATE="/usr/share/edk2/x64/OVMF_VARS.4m.fd"

BASE_NAME="$(basename "$ISO" .iso)"
$ENCRYPT && BASE_NAME+="-encrypted"
$PROVISION && BASE_NAME+="-provision"
BASE_DIR="$REPO_DIR/test-runs/$BASE_NAME"
RUN_DIR="$BASE_DIR/runs/$(date +%Y%m%d-%H%M%S)"
BASE_DISK="$BASE_DIR/base.qcow2"
BASE_OVMF="$BASE_DIR/OVMF_VARS.4m.fd"
SSH_KEY="$BASE_DIR/id_ed25519"
HTTP_PORT=$((SSH_PORT + 1))
HTTP_PID=""

mkdir -p "$BASE_DIR" "$RUN_DIR"

# Unix socket paths are capped at ~108 bytes, which a run directory nested in
# test-runs/<iso-name>/runs/<timestamp> can exceed. Keep the socket in /tmp.
QMP_SOCK=$(mktemp -u "${TMPDIR:-/tmp}/omarchy-iso-test-qmp.XXXXXX.sock")
PIDFILE="$RUN_DIR/qemu.pid"

log() {
  printf '\033[1;35m==> %s\033[0m\n' "$1"
}

# ---------------------------------------------------------------- vm control

vm_running() {
  [[ -f $PIDFILE ]] && kill -0 "$(cat "$PIDFILE")" 2>/dev/null
}

qmp() {
  printf '{"execute":"qmp_capabilities"}\n{"execute":%s}\n' "$1" |
    timeout 5 socat -t 2 - "UNIX-CONNECT:$QMP_SOCK" 2>/dev/null || true
}

screendump() {
  qmp "\"screendump\", \"arguments\": {\"filename\": \"$1\"}" >/dev/null
}

capture_console() {
  local name="$1" shot="$RUN_DIR/.capture.ppm"

  sleep 1
  screendump "$shot"
  [[ -s $shot ]] || return 0
  magick "$shot" "$RUN_DIR/$name.png" 2>/dev/null || true
  rm -f "$shot"
}

stop_vm() {
  vm_running || return 0

  if ! ssh_guest "echo $GUEST_PASSWORD | sudo -S systemctl poweroff" >/dev/null 2>&1; then
    qmp '"system_powerdown"' >/dev/null
  fi

  local waited=0
  while vm_running && ((waited < 15)); do
    sleep 1
    ((waited += 1))
  done

  if vm_running; then
    qmp '"quit"' >/dev/null
    sleep 1
  fi

  if vm_running; then
    kill "$(cat "$PIDFILE")" 2>/dev/null || true
    sleep 1
  fi
}

open_screenshots() {
  local entry
  local -a screenshots=()

  while IFS= read -r -d '' entry; do
    screenshots+=("${entry#* }")
  done < <(
    find "$RUN_DIR" -type f \( -name "success-*.png" -o -name "failure-*.png" \) -printf '%T@ %p\0' |
      sort -zn
  )

  (( ${#screenshots[@]} > 0 )) || return 0

  log "Visual review: ${#screenshots[@]} screenshots in $RUN_DIR"
  $NO_PREVIEW && return 0

  if command -v imv >/dev/null && [[ -n ${WAYLAND_DISPLAY:-}${DISPLAY:-} ]]; then
    setsid -f imv "${screenshots[@]}" >/dev/null 2>&1
  fi
}

cleanup() {
  local status=$?

  [[ -n $HTTP_PID ]] && kill "$HTTP_PID" 2>/dev/null || true

  if ! $KEEP_RUNNING; then
    vm_running && log "Stopping test VM"
    stop_vm
  fi

  rm -f "$RUN_DIR/.screen.ppm" "$RUN_DIR/.screen.png"
  open_screenshots

  return $status
}
trap cleanup EXIT

start_vm() {
  local disk="$1" serial="$2"
  shift 2

  qemu-system-x86_64 \
    -cpu host -enable-kvm -machine q35,accel=kvm \
    -smp "$(nproc)" \
    -m "$MEMORY" \
    -drive if=pflash,format=raw,readonly=on,file="$OVMF_CODE" \
    -drive if=pflash,format=raw,file="$BASE_OVMF" \
    -drive file="$disk",format=qcow2,if=none,id=drive0 \
    -device virtio-blk-pci,drive=drive0,bootindex=1 \
    -device virtio-vga \
    -display none \
    -usb -device usb-tablet \
    -netdev user,id=net0,hostfwd=tcp:127.0.0.1:$SSH_PORT-:22 \
    -device virtio-net-pci,netdev=net0 \
    -qmp "unix:$QMP_SOCK,server,nowait" \
    -serial "file:$serial" \
    -pidfile "$PIDFILE" \
    -daemonize \
    "$@"
}

# ------------------------------------------------------------ console driver
#
# The heart of the harness: read the VM's console with screendump + OCR, and
# type into it with QMP send-key. Together they make an expect-style loop
# that walks the same screens a user sees.

ocr_screen() {
  local shot="$RUN_DIR/.screen.ppm" prepped="$RUN_DIR/.screen.png"

  screendump "$shot"
  [[ -s $shot ]] || return 0
  magick "$shot" -colorspace gray -negate -resize 150% "$prepped" 2>/dev/null || return 0
  tesseract "$prepped" - --psm 6 2>/dev/null || true
}

# Wait until the console shows the given text. Saves the offending screen as
# an artifact on timeout.
wait_for_screen() {
  local text="$1" timeout="$2" waited=0 slug

  until ocr_screen | grep -qi "$text"; do
    if ! vm_running; then
      echo "VM exited while waiting for screen: $text" >&2
      return 1
    fi

    if ((waited >= timeout)); then
      slug=${text,,}
      slug=${slug// /-}
      slug=${slug//[^a-z0-9-]/}
      capture_console "failure-installer-waiting-for-$slug"
      echo "Timed out after ${timeout}s waiting for screen: $text" >&2
      return 1
    fi

    sleep 3
    ((waited += 3))
  done

  sleep 1 # let the screen settle before typing
}

# Press a key or chord, dash-separated: "ret", "ctrl-c", "ctrl-alt-f3".
press() {
  local part json=""
  local -a parts

  IFS='-' read -ra parts <<<"$1"
  for part in "${parts[@]}"; do
    json+="{\"type\":\"qcode\",\"data\":\"$part\"},"
  done

  qmp "\"send-key\", \"arguments\": {\"keys\": [${json%,}]}" >/dev/null
}

type_text() {
  local text="$1" ch i

  for ((i = 0; i < ${#text}; i++)); do
    ch=${text:i:1}
    case "$ch" in
      [a-z0-9]) press "$ch" ;;
      [A-Z]) press "shift-${ch,,}" ;;
      " ") press spc ;;
      .) press dot ;;
      ,) press comma ;;
      -) press minus ;;
      _) press shift-minus ;;
      /) press slash ;;
      :) press shift-semicolon ;;
      @) press shift-2 ;;
      \|) press shift-backslash ;;
      ";") press semicolon ;;
      "'") press apostrophe ;;
      '"') press shift-apostrophe ;;
      "=") press equal ;;
      "+") press shift-equal ;;
      "%") press shift-5 ;;
      "&") press shift-7 ;;
      "*") press shift-8 ;;
      "(") press shift-9 ;;
      ")") press shift-0 ;;
      "<") press shift-comma ;;
      ">") press shift-dot ;;
      "!") press shift-1 ;;
      "#") press shift-3 ;;
      '$') press shift-4 ;;
      "?") press shift-slash ;;
      "~") press shift-grave_accent ;;
      "\\") press backslash ;;
      "[") press bracket_left ;;
      "]") press bracket_right ;;
      "{") press shift-bracket_left ;;
      "}") press shift-bracket_right ;;
      *) echo "type_text: unsupported character: $ch" >&2; return 1 ;;
    esac
    sleep 0.05
  done
}

# --------------------------------------------------------------------- guest

ssh_guest() {
  ssh -i "$SSH_KEY" -p "$SSH_PORT" \
    -o BatchMode=yes \
    -o IdentitiesOnly=yes \
    -o StrictHostKeyChecking=no \
    -o UserKnownHostsFile=/dev/null \
    -o ConnectTimeout=5 \
    -o LogLevel=ERROR \
    "$GUEST_USER@127.0.0.1" "$@"
}

wait_for_ssh() {
  local timeout="$1" failure_name="${2:-failure-session-ssh-timeout}" waited=0

  while ! ssh_guest true 2>/dev/null; do
    if ! vm_running; then
      echo "VM exited while waiting for SSH" >&2
      return 1
    fi

    if ((waited >= timeout)); then
      capture_console "$failure_name"
      echo "Timed out after ${timeout}s waiting for SSH" >&2
      return 1
    fi

    sleep 5
    ((waited += 5))
  done
}

session_started() {
  ssh_guest "ls /run/user/1000/hypr 2>/dev/null | grep -q ." 2>/dev/null
}

ssh_session() {
  local command="$1"

  ssh_guest "export XDG_RUNTIME_DIR=/run/user/\$(id -u); \
    export DBUS_SESSION_BUS_ADDRESS=unix:path=\$XDG_RUNTIME_DIR/bus; \
    export LANG=\$(systemctl --user show-environment | sed -n 's/^LANG=//p' | head -1); \
    export LANG=\${LANG:-C.UTF-8}; \
    export HYPRLAND_INSTANCE_SIGNATURE=\$(ls -t \$XDG_RUNTIME_DIR/hypr | head -1); \
    export WAYLAND_DISPLAY=\$(find \$XDG_RUNTIME_DIR -maxdepth 1 -name 'wayland-*' ! -name '*.lock' -printf '%f\\n' | head -1); \
    export OMARCHY_PATH=/usr/share/omarchy; \
    export PATH=\$OMARCHY_PATH/bin:\$PATH; \
    $command" </dev/null
}

wait_for_guest_state() {
  local description="$1" timeout="$2"
  shift 2
  local deadline=$((SECONDS + timeout))

  until "$@" >/dev/null 2>&1; do
    if ((SECONDS >= deadline)); then
      printf 'not ok - %s\n' "$description" >&2
      return 1
    fi
    sleep 1
  done

  printf 'ok - %s\n' "$description"
}

guest_layer_present() {
  local namespace="$1"

  ssh_session "hyprctl -j layers | jq -e --arg namespace '$namespace' \
    '[.. | objects | select(.namespace? == \$namespace)] | length > 0'"
}

guest_layer_absent() {
  ! guest_layer_present "$1"
}

cleanup_shortcut_state() {
  ssh_session 'for plugin in omarchy.launcher omarchy.menu omarchy.emojis omarchy.clipboard omarchy.audio; do
    omarchy-shell shell hide "$plugin" >/dev/null 2>&1 || true
  done
  while read -r address; do
    hyprctl dispatch "hl.dsp.window.close({ window = \"address:$address\" })" >/dev/null 2>&1 ||
      hyprctl dispatch closewindow "address:$address" >/dev/null 2>&1 || true
  done < <(hyprctl -j clients | jq -r '\''.[] | select(.class == "foot") | .address'\'')
  hyprctl dispatch workspace 1 >/dev/null 2>&1 || true' >/dev/null 2>&1 || true
}

verify_surface_shortcut() {
  local name="$1" chord="$2" namespace="$3" plugin="$4"

  press "$chord"
  if ! wait_for_guest_state "$name shortcut opens" 15 guest_layer_present "$namespace"; then
    capture_console "failure-shortcut-$name-open"
    ssh_session "omarchy-shell shell hide '$plugin' >/dev/null 2>&1 || true" || true
    return 1
  fi
  capture_console "success-shortcut-$name"

  press esc
  if ! wait_for_guest_state "$name shortcut surface closes" 15 guest_layer_absent "$namespace"; then
    capture_console "failure-shortcut-$name-close"
    ssh_session "omarchy-shell shell hide '$plugin' >/dev/null 2>&1 || true" || true
    return 1
  fi
}

verify_desktop_shortcuts() {
  cleanup_shortcut_state

  press meta_l-ret
  if ! wait_for_guest_state "terminal shortcut opens a window" 30 \
    ssh_session "hyprctl -j clients | jq -e '[.[] | select(.class == \"foot\")] | length == 1'"; then
    capture_console "failure-shortcut-terminal-open"
    return 1
  fi

  press meta_l-ret
  if ! wait_for_guest_state "terminal shortcut opens a second window" 30 \
    ssh_session "hyprctl -j clients | jq -e '[.[] | select(.class == \"foot\")] | length == 2'"; then
    capture_console "failure-shortcut-terminal-second-window"
    return 1
  fi
  capture_console "success-shortcut-desktop-01-two-tiled-terminals"

  press meta_l-t
  if ! wait_for_guest_state "floating shortcut floats the active window" 15 \
    ssh_session "hyprctl -j activewindow | jq -e '.class == \"foot\" and .floating == true'"; then
    capture_console "failure-shortcut-window-float"
    return 1
  fi
  capture_console "success-shortcut-desktop-02-floating-terminal"

  press meta_l-t
  if ! wait_for_guest_state "floating shortcut returns the window to tiling" 15 \
    ssh_session "hyprctl -j activewindow | jq -e '.class == \"foot\" and .floating == false'"; then
    capture_console "failure-shortcut-window-retile"
    return 1
  fi

  press meta_l-shift-2
  if ! wait_for_guest_state "move shortcut follows the window to workspace 2" 15 \
    ssh_session "hyprctl -j activeworkspace | jq -e '.id == 2' && \
      hyprctl -j activewindow | jq -e '.class == \"foot\" and .workspace.id == 2'"; then
    capture_console "failure-shortcut-window-workspace-2"
    return 1
  fi
  capture_console "success-shortcut-desktop-03-terminal-on-workspace-2"

  press meta_l-1
  if ! wait_for_guest_state "workspace shortcut returns to workspace 1" 15 \
    ssh_session "hyprctl -j activeworkspace | jq -e '.id == 1'"; then
    capture_console "failure-shortcut-workspace-1"
    return 1
  fi
  capture_console "success-shortcut-desktop-04-workspace-1"

  press meta_l-w
  if ! wait_for_guest_state "close shortcut closes the workspace 1 terminal" 15 \
    ssh_session "hyprctl -j clients | jq -e '[.[] | select(.class == \"foot\")] | length == 1'"; then
    capture_console "failure-shortcut-close-workspace-1"
    return 1
  fi

  press meta_l-2
  if ! wait_for_guest_state "workspace shortcut visits workspace 2" 15 \
    ssh_session "hyprctl -j activeworkspace | jq -e '.id == 2'"; then
    capture_console "failure-shortcut-visit-workspace-2"
    return 1
  fi
  press meta_l-w
  if ! wait_for_guest_state "close shortcut closes the workspace 2 terminal" 15 \
    ssh_session "hyprctl -j clients | jq -e '[.[] | select(.class == \"foot\")] | length == 0'"; then
    capture_console "failure-shortcut-close-workspace-2"
    return 1
  fi
}

verify_universal_clipboard_shortcuts() {
  local copy_token="OmarchyUniversalCopy$(date +%s)"
  local paste_token="OmarchyUniversalPaste$(date +%s)"
  local terminal_token="OmarchyTerminalPaste$(date +%s)"
  local terminal_result="/tmp/omarchy-universal-terminal-paste"

  cleanup_shortcut_state
  ssh_session "setsid -f chromium --new-window about:blank >/dev/null 2>&1"
  if ! wait_for_guest_state "universal clipboard test opens Chromium" 45 \
    ssh_session "hyprctl -j clients | jq -e '[.[] | select(.class | test(\"(?i)chromium\"))] | length == 1'"; then
    capture_console "failure-shortcut-universal-clipboard-browser"
    return 1
  fi

  press ctrl-l
  type_text "$copy_token"
  press ctrl-a
  press meta_l-c
  if ! wait_for_guest_state "universal copy uses the active GUI application" 15 \
    ssh_session "[[ \$(wl-paste --no-newline) == '$copy_token' ]]"; then
    capture_console "failure-shortcut-universal-copy"
    return 1
  fi

  ssh_session "setsid -f bash -c 'printf \"%s\" \"$paste_token\" | wl-copy' >/dev/null 2>&1"
  if ! wait_for_guest_state "universal paste clipboard is seeded" 15 \
    ssh_session "[[ \$(wl-paste --no-newline) == '$paste_token' ]]"; then
    capture_console "failure-shortcut-universal-paste-seed"
    return 1
  fi
  press ctrl-l
  press meta_l-v
  sleep 1
  press ctrl-a
  press ctrl-c
  if ! wait_for_guest_state "universal paste uses the active GUI application" 15 \
    ssh_session "[[ \$(wl-paste --no-newline) == '$paste_token' ]]"; then
    capture_console "failure-shortcut-universal-paste-gui"
    return 1
  fi
  capture_console "success-shortcut-universal-copy-paste-gui"
  ssh_session 'while read -r address; do
    hyprctl dispatch "hl.dsp.window.close({ window = \"address:$address\" })" >/dev/null 2>&1 ||
      hyprctl dispatch closewindow "address:$address" >/dev/null 2>&1 || true
  done < <(hyprctl -j clients | jq -r '\''.[] | select(.class | test("(?i)chromium")) | .address'\'')' || true

  ssh_session "rm -f '$terminal_result'; \
    setsid -f bash -c 'printf \"%s\" \"$terminal_token\" | wl-copy' >/dev/null 2>&1; \
    setsid -f foot --app-id=foot -e bash -c 'read -r value; printf \"%s\" \"\$value\" > \"$terminal_result\"; sleep 30' >/dev/null 2>&1"
  if ! wait_for_guest_state "terminal paste clipboard is seeded" 15 \
    ssh_session "[[ \$(wl-paste --no-newline) == '$terminal_token' ]]"; then
    capture_console "failure-shortcut-universal-paste-terminal-seed"
    return 1
  fi
  if ! wait_for_guest_state "universal clipboard test opens a terminal" 30 \
    ssh_session "hyprctl -j clients | jq -e '[.[] | select(.class == \"foot\")] | length == 1'"; then
    capture_console "failure-shortcut-universal-clipboard-terminal"
    return 1
  fi

  press meta_l-v
  sleep 1
  press ret
  if ! wait_for_guest_state "universal paste uses the terminal-safe shortcut" 15 \
    ssh_session "[[ \$(cat '$terminal_result' 2>/dev/null) == '$terminal_token' ]]"; then
    capture_console "failure-shortcut-universal-paste-terminal"
    return 1
  fi
  capture_console "success-shortcut-universal-paste-terminal"
  cleanup_shortcut_state
}

shortcut_smoke_phase() {
  log "Exercising primary desktop shortcuts through the VM keyboard"

  local status=0
  local name chord namespace plugin
  local shortcuts='launcher|meta_l-spc|omarchy-launcher|omarchy.launcher
menu|meta_l-alt-spc|omarchy-menu|omarchy.menu
emojis|meta_l-ctrl-e|omarchy-emojis|omarchy.emojis
clipboard|meta_l-ctrl-v|omarchy-clipboard|omarchy.clipboard
audio-panel|meta_l-ctrl-a|omarchy-keyboard-panel|omarchy.audio'

  cleanup_shortcut_state
  while IFS='|' read -r name chord namespace plugin; do
    if ! verify_surface_shortcut "$name" "$chord" "$namespace" "$plugin"; then
      status=1
    fi
  done <<<"$shortcuts"

  if ! verify_desktop_shortcuts; then
    status=1
  fi
  if ! verify_universal_clipboard_shortcuts; then
    status=1
  fi
  cleanup_shortcut_state

  return $status
}

# ----------------------------------------------------- driving the installer

drive_configurator() {
  log "Driving the installer wizard"

  # A greeter (animated logo + tagline) opens the installer; Return begins.
  wait_for_screen "Opinionated" 300
  capture_console "success-installer-00-greeter"
  press ret

  # Keyboard is the first screen. Ctrl+C there is the hidden entry into deferred provisioning
  # mode; a normal install just selects the layout.
  wait_for_screen "keyboard layout" 300
  capture_console "success-installer-01-keyboard-layout"

  if $PROVISION; then
    press ctrl-c # arm "prepare for another owner"
    wait_for_screen "another owner" 30
    capture_console "success-installer-02-prepare-confirm"
    press ret # "Yes, prepare for another owner" is the affirmative default

    # Jumps straight to disk selection, then the overwrite confirm.
    wait_for_screen "install disk" 60
    capture_console "success-installer-03-prepare-install-disk"
    press ret

    wait_for_screen "recovery possible" 60
    capture_console "success-installer-04-prepare-disk-warning-encrypted"
    if ! $ENCRYPT; then
      press ctrl-c
      wait_for_screen "without encryption" 30
      capture_console "success-installer-05-prepare-disk-warning-unencrypted"
    fi
    press ret
    sleep 2
    capture_console "success-installer-13-install-started"
    return 0
  fi

  press ret # English (US) is preselected

  # Normal flow: user step comes before disk selection.
  wait_for_screen "Username" 120
  type_text "$GUEST_USER"
  capture_console "success-installer-02-username"
  press ret

  wait_for_screen "Password" 60
  type_text "$GUEST_PASSWORD"
  capture_console "success-installer-03-password"
  press ret

  wait_for_screen "Must match" 60
  type_text "$GUEST_PASSWORD"
  capture_console "success-installer-04-password-confirmation"
  press ret

  wait_for_screen "Full name" 60
  type_text "Omarchy Test"
  capture_console "success-installer-05-full-name"
  press ret

  wait_for_screen "Email address" 60
  type_text "test@omarchy.org"
  capture_console "success-installer-06-email"
  press ret

  wait_for_screen "Hostname" 60
  type_text "$GUEST_HOSTNAME"
  capture_console "success-installer-07-hostname"
  press ret

  wait_for_screen "Timezone" 180 # tzupdate geo-guess can take a moment
  capture_console "success-installer-08-timezone"
  press ret

  wait_for_screen "look right" 60
  capture_console "success-installer-09-review"
  press ret

  wait_for_screen "install disk" 60
  capture_console "success-installer-10-install-disk"
  press ret # the single virtio disk

  wait_for_screen "installation mode\|recovery possible" 60
  if ocr_screen | grep -qi "installation mode"; then
    capture_console "success-installer-11-install-mode"
    press ret # "Full disk install" is preselected
  fi

  wait_for_screen "recovery possible" 60
  capture_console "success-installer-12-disk-warning-encrypted"
  if ! $ENCRYPT; then
    press ctrl-c # toggles the confirm to the unencrypted flow
    wait_for_screen "without encryption" 30
    capture_console "success-installer-13-disk-warning-unencrypted"
  fi
  press ret

  sleep 2
  capture_console "success-installer-14-install-started"
}

# The deferred-provisioning first boot runs the user form on tty1 (omarchy-provision-owner) before
# SDDM. Encrypted deferred-provisioning installs auto-unlock from the staged keyfile, so there
# is no LUKS prompt to answer here either way.
drive_provision_owner() {
  log "Driving the deferred-provisioning first-boot setup"

  # Deferred provisioning defers the keyboard step to first boot; it comes before the user form.
  # Pick English (UK) (one row above the preselected English (US)) so the
  # verifier can prove the deferred step applied a non-default layout — UK
  # leaves every letter identical to US, so the typed password (sent as raw
  # keycodes) still matches the literal one the harness pipes to sudo.
  # A greeter (animated logo + tagline) opens first boot; Return begins.
  wait_for_screen "Opinionated" 600
  capture_console "success-provision-00-greeter"
  press ret

  wait_for_screen "keyboard layout" 120
  press up
  capture_console "success-provision-01-keyboard-layout-uk"
  press ret

  wait_for_screen "Username" 120
  capture_console "success-provision-02-username"
  type_text "$GUEST_USER"
  press ret

  wait_for_screen "Password" 60
  type_text "$GUEST_PASSWORD"
  capture_console "success-provision-03-password"
  press ret

  wait_for_screen "Conf *irm" 60
  type_text "$GUEST_PASSWORD"
  press ret

  wait_for_screen "Full name" 60
  type_text "Omarchy Test"
  press ret

  wait_for_screen "Email address" 60
  # The keyboard step above applied English (UK), where '@' is Shift+' (not
  # Shift+2 as on US); send that keycode so the typed email is correct under
  # the new layout — itself further proof the deferred keymap took effect.
  type_text "test"
  press shift-apostrophe
  type_text "omarchy.org"
  capture_console "success-provision-04-email"
  press ret

  # Hostname is deferred to first boot too.
  wait_for_screen "Hostname" 60
  type_text "$GUEST_HOSTNAME"
  capture_console "success-provision-05-hostname"
  press ret

  # Timezone is deferred to first boot too; accept the geo-guessed default.
  # Generous timeout: the tzupdate geo-guess over first-boot networking can lag.
  wait_for_screen "Timezone" 180
  capture_console "success-provision-06-timezone"
  press ret

  wait_for_screen "look right" 60
  capture_console "success-provision-07-review"
  press ret

  # Finalization (offline Node unpack, LUKS re-key on encrypted installs)
  # runs before SDDM appears.
  capture_console "success-provision-08-finalizing"
}

wait_for_install() {
  log "Waiting for the install to finish (timeout ${INSTALL_TIMEOUT}s)"

  local waited=0 text progress_name
  while true; do
    text=$(ocr_screen)

    if grep -qi "Reboot Now" <<<"$text"; then
      log "Install finished. Confirming the reboot prompt."
      capture_console "success-installer-15-reboot"
      press ret
      return 0
    fi

    # Deferred-provisioning installs skip the celebration and reboot on their own
    # (no Reboot Now prompt); the first-boot greeter tagline appearing means the
    # install finished and rebooted into setup.
    if $PROVISION && grep -qi "Opinionated" <<<"$text"; then
      log "Install finished and auto-rebooted into first-boot setup."
      capture_console "success-installer-15-autoreboot"
      return 0
    fi

    if grep -qi "installation stopped" <<<"$text"; then
      capture_console "failure-installer-installation-stopped"
      echo "Install failed — screenshot saved to $RUN_DIR" >&2
      return 1
    fi

    if ! vm_running; then
      echo "VM exited during install" >&2
      return 1
    fi

    if ((waited >= INSTALL_TIMEOUT)); then
      capture_console "failure-installer-install-timeout"
      echo "Timed out after ${INSTALL_TIMEOUT}s waiting for install" >&2
      return 1
    fi

    if ((waited % 60 == 0)); then
      printf -v progress_name 'success-installer-14-install-progress-%04ds' "$waited"
      capture_console "$progress_name"
      echo "    ... installing (${waited}s)"
    fi

    sleep 10
    ((waited += 10))
  done
}

# Encrypted boots stop at the LUKS passphrase prompt. Plymouth's prompt isn't
# reliably OCR-able, so type the passphrase on an interval until a console
# login becomes reachable; early keystrokes are harmlessly discarded.
unlock_luks() {
  log "Typing the LUKS passphrase until boot proceeds"

  local waited=0
  sleep 10
  capture_console "success-first-boot-01-luks-prompt"

  while true; do
    type_text "$GUEST_PASSWORD"
    if ((waited == 0)); then
      capture_console "success-first-boot-02-luks-passphrase"
    fi
    press ret
    sleep 20
    ((waited += 30))

    press ctrl-alt-f3
    sleep 3
    if ocr_screen | grep -qi "login:"; then
      return 0
    fi

    if ((waited >= 300)); then
      capture_console "failure-first-boot-luks-timeout"
      echo "Timed out waiting for LUKS unlock" >&2
      return 1
    fi
  done
}

# Log into a spare console TTY as the user and authorize the harness SSH key —
# the same thing a person would do to give themselves remote access.
bootstrap_ssh() {
  log "Authorizing SSH access via console login"

  mkdir -p "$BASE_DIR/www"
  cat >"$BASE_DIR/www/bootstrap" <<EOF
mkdir -p ~/.ssh && chmod 700 ~/.ssh
echo "$(cat "$SSH_KEY.pub")" >>~/.ssh/authorized_keys
chmod 600 ~/.ssh/authorized_keys
echo "$GUEST_PASSWORD" | sudo -S ufw allow from 10.0.2.2 to any port 22 proto tcp
echo "$GUEST_PASSWORD" | sudo -S systemctl enable --now sshd.service
EOF

  (cd "$BASE_DIR/www" && exec python3 -m http.server "$HTTP_PORT" --bind 127.0.0.1 >/dev/null 2>&1) &
  HTTP_PID=$!

  local waited=0
  while true; do
    press ctrl-alt-f3
    sleep 4
    ocr_screen | grep -qi "login:" && break
    ((waited += 8))

    if ((waited >= 300)); then
      capture_console "failure-first-boot-console-timeout"
      echo "Timed out waiting for a console login prompt" >&2
      return 1
    fi
    sleep 4
  done

  type_text "$GUEST_USER"
  capture_console "success-first-boot-03-console-username"
  press ret
  wait_for_screen "Password" 60
  type_text "$GUEST_PASSWORD"
  capture_console "success-first-boot-04-console-password"
  press ret
  sleep 3

  type_text "curl -fsS http://10.0.2.2:$HTTP_PORT/bootstrap -o /tmp/bs && bash /tmp/bs"
  capture_console "success-first-boot-05-bootstrap-command"
  press ret

  wait_for_ssh 360 "failure-first-boot-ssh-timeout"
  capture_console "success-first-boot-06-bootstrap-complete"

  kill "$HTTP_PID" 2>/dev/null || true
  HTTP_PID=""
  press ctrl-alt-f1
}

# Get a freshly booted system to a running Hyprland session, typing at the
# LUKS prompt and/or SDDM greeter exactly like a user would.
establish_session() {
  if $ENCRYPT; then
    local waited=0
    sleep 10
    capture_console "success-session-01-luks-prompt"
    until ssh_guest true 2>/dev/null; do
      type_text "$GUEST_PASSWORD"
      if ((waited == 0)); then
        capture_console "success-session-02-luks-passphrase"
      fi
      press ret
      sleep 15
      ((waited += 16))

      if ((waited >= BOOT_TIMEOUT)); then
        capture_console "failure-session-luks-timeout"
        echo "Timed out unlocking LUKS" >&2
        return 1
      fi
    done
    # Encrypted installs autologin once the disk is unlocked
  else
    wait_for_ssh "$BOOT_TIMEOUT" "failure-session-ssh-timeout"
    sleep 5
  fi

  local waited=0
  until session_started; do
    if ! $ENCRYPT; then
      # SDDM remembers the user and focuses the password field
      if ((waited == 0)); then
        capture_console "success-session-01-login"
      fi
      type_text "$GUEST_PASSWORD"
      if ((waited == 0)); then
        capture_console "success-session-02-password"
      fi
      press ret
    fi
    sleep 10
    ((waited += 11))

    if ((waited >= 300)); then
      capture_console "failure-session-start-timeout"
      echo "Timed out waiting for the Hyprland session" >&2
      return 1
    fi
  done

  capture_console "success-session-03-desktop"
}

# -------------------------------------------------------------------- phases

install_phase() {
  log "Installing $(basename "$ISO") into a fresh VM (headless, driven via console)"

  [[ -f $SSH_KEY ]] || ssh-keygen -t ed25519 -N "" -q -C "omarchy-iso-test" -f "$SSH_KEY"

  rm -f "$BASE_DISK"
  qemu-img create -f qcow2 "$BASE_DISK" 40G >/dev/null
  cp "$OVMF_VARS_TEMPLATE" "$BASE_OVMF"

  start_vm "$BASE_DISK" "$RUN_DIR/install-serial.log" \
    -drive "file=$ISO,media=cdrom,if=none,format=raw,id=cdrom0" \
    -device ide-cd,drive=cdrom0,bootindex=2

  drive_configurator
  wait_for_install

  if $PROVISION; then
    drive_provision_owner
  elif $ENCRYPT; then
    unlock_luks
  fi
  bootstrap_ssh

  # SSH can come up (console bootstrap logs in on a spare TTY) while
  # omarchy-provision-owner is still finalizing the user and re-keying LUKS. Saving
  # the base mid-finalization produces a flaky image, so wait for the pending
  # flag to clear before powering off.
  if $PROVISION; then
    log "Waiting for provisioning first-boot setup to finish"
    local waited=0 rc
    while true; do
      # Capture rc via `|| rc=$?` so a non-zero result doesn't trip `set -e`
      # (this command is intentionally allowed to "fail" — that's the signal).
      rc=0
      ssh_guest "test -f /var/lib/omarchy/provisioning/pending" 2>/dev/null || rc=$?
      # rc 0 = flag still present; rc 255 = SSH not up yet (finalization may
      # still be running) — keep waiting. Only a clean "file absent" (rc 1)
      # means provisioning finished, so we don't snapshot mid-finalization on a
      # transient SSH blip.
      (( rc == 0 || rc == 255 )) || break
      if ((waited >= 900)); then
        capture_console "failure-provision-finalization-timeout"
        echo "Timed out waiting for provisioning first-boot setup to complete" >&2
        return 1
      fi
      sleep 10
      ((waited += 10))
    done

    # The owner's hostname and timezone are set at first boot now, not baked
    # into the install — confirm they actually landed on the running system.
    # Hostname is deterministic (we typed it), so a mismatch is a hard failure;
    # timezone is geo-guessed, so just report what landed.
    local got_host="" got_tz=""
    got_host=$(ssh_guest "cat /etc/hostname" 2>/dev/null | tr -d '\r\n') || got_host=""
    got_tz=$(ssh_guest "timedatectl show -p Timezone --value" 2>/dev/null | tr -d '\r\n') || got_tz=""
    if [[ $got_host != "$GUEST_HOSTNAME" ]]; then
      capture_console "failure-provision-hostname-not-applied"
      echo "First-boot hostname is '${got_host:-<empty>}', expected '$GUEST_HOSTNAME'" >&2
      return 1
    fi
    log "hostname applied at first boot: $got_host"
    if [[ -n $got_tz && $got_tz != "UTC" ]]; then
      log "timezone applied at first boot: $got_tz"
    else
      echo "WARN: timezone is '${got_tz:-<empty>}' (expected the owner's choice, not UTC)" >&2
    fi

    # Deferred first-boot setup hands straight off to SDDM once the oneshot
    # service exits — no celebration screen, no Start Omarchy button to confirm.
    capture_console "success-provision-09-handoff"
  fi

  ssh_guest "echo $GUEST_PASSWORD | sudo -S cat /var/log/omarchy-install-timing.json 2>/dev/null" \
    >"$RUN_DIR/omarchy-install-timing.json" 2>/dev/null || true
  ssh_guest "echo $GUEST_PASSWORD | sudo -S cat /var/log/pacman.log 2>/dev/null" \
    >"$RUN_DIR/pacman.log" 2>/dev/null || true

  log "Installed system is up. Saving base image."
  stop_vm
}

sync_omarchy() {
  [[ -n $SYNC_DIR ]] || return 0

  local target=".local/share/omarchy"

  if $SYNC_ALL; then
    log "Syncing $SYNC_DIR into the guest"
    tar -C "$SYNC_DIR" --exclude .git -cf - . | ssh_guest "mkdir -p $target && tar -C $target -xf -"
  else
    log "Syncing $SYNC_DIR/test/acceptance into the guest"
    tar -C "$SYNC_DIR/test" -cf - acceptance acceptance.d | ssh_guest "mkdir -p $target/test && tar -C $target/test -xf -"
  fi
}

acceptance_phase() {
  log "Booting acceptance VM from base image overlay"

  local guest_omarchy_path="/usr/share/omarchy"
  $SYNC_ALL && guest_omarchy_path=".local/share/omarchy"

  qemu-img create -f qcow2 -b "$BASE_DISK" -F qcow2 "$RUN_DIR/run.qcow2" >/dev/null
  start_vm "$RUN_DIR/run.qcow2" "$RUN_DIR/serial.log"

  establish_session
  sync_omarchy

  local status=0
  shortcut_smoke_phase || status=1

  log "Running acceptance suite"
  ssh_guest "OMARCHY_PATH=$guest_omarchy_path OMARCHY_ACCEPTANCE_DIR=/tmp/omarchy-acceptance bash .local/share/omarchy/test/acceptance" || status=$?

  log "Collecting artifacts into $RUN_DIR"
  ssh_guest "tar -C /tmp -cf - omarchy-acceptance 2>/dev/null" | tar -C "$RUN_DIR" -xf - || true
  ssh_guest "echo $GUEST_PASSWORD | sudo -S cat /var/log/omarchy-install.log 2>/dev/null" >"$RUN_DIR/omarchy-install.log" 2>/dev/null || true
  if ((status == 0)); then
    capture_console "success-acceptance-final"
  else
    capture_console "failure-acceptance-final"
  fi

  return $status
}

# ---------------------------------------------------------------------- main

if $REUSE_BASE && [[ -f $BASE_DISK ]]; then
  log "Reusing existing base image: $BASE_DISK"
else
  install_phase
fi

if $INSTALL_ONLY; then
  log "Base image ready: $BASE_DISK"
  exit 0
fi

if acceptance_phase; then
  log "Acceptance suite passed. Artifacts: $RUN_DIR"
else
  log "Acceptance suite FAILED. Artifacts: $RUN_DIR"
  exit 1
fi
