#!/usr/bin/env bash
#
# Give a machine an identity in an llmgw cluster.
#
# Run on the host that holds the certificate authority, because that is the one
# thing that cannot be delegated: this script issues the credential by which the
# control plane will believe the node is who it says it is.
#
# What it does, in the order it matters:
#   1. issue a client certificate whose common name is the node id
#   2. copy the CA root, the certificate and its key to the machine
#   3. trust the root system-wide, so other tooling there stops complaining
#   4. install the binary, the configuration and the unit
#   5. start the agent
#
# The node id, the VM name, the hostname and the certificate common name are all
# the same string. They do not have to be, and every time they are not, someone
# spends an afternoon working out which of the four a message is talking about.
set -euo pipefail

usage() {
  cat >&2 <<USAGE
usage: $0 --node ID --addr HOST --server URL [options]

  --node ID         node id, hostname and certificate common name
  --addr HOST       how to reach the machine over ssh
  --server URL      the control plane's agent endpoint, e.g. https://10.0.0.140:7443;
                    several, comma separated, for a node that fails over
  --fabric NAME     InfiniBand fabric this node is attached to; repeatable as
                    --fabric NAME:hca1,hca2 . Omit for a node with no fabric,
                    which can host single-node replicas and will never be given
                    a multi-node one.
  --mgmt-addr ADDR  address engines are reachable on (default: the node's
                    address on the route to --server, found by the agent)
  --mgmt-if NAME    interface NCCL bootstraps over (default: the one holding
                    that address)
  --llamacpp-bin P  use this llama-server instead of the one the agent installs
  --vllm-bin P      path to the vllm entry point (not yet installed by the agent)
  --sglang-python P python interpreter of the SGLang virtualenv
  --nvidia-smi PATH where nvidia-smi lives, when it is not on PATH. Containerized
                    driver installs routinely put it somewhere else.
  --ib-root DIR     the sysfs InfiniBand directory (default /sys/class/infiniband)
  --pki DIR         where the authority lives (default /etc/llmgw/pki)
  --binary PATH     llmgw binary to install (default: the one running this).
                    Releases carry static linux/amd64 and linux/arm64 builds;
                    downloading one is better than copying whatever happens to
                    be on the machine you are running this from.
  --releases URL    where release binaries are downloaded from (default the
                    project's Gitea releases)
  --ssh-user NAME   (default ilker)
USAGE
  exit 2
}

NODE= ADDR= SERVER= MGMT_OVERRIDE= MGMT_IF_OVERRIDE=
NVIDIA_SMI= IB_ROOT=/sys/class/infiniband
LLAMACPP_BIN= VLLM_BIN= SGLANG_PYTHON=
SSH_WAIT=180
PKI=/etc/llmgw/pki BINARY= SSH_USER=ilker
RELEASES=https://gitea.mynodes.xyz/ilker/llmgw/releases/download
FABRIC_NAME= FABRIC_HCAS="mlx5_0,mlx5_1"

while [ $# -gt 0 ]; do
  case "$1" in
    --node) NODE="$2"; shift 2;;
    --addr) ADDR="$2"; shift 2;;
    --server) SERVER="$2"; shift 2;;
    --fabric) FABRIC_NAME="${2%%:*}"; [ "${2#*:}" != "$2" ] && FABRIC_HCAS="${2#*:}"; shift 2;;
    --mgmt-addr) MGMT_OVERRIDE="$2"; shift 2;;
    --mgmt-if) MGMT_IF_OVERRIDE="$2"; shift 2;;
    --llamacpp-bin) LLAMACPP_BIN="$2"; shift 2;;
    --vllm-bin) VLLM_BIN="$2"; shift 2;;
    --sglang-python) SGLANG_PYTHON="$2"; shift 2;;
    --nvidia-smi) NVIDIA_SMI="$2"; shift 2;;
    --ib-root) IB_ROOT="$2"; shift 2;;
    --pki) PKI="$2"; shift 2;;
    --binary) BINARY="$2"; shift 2;;
    --releases) RELEASES="$2"; shift 2;;
    --ssh-user) SSH_USER="$2"; shift 2;;
    *) usage;;
  esac
done
[ -n "$NODE" ] && [ -n "$ADDR" ] && [ -n "$SERVER" ] || usage
BINARY="${BINARY:-$(command -v llmgw)}"
[ -x "$BINARY" ] || { echo "no llmgw binary; pass --binary" >&2; exit 1; }
[ -f "$PKI/ca-key.pem" ] || { echo "no authority in $PKI; run 'llmgw ctl cert init-ca' first" >&2; exit 1; }

# accept-new, not "no": a machine provisioned a minute ago has a host key
# nobody has seen, and refusing it stopped onboarding dead with "Host key
# verification failed" — an error about ssh, in the middle of a script whose
# whole job is to talk to a new machine. accept-new learns an unknown key and
# still refuses a changed one, which is the distinction that matters: the risk
# being guarded against is a substituted machine, not a new one.
SSHOPTS="-o BatchMode=yes -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new"
SSH="ssh $SSHOPTS $SSH_USER@$ADDR"
say() { printf '  %s\n' "$*"; }

