#!/bin/bash
# Runs LOCALLY. Prepares the target VM for flashing and boots it into
# initramfs SSH mode:
#   1. Installs dropbear-initramfs
#   2. Copies SSH authorized key into dropbear's keystore
#   3. Configures initramfs networking (IP=dhcp, DEVICE=eth0)
#   4. Installs scripts/init-premount/99-pause hook (loops forever so
#      /dev/vda is never mounted and is free to flash)
#   5. Rebuilds /boot/initrd.img with dropbear + hook
#   6. Cycles the VM (Shutdown → Start) and waits for initramfs SSH
#
# After this script, run 02_flash.sh immediately — /dev/vda is free.

set -euo pipefail

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

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

echo "==> Installing dropbear-initramfs on $REMOTE"
ssh $SSH_OPTS "$REMOTE" "DEBIAN_FRONTEND=noninteractive apt-get install -y -q dropbear-initramfs"

echo ""
echo "==> Configuring dropbear"
ssh $SSH_OPTS "$REMOTE" "
  set -euo pipefail

  # Authorized key for SSH access in initramfs
  mkdir -p /etc/dropbear/initramfs
  cp /root/.ssh/authorized_keys /etc/dropbear/initramfs/authorized_keys
  chmod 600 /etc/dropbear/initramfs/authorized_keys

  grep -q '^DEVICE=' /etc/initramfs-tools/initramfs.conf \
    || echo 'DEVICE=eth0' >> /etc/initramfs-tools/initramfs.conf
  grep -q '^IP=' /etc/initramfs-tools/initramfs.conf \
    || echo 'IP=dhcp' >> /etc/initramfs-tools/initramfs.conf

  echo '  dropbear configured'
"

echo ""
echo "==> Installing init-premount/99-pause hook"
ssh $SSH_OPTS "$REMOTE" 'cat > /etc/initramfs-tools/scripts/init-premount/99-pause <<'"'"'HOOK'"'"'
#!/bin/sh
# mkinitramfs calls every hook with "prereqs" to resolve ordering.
# Must handle it and exit immediately, or the build hangs.
case "$1" in
  prereqs) echo "dropbear"; exit 0 ;;
esac
# Pause here forever: /dev/vda is not yet mounted and is free to flash.
# Dropbear SSH is already running (started by init-premount/dropbear).
log_begin_msg "Initramfs SSH active — port 22 open, /dev/vda not mounted"
while true; do sleep 3600; done
HOOK
chmod +x /etc/initramfs-tools/scripts/init-premount/99-pause
echo "  99-pause hook installed"'

echo ""
echo "==> Rebuilding initrd with dropbear + pause hook"
ssh $SSH_OPTS "$REMOTE" "update-initramfs -u -k \$(uname -r) 2>&1 | tail -2"

echo ""
echo "==> Powering off $REMOTE (sysrq)..."
echo "    CC VMs cannot reboot from inside — a manual Start from Verda is required."
ssh $SSH_OPTS "$REMOTE" "nohup sh -c 'sleep 2 && echo o > /proc/sysrq-trigger' >/dev/null 2>&1 &" || true

cat <<'EOF'

>>> In the Verda console: once the VM is offline, Start it.
    Note: if the VM is already running, Shutdown it first, then Start it.

Waiting for initramfs SSH (dropbear on port 22)...
EOF

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
    pid1=$(ssh -o ConnectTimeout=4 -o BatchMode=yes \
           $SSH_OPTS "$REMOTE" "cat /proc/1/comm" 2>/dev/null || echo "unknown")
    if [ "$pid1" = "systemd" ]; then
      echo "  ... PID 1 is 'systemd' (regular OS boot) — Shutdown the VM from Verda and Start it again."
    else
      echo ""
      echo "==> $REMOTE is up in initramfs SSH mode (PID 1: $pid1)."
      echo ""
      echo "Next step:"
      echo "  ./02_flash.sh"
      exit 0
    fi
  else
    echo "  ... still waiting ($((i * 5))s) — did you Start the VM from Verda?"
  fi
done

echo ""
echo "ERROR: SSH did not come up within 600 s."
echo "Check that the VM was Started from the Verda console."
exit 1