#!/bin/sh
set -eu

project_root=${PROJECT_ROOT:-$(pwd)}
script_dir=$(CDPATH= cd -- "$(dirname "$0")" && pwd)
fruix_cmd=$project_root/bin/fruix
os_template=${OS_TEMPLATE:-$script_dir/phase19-generation-rollback-operating-system.scm.in}
system_name=${SYSTEM_NAME:-phase19-operating-system}
store_dir=${STORE_DIR:-/frx/store}
disk_capacity=${DISK_CAPACITY:-12g}
root_size=${ROOT_SIZE:-10g}
qemu_smp=${QEMU_SMP:-2}
ssh_port=${QEMU_SSH_PORT:-10024}
base_name=${BASE_NAME:-phase19-generation-layout}
base_version_label=${BASE_VERSION_LABEL:-15.0-STABLE-generation-layout}
base_release=${BASE_RELEASE:-15.0-STABLE}
base_branch=${BASE_BRANCH:-stable/15}
source_name=${SOURCE_NAME:-stable15-generation-layout-source}
source_ref=${SOURCE_REF:-stable/15}
source_commit=${SOURCE_COMMIT:-332708a606f6bf0841c1d4a74c0d067f5640fe89}
declared_source_root=${DECLARED_SOURCE_ROOT:-/var/empty/fruix-unused-source-root-generation-layout}
current_host_name=${CURRENT_HOST_NAME:-fruix-freebsd-current}
candidate_host_name=${CANDIDATE_HOST_NAME:-fruix-freebsd-canary}
metadata_target=${METADATA_OUT:-}
root_authorized_key_file=${ROOT_AUTHORIZED_KEY_FILE:-$HOME/.ssh/id_ed25519.pub}
root_ssh_private_key_file=${ROOT_SSH_PRIVATE_KEY_FILE:-$HOME/.ssh/id_ed25519}
cleanup=0

[ -x "$fruix_cmd" ] || {
  echo "fruix command is not executable: $fruix_cmd" >&2
  exit 1
}
[ -f "$os_template" ] || {
  echo "missing operating-system template: $os_template" >&2
  exit 1
}
[ -f "$root_authorized_key_file" ] || {
  echo "missing root authorized key file: $root_authorized_key_file" >&2
  exit 1
}
[ -f "$root_ssh_private_key_file" ] || {
  echo "missing root SSH private key file: $root_ssh_private_key_file" >&2
  exit 1
}
command -v qemu-system-x86_64 >/dev/null 2>&1 || {
  echo "qemu-system-x86_64 is required" >&2
  exit 1
}
[ -f /usr/local/share/edk2-qemu/QEMU_UEFI_CODE-x86_64.fd ] || {
  echo "missing QEMU UEFI firmware" >&2
  exit 1
}

if [ -n "${WORKDIR:-}" ]; then
  workdir=$WORKDIR
  mkdir -p "$workdir"
else
  workdir=$(mktemp -d /tmp/fruix-phase19-installed-system-rollback-qemu.XXXXXX)
  cleanup=1
fi
if [ "${KEEP_WORKDIR:-0}" -eq 1 ]; then
  cleanup=0
fi

current_os_file=$workdir/current-operating-system.scm
candidate_os_file=$workdir/candidate-operating-system.scm
current_install_out=$workdir/current-install.txt
candidate_build_out=$workdir/candidate-build.txt
target_image=$workdir/installed.img
candidate_store_items=$workdir/candidate-store-items.txt
stage_log=$workdir/stage-candidate-store.txt
serial_log=$workdir/serial.log
qemu_pidfile=$workdir/qemu.pid
uefi_vars=$workdir/QEMU_UEFI_VARS.fd
metadata_file=$workdir/phase19-installed-system-rollback-qemu-metadata.txt
switch_status_file=$workdir/switch-status.txt
rollback_status_file=$workdir/rollback-status.txt
boot_current_status_file=$workdir/boot-current-status.txt
boot_candidate_status_file=$workdir/boot-candidate-status.txt
boot_rollback_status_file=$workdir/boot-rollback-status.txt
generation2_metadata_file=$workdir/generation2-metadata.scm
generation2_install_file=$workdir/generation2-install.scm
mnt_root=$workdir/mnt-root
md_unit=