echo "onboarding $NODE at $ADDR"

# Waiting rather than failing: cloud-init restarts sshd after it regenerates
# the host keys, so a machine that answered a moment ago refuses connections
# for a few seconds. Observed while onboarding node05 — the agent was already
# registered and healthy while ssh was refusing.
wait_for_ssh() {
  deadline=$(( $(date +%s) + SSH_WAIT ))
  until $SSH true 2>/dev/null; do
    [ "$(date +%s)" -lt "$deadline" ] || {
      echo "cannot reach $SSH_USER@$ADDR over ssh after ${SSH_WAIT}s" >&2
      exit 1
    }
    sleep 5
  done
}
wait_for_ssh

# 1. The certificate. Issued fresh every time: re-onboarding a machine should
#    not depend on finding an old key, and the old one keeps working until it
#    is revoked, which is a separate decision.
say "issuing a certificate"
TMP="$(mktemp -d)"
trap 'rm -rf "$TMP"' EXIT
"$BINARY" ctl cert issue-client "$NODE" --dir "$PKI" --out "$TMP" >/dev/null

# 2 and 3. Identity, and trusting the root so future tooling on that machine
#    does not treat our certificates as strangers.
say "installing the certificate and trusting the root"
scp -q $SSHOPTS "$PKI/ca.pem" "$TMP/$NODE.pem" "$TMP/$NODE-key.pem" "$SSH_USER@$ADDR:/tmp/"
$SSH "set -e
  sudo install -d -m 0755 /etc/llmgw
  sudo install -m 0644 /tmp/ca.pem /etc/llmgw/ca.pem
  sudo install -m 0644 /tmp/$NODE.pem /etc/llmgw/agent.pem
  sudo install -m 0600 /tmp/$NODE-key.pem /etc/llmgw/agent-key.pem
  sudo chown root:root /etc/llmgw/*
  sudo install -m 0644 /tmp/ca.pem /usr/local/share/ca-certificates/llmgw-ca.crt
  sudo update-ca-certificates >/dev/null 2>&1 || true
  rm -f /tmp/ca.pem /tmp/$NODE.pem /tmp/$NODE-key.pem"

# 4. The binary, the configuration and the unit.
# The node fetches the release CI built for this version, for its own
# architecture, and checks it against the release's checksums — rather than
# being handed whatever binary is on the machine running this script. A build
# that is not a release (a developer's, with commits past the tag) has no such
# file, and is copied as before.
TAG="$("$BINARY" --version 2>/dev/null | awk '{print $NF}')"
if printf '%s' "$TAG" | grep -Eq '^v[0-9]+\.[0-9]+\.[0-9]+$'; then
  say "installing llmgw $TAG from the release"
  $SSH "set -e
    arch=\$(dpkg --print-architecture)
    name=llmgw-$TAG-linux-\$arch
    cd /tmp && rm -f \$name checksums.txt
    curl -fsSL -o \$name '$RELEASES/$TAG/'\$name
    curl -fsSL -o checksums.txt '$RELEASES/$TAG/checksums.txt'
    grep \" \$name\\\$\" checksums.txt | sha256sum -c - >/dev/null
    sudo install -m 0755 \$name /usr/local/bin/llmgw
    rm -f \$name checksums.txt"
else
  say "installing the binary (${TAG:-unversioned}, not a release)"
  scp -q $SSHOPTS "$BINARY" "$SSH_USER@$ADDR:/tmp/llmgw"
  $SSH 'sudo install -m 0755 /tmp/llmgw /usr/local/bin/llmgw && rm -f /tmp/llmgw'
fi

say "writing the configuration"
NODE_IS_HOSTNAME=no
[ "$($SSH hostname 2>/dev/null)" = "$NODE" ] && NODE_IS_HOSTNAME=yes
# Only what the node cannot find out for itself. Its name is its hostname, its
# address is the one on the route to the control plane, llama-server is looked
# for on PATH and where release archives are unpacked, GPUs are there when
# nvidia-smi is, and memory, disk and engines are re-read while it runs. Every
# setting written here is one that has to be kept true by hand — so an option
# below is written only when it was given, as an override.
{
  echo "# $NODE, written by onboard-node.sh."
  # Several addresses, comma separated, are every control plane the node may
  # fail over to; a standby refuses it and it moves on to the next.
  if [ "${SERVER#*,}" != "$SERVER" ]; then
    echo "servers:"
    printf '%s\n' "$SERVER" | tr ',' '\n' | sed 's/^ *//; s/^/  - "/; s/$/"/'
  else
    echo "server: \"$SERVER\""
  fi
  echo
  echo "tls:"
  echo "  ca: /etc/llmgw/ca.pem"
  echo "  cert: /etc/llmgw/agent.pem"
  echo "  key: /etc/llmgw/agent-key.pem"
  # The certificate's common name is --node, and the control plane refuses a
  # stream whose node id differs from it. Said explicitly only when the
  # hostname would say something else.
  [ "$NODE_IS_HOSTNAME" = yes ] || echo "node_id: $NODE"
  [ -n "$MGMT_OVERRIDE" ] && echo "mgmt_addr: \"$MGMT_OVERRIDE\""
  [ -n "$MGMT_IF_OVERRIDE" ] && echo "mgmt_interface: $MGMT_IF_OVERRIDE"
  if [ -n "$FABRIC_NAME" ]; then
    # Fabric membership is the one thing that cannot be found: OpenSM's default
    # subnet prefix is the same on every fabric, so a matching prefix proves
    # nothing and a wrong guess hangs the first all-reduce.
    echo "fabrics:"
    echo "  - name: $FABRIC_NAME"
    echo "    hcas: [$(echo "$FABRIC_HCAS" | sed 's/,/, /g')]"
  fi
  if [ -n "$NVIDIA_SMI" ] || [ "$IB_ROOT" != /sys/class/infiniband ]; then
    echo "discovery:"
    [ -n "$NVIDIA_SMI" ] && echo "  nvidia_smi: \"$NVIDIA_SMI\""
    [ "$IB_ROOT" != /sys/class/infiniband ] && echo "  ib_root: $IB_ROOT"
  fi
  if [ -n "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
    echo "engines:"
    [ -n "$LLAMACPP_BIN" ] && printf '  llamacpp:\n    bin: "%s"\n' "$LLAMACPP_BIN"
    [ -n "$VLLM_BIN" ] && printf '  vllm:\n    bin: "%s"\n' "$VLLM_BIN"
    [ -n "$SGLANG_PYTHON" ] && printf '  sglang:\n    python: "%s"\n' "$SGLANG_PYTHON"
  fi
  true
} > "$TMP/agent.yaml"
scp -q "$TMP/agent.yaml" "$SSH_USER@$ADDR:/tmp/agent.yaml"

