#!/bin/sh
# add-node.sh — build a node and put it in the fleet, in one command.
#
# provision-node.sh and onboard-node.sh stay separate because the two halves
# need different access: one the hypervisor, one the certificate authority.
# That is a real boundary and this script does not erase it — it drives both
# where one person holds both, which is the common case and was the case here.
#
# The reason to have it is not keystrokes. The VM name, the hostname, the node
# id and the certificate's common name must all be the same string, and the
# address appears in three places; split across two invocations run minutes
# apart, that agreement is maintained by whoever is typing. Here it is stated
# once.
set -eu

usage() {
  cat >&2 <<'USAGE'
usage: add-node.sh --name NODE --ip ADDR/CIDR --server URL [options]

  --name NODE       node id, hostname and VM name; also the certificate's CN
  --ip ADDR/CIDR    e.g. 10.0.0.145/24
  --server URL      control plane, e.g. https://10.0.0.140:7443

  --vmid N          Proxmox VMID (default: the last octet of --ip)
  --cores N         (default 4)
  --cpu TYPE        QEMU CPU model (default host; see provision-node.sh)
  --memory MB       (default 8192)
  --disk GB         (default 20)
  --fabric NAME     InfiniBand fabric; repeatable as NAME:hca1,hca2.
                    Omit for a node with no fabric.
  --llamacpp-bin P  path to llama-server on the node
  --vllm-bin P      path to vllm
  --sglang-python P path to the python that has sglang
  --binary PATH     llmgw binary to install (default: the one on PATH)
  --pki DIR         certificate authority (default /etc/llmgw/pki)
  --ssh-user U      (default ilker)
  --provision-only  build the VM and stop
  --onboard-only    skip provisioning; the machine already exists
USAGE
  exit 2
}

NAME= IP= SERVER= VMID= CORES=4 MEMORY=8192 DISK=20 CPU=
BINARY= PKI=/etc/llmgw/pki SSH_USER=ilker
PROVISION=1 ONBOARD=1
set -- "$@"
PASS_THROUGH=
while [ $# -gt 0 ]; do
  case "$1" in
    --name) NAME="$2"; shift 2;;
    --ip) IP="$2"; shift 2;;
    --server) SERVER="$2"; shift 2;;
    --vmid) VMID="$2"; shift 2;;
    --cores) CORES="$2"; shift 2;;
    --cpu) CPU="$2"; shift 2;;
    --memory) MEMORY="$2"; shift 2;;
    --disk) DISK="$2"; shift 2;;
    --binary) BINARY="$2"; shift 2;;
    --pki) PKI="$2"; shift 2;;
    --ssh-user) SSH_USER="$2"; shift 2;;
    --provision-only) ONBOARD=0; shift;;
    --onboard-only) PROVISION=0; shift;;
    # Everything the onboarding step understands and this one does not need to.
    --fabric|--llamacpp-bin|--vllm-bin|--sglang-python|--ib-root|--nvidia-smi|--mgmt-addr|--mgmt-if)
      PASS_THROUGH="$PASS_THROUGH $1 $2"; shift 2;;
    -h|--help) usage;;
    *) echo "unknown option: $1" >&2; usage;;
  esac
done
[ -n "$NAME" ] && [ -n "$IP" ] && [ -n "$SERVER" ] || usage

# The address without its prefix, which is what onboarding and ssh want.
ADDR="${IP%%/*}"
# The last octet as the VMID by default: it keeps `qm list` and `ctl nodes` in
# the same order, and one fewer number to choose is one fewer to get wrong.
[ -n "$VMID" ] || VMID="${ADDR##*.}"

here() { dirname "$0"; }
PROVISION_SH="$(here)/provision-node.sh"
ONBOARD_SH="$(here)/onboard-node.sh"
[ -x "$PROVISION_SH" ] || PROVISION_SH="$(command -v llmgw-provision-node || true)"
[ -x "$ONBOARD_SH" ] || ONBOARD_SH="$(command -v llmgw-onboard-node || true)"

if [ "$PROVISION" = 1 ]; then
  [ -x "$PROVISION_SH" ] || { echo "cannot find provision-node.sh" >&2; exit 1; }
  echo "== 1/2  building $NAME as VM $VMID"
  # shellcheck disable=SC2086
  "$PROVISION_SH" --vmid "$VMID" --name "$NAME" --ip "$IP" \
    --cores "$CORES" --memory "$MEMORY" --disk "$DISK" ${CPU:+--cpu "$CPU"}
fi

if [ "$ONBOARD" = 0 ]; then
  echo "$NAME built. Onboard it from the CA host when you are ready."
  exit 0
fi

[ -x "$ONBOARD_SH" ] || { echo "cannot find onboard-node.sh" >&2; exit 1; }
echo
echo "== 2/2  onboarding $NAME"
# Onboarding waits for ssh itself, which is what makes running the two halves
# back to back safe: cloud-init is still regenerating host keys and restarting
# sshd for the first minute of a new machine's life.
# shellcheck disable=SC2086
"$ONBOARD_SH" --node "$NAME" --addr "$ADDR" --server "$SERVER" \
  --pki "$PKI" --ssh-user "$SSH_USER" \
  ${BINARY:+--binary "$BINARY"} $PASS_THROUGH

echo
echo "Check it with: llmgw ctl nodes"
echo
echo "A node is not finished until the engines its profiles name are installed"
echo "and the weights are on disk. Onboarding checks that a configured engine"
echo "runs, but it installs neither — see docs/node-onboarding.md."