#!/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)
#   --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
#   --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

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
REUSE_BASE=false
INSTALL_ONLY=false
KEEP_RUNNING=false
SYNC_DIR=""
SYNC_ALL=false

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

while (($#)); do
  case "$1" in
    --encrypt) ENCRYPT=true ;;
    --reuse-base) REUSE_BASE=true ;;
    --install-only) INSTALL_ONLY=true ;;
    --keep-running) KEEP_RUNNING=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 $ISO ]]; then
  ISO=$(ls -t "$REPO_DIR"/release/*.iso 2>/dev/null | head -1) || true
fi

if [[ ! -f ${ISO:-} ]]; then
  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"
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"

QMP_SOCK="$RUN_DIR/qmp.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
}

stop_vm() {
  vm_running || return 0

  qmp '"system_powerdown"' >/dev/null

  local waited=0
  while vm_running && ((waited < 60)); do
    sleep 2
    ((waited += 2))
  done

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

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

  if ! $KEEP_RUNNING; then
    stop_vm
  fi
}
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

  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
      screendump "$RUN_DIR/timeout-${text// /-}.ppm"
      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 ;;
      *) 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 StrictHostKeyChecking=no \
    -o UserKnownHostsFile=/dev/null \
    -o ConnectTimeout=5 \
    -o LogLevel=ERROR \
    "$GUEST_USER@127.0.0.1" "$@"
}

wait_for_ssh() {
  local timeout="$1" 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
      screendump "$RUN_DIR/ssh-timeout.ppm"
      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
}

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

drive_configurator() {
  log "Driving the installer wizard"

  wait_for_screen "keyboard layout" 300
  press ret # English (US) is preselected

  wait_for_screen "Username" 120
  type_text "$GUEST_USER"
  press ret

  wait_for_screen "Password" 60
  type_text "$GUEST_PASSWORD"
  press ret

  wait_for_screen "Must match" 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
  type_text "test@omarchy.org"
  press ret

  wait_for_screen "Hostname" 60
  type_text "$GUEST_HOSTNAME"
  press ret

  wait_for_screen "Timezone" 180 # tzupdate geo-guess can take a moment
  press ret

  wait_for_screen "look right" 60
  press ret

  wait_for_screen "install disk" 60
  press ret # the single virtio disk

  wait_for_screen "overwritten" 60
  if ! $ENCRYPT; then
    press ctrl-c # toggles the confirm to the unencrypted flow
    wait_for_screen "without encryption" 30
  fi
  press ret

  screendump "$RUN_DIR/install-started.ppm"
}

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

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

    if grep -qi "Reboot Now" <<<"$text"; then
      log "Install finished. Confirming the reboot prompt."
      screendump "$RUN_DIR/install-finished.ppm"
      press ret
      return 0
    fi

    if grep -qi "installation stopped" <<<"$text"; then
      screendump "$RUN_DIR/install-failed.ppm"
      echo "Install failed — screen saved to $RUN_DIR/install-failed.ppm" >&2
      return 1
    fi

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

    if ((waited >= INSTALL_TIMEOUT)); then
      screendump "$RUN_DIR/install-timeout.ppm"
      echo "Timed out after ${INSTALL_TIMEOUT}s waiting for install" >&2
      return 1
    fi

    if ((waited % 60 == 0)); then
      screendump "$RUN_DIR/progress.ppm"
      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

  while true; do
    type_text "$GUEST_PASSWORD"
    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
      screendump "$RUN_DIR/luks-timeout.ppm"
      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 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
      screendump "$RUN_DIR/console-timeout.ppm"
      echo "Timed out waiting for a console login prompt" >&2
      return 1
    fi
    sleep 4
  done

  type_text "$GUEST_USER"
  press ret
  wait_for_screen "Password" 60
  type_text "$GUEST_PASSWORD"
  press ret
  sleep 3

  type_text "curl -fsS http://10.0.2.2:$HTTP_PORT/bootstrap | bash"
  press ret

  wait_for_ssh 180

  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
    until ssh_guest true 2>/dev/null; do
      type_text "$GUEST_PASSWORD"
      press ret
      sleep 15
      ((waited += 16))

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

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

    if ((waited >= 300)); then
      screendump "$RUN_DIR/session-timeout.ppm"
      echo "Timed out waiting for the Hyprland session" >&2
      return 1
    fi
  done
}

# -------------------------------------------------------------------- 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 $ENCRYPT; then
    unlock_luks
  fi
  bootstrap_ssh

  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"

  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

  log "Running acceptance suite"
  local status=0
  ssh_guest "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
  screendump "$RUN_DIR/final.ppm"

  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