cleanup_workdir() {
  if [ -f "$qemu_pidfile" ]; then
    sudo kill "$(sudo cat "$qemu_pidfile")" >/dev/null 2>&1 || true
  fi
  if [ -n "$md_unit" ]; then
    sudo umount "$mnt_root" >/dev/null 2>&1 || true
    sudo mdconfig -d -u "$md_unit" >/dev/null 2>&1 || true
  fi
  if [ "$cleanup" -eq 1 ]; then
    rm -rf "$workdir" 2>/dev/null || sudo rm -rf "$workdir"
  fi
}
trap cleanup_workdir EXIT INT TERM

render_os() {
  output=$1
  host_name=$2
  root_authorized_key=$(tr -d '\n' < "$root_authorized_key_file")
  sed \
    -e "s|__BASE_NAME__|$base_name|g" \
    -e "s|__BASE_VERSION_LABEL__|$base_version_label|g" \
    -e "s|__BASE_RELEASE__|$base_release|g" \
    -e "s|__BASE_BRANCH__|$base_branch|g" \
    -e "s|__SOURCE_NAME__|$source_name|g" \
    -e "s|__SOURCE_REF__|$source_ref|g" \
    -e "s|__SOURCE_COMMIT__|$source_commit|g" \
    -e "s|__DECLARED_SOURCE_ROOT__|$declared_source_root|g" \
    -e "s|__HOST_NAME__|$host_name|g" \
    -e "s|__ROOT_AUTHORIZED_KEY__|$root_authorized_key|g" \
    "$os_template" > "$output"
}

render_os "$current_os_file" "$current_host_name"
render_os "$candidate_os_file" "$candidate_host_name"
cp /usr/local/share/edk2-qemu/QEMU_UEFI_VARS-x86_64.fd "$uefi_vars"
mkdir -p "$mnt_root"

action_env() {
  sudo env \
    HOME="$HOME" \
    GUILE_AUTO_COMPILE=0 \
    FRUIX_FREEBSD_BUILD_JOBS="${FRUIX_FREEBSD_BUILD_JOBS:-8}" \
    GUIX_SOURCE_DIR="${GUIX_SOURCE_DIR:-$HOME/repos/guix}" \
    GUILE_BIN="${GUILE_BIN:-/tmp/guile-freebsd-validate-install/bin/guile}" \
    GUILE_EXTRA_PREFIX="${GUILE_EXTRA_PREFIX:-/tmp/guile-gnutls-freebsd-validate-install}" \
    SHEPHERD_PREFIX="${SHEPHERD_PREFIX:-/tmp/shepherd-freebsd-validate-install}" \
    "$@"
}

field() {
  name=$1
  file=$2
  sed -n "s/^$name=//p" "$file" | tail -n 1
}

status_field() {
  name=$1
  file=$2
  sed -n "s/^$name=//p" "$file" | tail -n 1
}

wait_for_ssh() {
  for attempt in $(jot 120 1 120); do
    if ssh_guest 'service sshd onestatus >/dev/null 2>&1' >/dev/null 2>&1; then
      return 0
    fi
    sleep 2
  done
  return 1
}

ssh_guest() {
  ssh -p "$ssh_port" -i "$root_ssh_private_key_file" \
    -o BatchMode=yes \
    -o StrictHostKeyChecking=no \
    -o UserKnownHostsFile=/dev/null \
    -o ConnectTimeout=5 \
    root@127.0.0.1 "$@"
}

reboot_guest() {
  ssh_guest 'shutdown -r now >/dev/null 2>&1 || reboot >/dev/null 2>&1 || true' >/dev/null 2>&1 || true
  sleep 5
  wait_for_ssh || {
    echo "guest did not return over SSH after reboot" >&2
    exit 1
  }
}

capture_status() {
  output_file=$1
  ssh_guest '/usr/local/bin/fruix system status' > "$output_file"
}

action_env "$fruix_cmd" system install "$current_os_file" \
  --system "$system_name" \
  --store "$store_dir" \
  --target "$target_image" \
  --disk-capacity "$disk_capacity" \
  --root-size "$root_size" > "$current_install_out"

current_closure_path=$(field closure_path "$current_install_out")
install_metadata_path=$(field install_metadata_path "$current_install_out")
materialized_source_store=$(field materialized_source_stores "$current_install_out")
[ -n "$current_closure_path" ] || { echo "missing current closure path" >&2; exit 1; }
[ "$install_metadata_path" = /var/lib/fruix/install.scm ] || { echo "unexpected install metadata path: $install_metadata_path" >&2; exit 1; }

action_env "$fruix_cmd" system build "$candidate_os_file" \
  --system "$system_name" \
  --store "$store_dir" > "$candidate_build_out"

