#!/bin/bash
# Runs LOCALLY. Streams the OS image directly into /dev/vda on the rescue
# machine and then triggers an in-guest poweroff. CC VMs cannot
# reboot from inside the guest, and an in-guest poweroff often leaves the
# VM marked "running" in Verda — so the user must Shutdown AND then Start
# the VM from the Verda console to boot into the freshly flashed OS.
#
# Prerequisites:
#   - Machine must be in dropbear rescue mode (run 02_reboot_to_ramdisk.sh)
#   - OS image built by 00_build_image.sh
#
# The image is streamed via: cat image | ssh remote "dd of=/dev/vda"
# No intermediate copy to the remote — the initramfs tmpfs is never filled.
#
# THIS IS DESTRUCTIVE: /dev/vda is overwritten. Ensure console access
# (VNC) is available before proceeding.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "$SCRIPT_DIR/.env"

: "${REMOTE:?Edit .env and set REMOTE}"

IMG=${IMG:-"$SCRIPT_DIR/questing-server-raw.img"}
DISK=${DISK:-/dev/vda}

# ---- build image if absent ----
if [ ! -f "$IMG" ]; then
  echo "==> Image not found at $IMG — building now"
  bash "$SCRIPT_DIR/00_build_image.sh"
fi

IMG_SIZE=$(du -h "$IMG" | cut -f1)
echo "==> Image: $IMG ($IMG_SIZE)"

# ---- verify rescue mode ----
echo "==> Verifying rescue mode on $REMOTE"
ROOT_MOUNTS=$(ssh $SSH_OPTS "$REMOTE" \
  "grep -c '/dev/vda' /proc/mounts 2>/dev/null || true")
if [ "$ROOT_MOUNTS" -gt 0 ]; then
  echo ""
  echo "ERROR: /dev/vda appears to be mounted on $REMOTE."
  echo "The machine is NOT in rescue mode — aborting to avoid data loss."
  echo "Run 02_reboot_to_ramdisk.sh first."
  exit 1
fi
echo "  /dev/vda not mounted — rescue mode confirmed."

# ---- stream and flash ----
echo ""
echo "==> Flashing $IMG -> $REMOTE:$DISK"
echo "    Streaming via ssh | dd. Connection drops when machine reboots."
echo ""

# Compress on the fly to reduce transfer size. Raw Ubuntu images are ~70%
# compressible (mostly zeros), so gzip -1 typically cuts 2.1 GB to ~650 MB.
# The remote decompresses with gzip before writing to disk.
# Use pv for a progress bar if available.
if command -v pv >/dev/null 2>&1; then
  STREAMER="pv --rate --eta --progress"
else
  echo "    (install pv for a progress bar: sudo apt-get install pv)"
  STREAMER="cat"
fi

$STREAMER "$IMG" | gzip -1 | ssh $SSH_OPTS "$REMOTE" \
  "gzip -d | dd of=$DISK bs=4194304 && sync && echo 'flash complete' && sleep 2 && echo o > /proc/sysrq-trigger" \
  || true  # ssh exits non-zero when the remote powers off and drops the connection
# Note: busybox dd (used in the Ubuntu initramfs) does not support suffix notation
# (bs=4M) or conv=fsync — use numeric bytes and a separate sync.
# sysrq 'o' powers the VM off (not 'b' reboot) because CC VMs
# cannot reboot from inside the guest.

echo ""
echo "Flash complete. In-guest poweroff issued, but Verda may still show"
echo "the VM as 'running'."
echo ">>> In the Verda console: Shutdown the VM, then Start it again. <<<"
echo "    (Both actions are required to actually cycle the VM into the new OS.)"
echo ""
echo "Waiting for new OS to come up once the VM is back up..."
sleep 10
for i in $(seq 1 120); do
  sleep 5
  if ssh -o ConnectTimeout=4 -o BatchMode=yes \
        $SSH_OPTS "$REMOTE" true 2>/dev/null; then
    echo ""
    echo "==> New OS is up on $REMOTE."
    exit 0
  fi
  echo "  ... waiting for new OS ($((i * 5 + 10))s) — did you Shutdown then Start the VM in the Verda console?"
done

echo ""
echo "New OS did not answer SSH within 610 s."
echo "Check that the VM was Shutdown AND then Started from the Verda console."