#!/usr/bin/env bash
# Creates the four-card serving VM on a Proxmox node. Runs ON the node as root:
#
#   ssh root@<proxmox-node> "VM_SSH_PUBKEY='$(cat ~/.ssh/<key>.pub)' bash -s -- <vmid> <name> <source-vmid>" \
#     < bench/gpu-vm/create-vm.sh
#
# The GPUs are taken from an existing, stopped VM (<source-vmid>) that already has the
# cards passed through: its hostpci entries and its QEMU args (the 64-bit PCI hole that
# four 24 GB cards need) are copied, nothing else of it is touched. The two VMs share
# those cards, so only one of them can run at a time; the source VM must keep
# onboot 0, and the new VM's description says which cards it borrows.
#
# Ubuntu 24.04 cloud image, q35 + OVMF, DHCP on the node's bridge (the address it got
# is printed at the end). The OS disk is kept small: checkpoints live on the NAS, so
# the node's thin pool (shared with other tenants' disks) is not asked for model space.
set -euo pipefail
main() {

VMID="${1:?usage: create-vm.sh <vmid> <name> <source-vmid>}"
NAME="${2:?usage: create-vm.sh <vmid> <name> <source-vmid>}"
SRC="${3:?usage: create-vm.sh <vmid> <name> <source-vmid>}"
: "${VM_SSH_PUBKEY:?set VM_SSH_PUBKEY to the public key line that may log in}"

CORES="${VM_CORES:-32}"
MEMORY_MB="${VM_MEMORY_MB:-163840}"
DISK_GB="${VM_DISK_GB:-120}"
STORAGE="${VM_STORAGE:-local-lvm}"
BRIDGE="${VM_BRIDGE:-vmbr0}"
IMAGE="${VM_IMAGE:-/var/lib/vz/template/iso/noble-server-cloudimg-amd64.img}"
USER_NAME="${VM_USER:-ubuntu}"

if qm status "$VMID" >/dev/null 2>&1; then
  echo "vmid $VMID already exists on $(hostname); pick another" >&2
  exit 1
fi
if [ "$(qm status "$SRC" | awk '{print $2}')" != stopped ]; then
  echo "source vm $SRC must be stopped: it holds the same cards" >&2
  exit 1
fi
test -f "$IMAGE" || { echo "no cloud image at $IMAGE" >&2; exit 1; }

src_conf="$(qm config "$SRC")"
hostpci="$(printf '%s\n' "$src_conf" | grep '^hostpci[0-9]*:')"
args="$(printf '%s\n' "$src_conf" | sed -n 's/^args: //p')"
[ -n "$hostpci" ] || { echo "source vm $SRC has no hostpci entries" >&2; exit 1; }
cards="$(printf '%s\n' "$hostpci" | awk '{print $2}' | cut -d, -f1 | paste -sd' ')"

keyfile="$(mktemp)"
trap 'rm -f "$keyfile"' EXIT
printf '%s\n' "$VM_SSH_PUBKEY" > "$keyfile"

qm create "$VMID" --name "$NAME" \
  --machine q35 --bios ovmf --cpu host --sockets 1 --cores "$CORES" --numa 1 \
  --memory "$MEMORY_MB" --balloon 0 \
  --efidisk0 "$STORAGE:1,efitype=4m,pre-enrolled-keys=0" \
  --scsihw virtio-scsi-single \
  --net0 "virtio,bridge=$BRIDGE" \
  --serial0 socket --vga none --agent enabled=1 --ostype l26 \
  --ide2 "$STORAGE:cloudinit" \
  --ciuser "$USER_NAME" --sshkeys "$keyfile" --ipconfig0 ip=dhcp \
  --onboot 1 \
  --description "Four-card model-serving replica. Borrows the GPUs of VM $SRC ($cards); only one of the two may run at a time, and VM $SRC stays onboot 0."
qm importdisk "$VMID" "$IMAGE" "$STORAGE" --format raw
qm set "$VMID" --scsi0 "$STORAGE:vm-$VMID-disk-1,discard=on,iothread=1,ssd=1" --boot order=scsi0
qm resize "$VMID" scsi0 "${DISK_GB}G"
[ -z "$args" ] || qm set "$VMID" --args "$args"
while read -r key val; do
  qm set "$VMID" "--${key%:}" "$val"
done <<< "$hostpci"

qm start "$VMID"
# The stock cloud image has no guest agent, so the DHCP address is found by the VM's MAC
# in the node's neighbour table after a ping sweep of the bridge's /24.
mac="$(qm config "$VMID" | sed -n 's/^net0: virtio=\([^,]*\),.*/\1/p' | tr 'A-F' 'a-f')"
net="$(ip -4 -o addr show "$BRIDGE" | awk '{print $4}' | head -1)"
prefix="${net%.*}"
echo "vm $VMID mac $mac; waiting for it to take an address on ${prefix}.0/24"
for _ in $(seq 1 60); do
  for i in $(seq 1 254); do ping -c1 -W1 "$prefix.$i" >/dev/null 2>&1 & done
  wait
  addr="$(ip neigh show dev "$BRIDGE" | awk -v m="$mac" 'tolower($3)==m {print $1; exit}')"
  [ -n "$addr" ] && { echo "$addr"; return 0; }
  sleep 5
done
echo "no address seen for $mac yet; check the serial console (qm terminal $VMID)" >&2

}
main "$@"