candidate_closure_path=$(field closure_path "$candidate_build_out")
[ -n "$candidate_closure_path" ] || { echo "missing candidate closure path" >&2; exit 1; }
[ "$candidate_closure_path" != "$current_closure_path" ] || { echo "candidate closure unexpectedly matches current closure" >&2; exit 1; }

{
  printf '%s\n' "$candidate_closure_path"
  cat "$candidate_closure_path/.references"
} | awk 'NF { print }' | sort -u > "$candidate_store_items"

candidate_store_item_count=$(wc -l < "$candidate_store_items" | tr -d ' ')

md=$(sudo mdconfig -a -t vnode -f "$target_image")
md_unit=${md#md}
sudo mount -t ufs "/dev/${md}p2" "$mnt_root"
sudo mkdir -p "$mnt_root/frx/store"
: > "$stage_log"
while IFS= read -r item; do
  [ -n "$item" ] || continue
  item_base=$(basename "$item")
  if [ -e "$mnt_root/frx/store/$item_base" ]; then
    printf 'already-present=%s\n' "$item_base" >> "$stage_log"
    continue
  fi
  printf 'copy=%s\n' "$item_base" >> "$stage_log"
  sudo sh -c "cd '$store_dir' && pax -rw -pe '$item_base' '$mnt_root/frx/store'"
done < "$candidate_store_items"
sudo sync
sudo umount "$mnt_root"
sudo mdconfig -d -u "$md_unit"
md_unit=

sudo qemu-system-x86_64 \
  -machine q35,accel=tcg \
  -cpu max \
  -m 2048 \
  -smp "$qemu_smp" \
  -display none \
  -serial "file:$serial_log" \
  -monitor none \
  -pidfile "$qemu_pidfile" \
  -daemonize \
  -drive if=pflash,format=raw,readonly=on,file=/usr/local/share/edk2-qemu/QEMU_UEFI_CODE-x86_64.fd \
  -drive if=pflash,format=raw,file="$uefi_vars" \
  -drive if=virtio,format=raw,file="$target_image" \
  -netdev user,id=net0,hostfwd=tcp::${ssh_port}-:22 \
  -device virtio-net-pci,netdev=net0

wait_for_ssh || {
  echo "guest never became reachable over SSH" >&2
  exit 1
}

capture_status "$boot_current_status_file"
boot_current_closure=$(ssh_guest 'readlink /run/current-system')
boot_current_hostname=$(ssh_guest 'hostname')
boot_current_generation=$(status_field current_generation "$boot_current_status_file")
boot_current_link=$(status_field current_link "$boot_current_status_file")
boot_current_rollback_generation=$(status_field rollback_generation "$boot_current_status_file")
[ "$boot_current_closure" = "$current_closure_path" ] || { echo "unexpected current closure after initial boot: $boot_current_closure" >&2; exit 1; }
[ "$boot_current_hostname" = "$current_host_name" ] || { echo "unexpected hostname after initial boot: $boot_current_hostname" >&2; exit 1; }
[ "$boot_current_generation" = 1 ] || { echo "unexpected initial generation: $boot_current_generation" >&2; exit 1; }
[ "$boot_current_link" = generations/1 ] || { echo "unexpected initial current link: $boot_current_link" >&2; exit 1; }
[ -z "$boot_current_rollback_generation" ] || { echo "rollback generation should be empty before switch" >&2; exit 1; }
ssh_guest 'test -x /usr/local/bin/fruix'

ssh_guest "/usr/local/bin/fruix system switch $candidate_closure_path" > "$switch_status_file"
switch_current_generation=$(status_field current_generation "$switch_status_file")
switch_current_link=$(status_field current_link "$switch_status_file")
switch_current_closure=$(status_field current_closure "$switch_status_file")
switch_rollback_generation=$(status_field rollback_generation "$switch_status_file")
switch_rollback_link=$(status_field rollback_link "$switch_status_file")
switch_rollback_closure=$(status_field rollback_closure "$switch_status_file")
[ "$switch_current_generation" = 2 ] || { echo "unexpected generation after switch: $switch_current_generation" >&2; exit 1; }
[ "$switch_current_link" = generations/2 ] || { echo "unexpected current link after switch: $switch_current_link" >&2; exit 1; }
[ "$switch_current_closure" = "$candidate_closure_path" ] || { echo "unexpected current closure after switch: $switch_current_closure" >&2; exit 1; }
[ "$switch_rollback_generation" = 1 ] || { echo "unexpected rollback generation after switch: $switch_rollback_generation" >&2; exit 1; }
[ "$switch_rollback_link" = generations/1 ] || { echo "unexpected rollback link after switch: $switch_rollback_link" >&2; exit 1; }
[ "$switch_rollback_closure" = "$current_closure_path" ] || { echo "unexpected rollback closure after switch: $switch_rollback_closure" >&2; exit 1; }
[ "$(ssh_guest 'readlink /frx/var/fruix/gcroots/current-system')" = "$candidate_closure_path" ] || { echo "unexpected current-system gc root after switch" >&2; exit 1; }
[ "$(ssh_guest 'readlink /frx/var/fruix/gcroots/rollback-system')" = "$current_closure_path" ] || { echo "unexpected rollback-system gc root after switch" >&2; exit 1; }
[ "$(ssh_guest 'readlink /frx/var/fruix/gcroots/system-2')" = "$candidate_closure_path" ] || { echo "unexpected system-2 gc root after switch" >&2; exit 1; }
ssh_guest 'test -f /var/lib/fruix/system/generations/2/metadata.scm'
ssh_guest 'test -f /var/lib/fruix/system/generations/2/provenance.scm'
ssh_guest 'test -f /var/lib/fruix/system/generations/2/install.scm'
ssh_guest "cat /var/lib/fruix/system/generations/2/metadata.scm" > "$generation2_metadata_file"
ssh_guest "cat /var/lib/fruix/system/generations/2/install.scm" > "$generation2_install_file"
case "$(cat "$generation2_metadata_file")" in
  *"$candidate_closure_path"*"$current_closure_path"*) : ;;
  *) echo "generation 2 metadata does not record both candidate and previous closure paths" >&2; exit 1 ;;