cat > "$TMP/llmgw-agent.service" <<'UNIT'
[Unit]
Description=llmgw node agent
After=network-online.target
Wants=network-online.target

[Service]
ExecStart=/usr/local/bin/llmgw agent --config /etc/llmgw/agent.yaml
# Root, because the agent launches engines and needs device access. Engines
# outlive a restart of this unit and are adopted again on the next start, so
# KillMode must not take the whole process group with it.
KillMode=process
Restart=on-failure
RestartSec=2s
StateDirectory=llmgw-agent

[Install]
WantedBy=multi-user.target
UNIT
scp -q "$TMP/llmgw-agent.service" "$SSH_USER@$ADDR:/tmp/"

# An engine path that is wrong, or right but unrunnable, is the failure this
# node type is most likely to have, and it is invisible: the agent registers
# healthy and every shard placed on it dies with "executable file not found"
# or a missing shared library. Checked here because it is cheap here and
# expensive later — llama-server on node05 was present, executable, and could
# not start for want of libgomp1, and the node reported ready throughout.
check_engine() {
  [ -n "$2" ] || return 0
  if $SSH "test -x '$2'" 2>/dev/null; then
    if $SSH "'$2' --version >/dev/null 2>&1 || '$2' --help >/dev/null 2>&1"; then
      say "$1: $2 runs"
      return 0
    fi
    ENGINE_PROBLEMS="$ENGINE_PROBLEMS
  $1 at $2 exists but will not run. Usually a missing shared library:
    ssh $SSH_USER@$ADDR ldd $2 | grep 'not found'"
    return 0
  fi
  ENGINE_PROBLEMS="$ENGINE_PROBLEMS
  $1 is configured at $2, which is not an executable on $NODE."
}
ENGINE_PROBLEMS=
if [ -n "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
  say "checking the engines"
  check_engine llamacpp "$LLAMACPP_BIN"
  check_engine vllm "$VLLM_BIN"
  check_engine sglang "$SGLANG_PYTHON"
fi

say "starting the agent"
$SSH 'set -e
  sudo install -m 0644 -o root -g root /tmp/agent.yaml /etc/llmgw/agent.yaml
  sudo install -m 0644 /tmp/llmgw-agent.service /etc/systemd/system/llmgw-agent.service
  rm -f /tmp/agent.yaml /tmp/llmgw-agent.service
  sudo systemctl daemon-reload
  sudo systemctl enable --now llmgw-agent >/dev/null
  sudo systemctl restart llmgw-agent'

echo
if [ -z "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
  say "the agent now installs the engines the control plane asks for: the"
  say "build for this machine, checked against its publisher's checksum and run"
  say "before the node is offered any model. 'llmgw ctl nodes' shows progress."
fi
if [ -n "$ENGINE_PROBLEMS" ]; then
  echo "$NODE is onboarded but an engine will not run:$ENGINE_PROBLEMS"
  echo
  echo "The node will register as ready regardless. Fix this before a replica"
  echo "is placed on it, or the shard fails and the reason is three logs away."
  echo
fi
echo "$NODE onboarded. It should appear in 'llmgw ctl nodes' within a few seconds."
echo "If it does not, the control plane's log says why — a certificate whose"
echo "common name does not match the claimed node id is refused by design."