#!/usr/bin/env bash
# Install Firecracker + jailer on a fleet node, and prove the node can boot a
# microVM before declaring it ready.
#
# Phase B runs mission workloads in microVMs on the KVM-capable nodes (tank /
# morpheus / architect). gw-04 is itself a VM without nested virtualisation, so
# it has no /dev/kvm and can never host one — KVM availability is a placement
# predicate, not an assumption.
#
# Installing is not the same as working, so this does both and only reports
# success when a VM has actually booted, run our code, and exited. The same
# rule as everywhere else in this repo: a capability that has never been
# observed working has not been shown to work.
#
# Usage:
#   scripts/fc-node-setup.sh <ssh-host> [...]      # e.g. osobh@tank morpheus
#   FC_VERSION=v1.16.1 scripts/fc-node-setup.sh osobh@tank

set -uo pipefail

FC_VERSION="${FC_VERSION:-v1.16.1}"
# Firecracker CI publishes a kernel and an Ubuntu rootfs; both are stand-ins
# for the images we will bake ourselves (one per CLI, per the A6/B4 plan).
CI_BASE="${FC_CI_BASE:-https://s3.amazonaws.com/spec.ccfc.min/firecracker-ci/v1.12/x86_64}"
KERNEL="${FC_KERNEL:-vmlinux-6.1.128}"
ROOTFS="${FC_ROOTFS:-ubuntu-24.04.squashfs}"
WORK="${FC_WORK:-/opt/clawmates-fc}"

FAILURES=0
# Single-quoted $HOME: it must expand on the NODE, not here.
AGENT_BIN=${FC_AGENT_BIN:-'$HOME/clawmates/target/x86_64-unknown-linux-musl/release/fcagent'}
pass() { printf 'PASS  %-14s %s\n' "$1" "$2"; }
fail() { printf 'FAIL  %-14s %s\n' "$1" "$2"; FAILURES=$((FAILURES + 1)); }