esac
case "$(cat "$generation2_install_file")" in
  *"(deployment-kind . \"switch\")"*"$candidate_closure_path"*) : ;;
  *) echo "generation 2 install metadata does not record switch provenance" >&2; exit 1 ;;
esac

reboot_guest
capture_status "$boot_candidate_status_file"
boot_candidate_closure=$(ssh_guest 'readlink /run/current-system')
boot_candidate_hostname=$(ssh_guest 'hostname')
boot_candidate_generation=$(status_field current_generation "$boot_candidate_status_file")
boot_candidate_rollback_generation=$(status_field rollback_generation "$boot_candidate_status_file")
boot_candidate_rollback_closure=$(status_field rollback_closure "$boot_candidate_status_file")
[ "$boot_candidate_closure" = "$candidate_closure_path" ] || { echo "unexpected closure after switch reboot: $boot_candidate_closure" >&2; exit 1; }
[ "$boot_candidate_hostname" = "$candidate_host_name" ] || { echo "unexpected hostname after switch reboot: $boot_candidate_hostname" >&2; exit 1; }
[ "$boot_candidate_generation" = 2 ] || { echo "unexpected generation after switch reboot: $boot_candidate_generation" >&2; exit 1; }
[ "$boot_candidate_rollback_generation" = 1 ] || { echo "unexpected rollback generation after switch reboot: $boot_candidate_rollback_generation" >&2; exit 1; }
[ "$boot_candidate_rollback_closure" = "$current_closure_path" ] || { echo "unexpected rollback closure after switch reboot: $boot_candidate_rollback_closure" >&2; exit 1; }

ssh_guest '/usr/local/bin/fruix system rollback' > "$rollback_status_file"
rollback_current_generation=$(status_field current_generation "$rollback_status_file")
rollback_current_link=$(status_field current_link "$rollback_status_file")
rollback_current_closure=$(status_field current_closure "$rollback_status_file")
rollback_rollback_generation=$(status_field rollback_generation "$rollback_status_file")
rollback_rollback_link=$(status_field rollback_link "$rollback_status_file")
rollback_rollback_closure=$(status_field rollback_closure "$rollback_status_file")
[ "$rollback_current_generation" = 1 ] || { echo "unexpected generation after rollback: $rollback_current_generation" >&2; exit 1; }
[ "$rollback_current_link" = generations/1 ] || { echo "unexpected current link after rollback: $rollback_current_link" >&2; exit 1; }
[ "$rollback_current_closure" = "$current_closure_path" ] || { echo "unexpected current closure after rollback: $rollback_current_closure" >&2; exit 1; }
[ "$rollback_rollback_generation" = 2 ] || { echo "unexpected rollback generation after rollback: $rollback_rollback_generation" >&2; exit 1; }
[ "$rollback_rollback_link" = generations/2 ] || { echo "unexpected rollback link after rollback: $rollback_rollback_link" >&2; exit 1; }
[ "$rollback_rollback_closure" = "$candidate_closure_path" ] || { echo "unexpected rollback closure after rollback: $rollback_rollback_closure" >&2; exit 1; }
[ "$(ssh_guest 'readlink /frx/var/fruix/gcroots/current-system')" = "$current_closure_path" ] || { echo "unexpected current-system gc root after rollback" >&2; exit 1; }
[ "$(ssh_guest 'readlink /frx/var/fruix/gcroots/rollback-system')" = "$candidate_closure_path" ] || { echo "unexpected rollback-system gc root after rollback" >&2; exit 1; }