[ $# -ge 1 ] || { echo "usage: $0 <ssh-host> [...]" >&2; exit 2; }

for host in "$@"; do
  echo "── $host ────────────────────────────────────────────"

  # 1. KVM. Without it there is nothing to install for.
  if ! ssh -o ConnectTimeout=10 "$host" 'test -e /dev/kvm' 2>/dev/null; then
    fail "$host" "no /dev/kvm — this node cannot host microVMs"
    continue
  fi
  pass "$host" "/dev/kvm present"

  # Present is not the same as usable. /dev/kvm is `crw-rw---- root:kvm`, and
  # the node daemon runs as an ordinary user — on tank and morpheus the kvm
  # group was EMPTY, so the daemon reported `kvm: false` while the device sat
  # right there. The B0 spike missed it entirely because it ran under sudo.
  # This is why the capability probe opens the device instead of stat-ing it.
  if ! ssh "$host" 'test -r /dev/kvm && test -w /dev/kvm' 2>/dev/null; then
    if ssh "$host" 'sudo -n usermod -aG kvm $(id -un) && sudo -n systemctl restart clawmates-node' 2>/dev/null; then
      pass "$host" "added $(ssh "$host" 'id -un') to the kvm group (daemon restarted)"
    else
      fail "$host" "/dev/kvm is not readable/writable by the daemon user and could not be fixed — run: sudo usermod -aG kvm <user> && sudo systemctl restart clawmates-node"
    fi
  else
    pass "$host" "/dev/kvm is openable by the daemon user"
  fi

  # 2. Install, verifying the checksum by hand.
  #
  # `sha256sum -c` against the published file compares by FILENAME, so saving
  # the download under any other name makes it print a warning and exit
  # non-zero for a reason that has nothing to do with integrity — a check that
  # fails for the wrong reason teaches you to ignore it. Compare the hashes.
  installed=$(ssh "$host" "
    set -e
    sudo mkdir -p '$WORK' && sudo chown \$(id -u):\$(id -g) '$WORK'
    cd '$WORK'
    if [ \"\$(firecracker --version 2>/dev/null | head -1)\" = 'Firecracker $FC_VERSION' ]; then
      echo already; exit 0
    fi
    curl -sSL -o fc.tgz 'https://github.com/firecracker-microvm/firecracker/releases/download/$FC_VERSION/firecracker-$FC_VERSION-x86_64.tgz'
    want=\$(curl -sSL 'https://github.com/firecracker-microvm/firecracker/releases/download/$FC_VERSION/firecracker-$FC_VERSION-x86_64.tgz.sha256.txt' | awk '{print \$1}')
    got=\$(sha256sum fc.tgz | awk '{print \$1}')
    [ -n \"\$want\" ] && [ \"\$want\" = \"\$got\" ] || { echo \"checksum \$got != \$want\" >&2; exit 1; }
    tar xzf fc.tgz
    sudo install -m0755 release-$FC_VERSION-x86_64/firecracker-$FC_VERSION-x86_64 /usr/local/bin/firecracker
    sudo install -m0755 release-$FC_VERSION-x86_64/jailer-$FC_VERSION-x86_64 /usr/local/bin/jailer
    echo installed
  " 2>&1) || { fail "$host" "install failed: $installed"; continue; }
  pass "$host" "firecracker $FC_VERSION ($installed)"

  # 3. Kernel + rootfs, converted to a writable ext4.
  ssh "$host" "
    set -e
    cd '$WORK'
    [ -s vmlinux ] || curl -sSL -o vmlinux '$CI_BASE/$KERNEL'
    if [ ! -s rootfs.ext4 ]; then
      curl -sSL -o rootfs.squashfs '$CI_BASE/$ROOTFS'
      command -v unsquashfs >/dev/null || sudo apt-get install -y -qq squashfs-tools >/dev/null 2>&1
      rm -rf squashfs-root && unsquashfs -q -d squashfs-root rootfs.squashfs
      truncate -s 1G rootfs.ext4 && mkfs.ext4 -q -d squashfs-root rootfs.ext4
      rm -rf squashfs-root rootfs.squashfs
    fi
  " >/dev/null 2>&1 || { fail "$host" "could not stage kernel/rootfs"; continue; }
  pass "$host" "kernel + rootfs staged in $WORK"

  # 4. Install the static guest agent into the shared rootfs.
  #
  # A loop mount needs root, and the node daemon deliberately runs as an
  # ordinary user, so this is the one place root is available.
  #
  # The agent is a STATIC musl binary (crates/bins/fcagent). It used to be a
  # python script, which only worked because Firecracker's CI Ubuntu image
  # happens to ship python3 — no image of ours does, so it could never have run
  # in a real mission rootfs. An agent that dictates what must be installed in
  # the image has the dependency backwards.
  #
  # Control is length-prefixed JSON over vsock, NOT the serial console: stdin
  # races the guest's startup and arrives half-consumed.
  ssh "$host" "
    set -e
    cd '$WORK'
    test -x $AGENT_BIN || { echo NO-AGENT-BINARY; exit 1; }
    sudo mkdir -p /mnt/fcroot && sudo mount -o loop rootfs.ext4 /mnt/fcroot
    sudo install -m0755 $AGENT_BIN /mnt/fcroot/usr/local/bin/fcagent
    printf '%s\\n' '#!/bin/sh' 'echo FC-GUEST-ALIVE kernel=\$(uname -r) cpus=\$(nproc)' \
      'exec /usr/local/bin/fcagent' | sudo tee /mnt/fcroot/usr/local/bin/fcinit >/dev/null
    sudo chmod 0755 /mnt/fcroot/usr/local/bin/fcinit
    sudo umount /mnt/fcroot
  " >/dev/null 2>&1 || { fail "$host" "could not install the guest agent (build it: cargo build --release -p fcagent --target x86_64-unknown-linux-musl)"; continue; }
  pass "$host" "guest agent baked into rootfs"

  # The daemon copies the shared rootfs per VM, so it must be readable by the
  # daemon user without sudo.
  ssh "$host" "sudo chmod 0644 '$WORK/rootfs.ext4' && sudo chmod 0644 '$WORK/vmlinux' && sudo mkdir -p '$WORK/vms' && sudo chown \$(id -u):\$(id -g) '$WORK/vms'" >/dev/null 2>&1 \
    && pass "$host" "rootfs/vmlinux readable, $WORK/vms writable by the daemon user" \
    || fail "$host" "could not make $WORK usable by the daemon user"

  # 5. Boot one. Installing proves nothing; this is the check that counts.
  #
  # The guest marker is printed by an init script rather than typed at a shell
  # over the serial console — feeding stdin races the shell's startup and
  # arrives half-consumed (observed: `# ho FC-GUEST-ALIVE`, the first two
  # characters eaten).
  # It does NOT rewrite fcinit — the baked one prints the marker and then execs
  # the agent, and an earlier version of this script clobbered the agent here,
  # which would have left every VM booting into a dead end.
  #
  # Booted is not the same as reachable, so the check goes all the way to an
  # agent round trip. It runs as the DAEMON USER, without sudo, because that is
  # who will actually be starting VMs; the B0 spike passed under sudo and hid a
  # /dev/kvm permission problem for exactly this reason.
  boot=$(ssh "$host" "
    cd '$WORK'
    cat > selftest-vm.json <<JSON
{
  \"boot-source\": {
    \"kernel_image_path\": \"$WORK/vmlinux\",
    \"boot_args\": \"console=ttyS0 reboot=k panic=1 pci=off init=/usr/local/bin/fcinit\"
  },
  \"drives\": [{\"drive_id\":\"rootfs\",\"path_on_host\":\"$WORK/selftest-rootfs.ext4\",\"is_root_device\":true,\"is_read_only\":false}],
  \"machine-config\": {\"vcpu_count\":2,\"mem_size_mib\":1024,\"smt\":false},
  \"vsock\": {\"guest_cid\": 3, \"uds_path\": \"$WORK/selftest.sock\"}
}
JSON
    cp --sparse=always rootfs.ext4 selftest-rootfs.ext4
    rm -f selftest.sock boot.log
    S=\$(date +%s%N)
    setsid timeout 30 firecracker --no-api --config-file selftest-vm.json > boot.log 2>&1 &
    FCPGID=\$!
    for i in \$(seq 1 200); do grep -qa FC-AGENT-LISTENING boot.log && break; sleep 0.05; done
    E=\$(date +%s%N)
    echo \"ms=\$(( (E - S) / 1000000 ))\"
    python3 - <<'PY' 2>&1 | tail -2
import socket, struct, json
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
s.connect('$WORK/selftest.sock')
s.sendall(b'CONNECT 9001\n'); s.recv(64)
req = json.dumps({'op': 'exec', 'cmd': 'echo AGENT-RPC-OK'}).encode()
s.sendall(struct.pack('>I', len(req)) + req)
n = struct.unpack('>I', s.recv(4))[0]
buf = b''
while len(buf) < n: buf += s.recv(n - len(buf))
print(json.loads(buf).get('stdout', '').strip())
PY
    kill -- -\$FCPGID 2>/dev/null
    rm -f selftest.sock selftest-rootfs.ext4 selftest-vm.json
    grep -a FC-GUEST-ALIVE boot.log | head -1 || echo MARKER-ABSENT
  " 2>&1)

  case "$boot" in
    *AGENT-RPC-OK*) pass "$host" "microVM booted and its agent answered over vsock ($(printf '%s' "$boot" | grep -o 'ms=[0-9]*' | head -1))" ;;
    *FC-GUEST-ALIVE*) fail "$host" "microVM booted but its agent did not answer: $(printf '%s' "$boot" | tail -2 | tr '\n' ' ')" ;;
    *)              fail "$host" "microVM did not boot: $(printf '%s' "$boot" | tail -3 | tr '\n' ' ')" ;;
  esac

  # 5. Teardown hygiene. Firecracker does NOT unlink its vsock UDS on exit, and
  # leaves it owned by whoever ran the VM (root, here) — so a driver running as
  # anyone else cannot clean it up. That is the same uid trap that cost this
  # codebase four bugs on the mission checkout; the driver must own the socket
  # path lifecycle explicitly rather than assume the VM tidies up after itself.
  # `pgrep -c` prints its count AND exits 1 when the count is zero, so
  # `|| echo 0` appended a second line and the comparison saw "0\n0" — this
  # check reported FAIL on a perfectly clean teardown the first time it ran.
  leaked=$(ssh "$host" "pgrep -c firecracker 2>/dev/null; true" | head -1 | tr -dc '0-9')
  leaked=${leaked:-0}
  if [ "$leaked" = "0" ]; then
    pass "$host" "no firecracker process survived teardown"
  else
    fail "$host" "$leaked firecracker process(es) still running after teardown"
  fi
done

echo
if [ "$FAILURES" -gt 0 ]; then
  echo "$FAILURES check(s) failed"
  exit 1
fi
echo "all nodes ready"