reboot_guest
capture_status "$boot_rollback_status_file"
boot_rollback_closure=$(ssh_guest 'readlink /run/current-system')
boot_rollback_hostname=$(ssh_guest 'hostname')
boot_rollback_generation=$(status_field current_generation "$boot_rollback_status_file")
boot_rollback_rollback_generation=$(status_field rollback_generation "$boot_rollback_status_file")
boot_rollback_rollback_closure=$(status_field rollback_closure "$boot_rollback_status_file")
shepherd_status=$(ssh_guest '/usr/local/etc/rc.d/fruix-shepherd onestatus >/dev/null 2>&1 && echo running || echo stopped')
sshd_status=$(ssh_guest 'service sshd onestatus >/dev/null 2>&1 && echo running || echo stopped')
activate_log=$(ssh_guest 'cat /var/log/fruix-activate.log 2>/dev/null || true' | tr '\n' ' ')
[ "$boot_rollback_closure" = "$current_closure_path" ] || { echo "unexpected closure after rollback reboot: $boot_rollback_closure" >&2; exit 1; }
[ "$boot_rollback_hostname" = "$current_host_name" ] || { echo "unexpected hostname after rollback reboot: $boot_rollback_hostname" >&2; exit 1; }
[ "$boot_rollback_generation" = 1 ] || { echo "unexpected generation after rollback reboot: $boot_rollback_generation" >&2; exit 1; }
[ "$boot_rollback_rollback_generation" = 2 ] || { echo "unexpected rollback generation after rollback reboot: $boot_rollback_rollback_generation" >&2; exit 1; }
[ "$boot_rollback_rollback_closure" = "$candidate_closure_path" ] || { echo "unexpected rollback closure after rollback reboot: $boot_rollback_rollback_closure" >&2; exit 1; }
[ "$shepherd_status" = running ] || { echo "fruix-shepherd is not running after rollback reboot" >&2; exit 1; }
[ "$sshd_status" = running ] || { echo "sshd is not running after rollback reboot" >&2; exit 1; }
case "$activate_log" in
  *fruix-activate:done*) : ;;
  *) echo "activation log does not show success after rollback workflow" >&2; exit 1 ;;
esac

cat >"$metadata_file" <<EOF
workdir=$workdir
current_os_file=$current_os_file
candidate_os_file=$candidate_os_file
target_image=$target_image
current_closure_path=$current_closure_path
candidate_closure_path=$candidate_closure_path
current_host_name=$current_host_name
candidate_host_name=$candidate_host_name
candidate_store_item_count=$candidate_store_item_count
candidate_store_items=$candidate_store_items
stage_log=$stage_log
serial_log=$serial_log
install_metadata_path=$install_metadata_path
materialized_source_store=$materialized_source_store
switch_status_file=$switch_status_file
rollback_status_file=$rollback_status_file
boot_current_status_file=$boot_current_status_file
boot_candidate_status_file=$boot_candidate_status_file
boot_rollback_status_file=$boot_rollback_status_file
generation2_metadata_file=$generation2_metadata_file
generation2_install_file=$generation2_install_file
final_current_generation=$boot_rollback_generation
final_current_closure=$boot_rollback_closure
final_rollback_generation=$boot_rollback_rollback_generation
final_rollback_closure=$boot_rollback_rollback_closure
shepherd_status=$shepherd_status
sshd_status=$sshd_status
boot_backend=qemu-uefi-tcg
installed_system_switch=ok
installed_system_rollback=ok
EOF

if [ -n "$metadata_target" ]; then
  mkdir -p "$(dirname "$metadata_target")"
  cp "$metadata_file" "$metadata_target"
fi

printf 'PASS phase19-installed-system-rollback-qemu\n'
printf 'Work directory: %s\n' "$workdir"
printf 'Metadata file: %s\n' "$metadata_file"
if [ -n "$metadata_target" ]; then
  printf 'Copied metadata to: %s\n' "$metadata_target"
fi
printf '%s\n' '--- metadata ---'
cat "$metadata_file"