#!/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
  --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: --addr)
  --mgmt-if NAME    interface NCCL bootstraps over (default eth0)
  --pki DIR         where the authority lives (default /etc/llmgw/pki)
  --binary PATH     llmgw binary to install (default: the one running this)
  --ssh-user NAME   (default ilker)
USAGE
  exit 2
}

NODE= ADDR= SERVER= MGMT= MGMT_IF=eth0
PKI=/etc/llmgw/pki BINARY= SSH_USER=ilker
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="$2"; shift 2;;
    --mgmt-if) MGMT_IF="$2"; shift 2;;
    --pki) PKI="$2"; shift 2;;
    --binary) BINARY="$2"; shift 2;;
    --ssh-user) SSH_USER="$2"; shift 2;;
    *) usage;;
  esac
done
[ -n "$NODE" ] && [ -n "$ADDR" ] && [ -n "$SERVER" ] || usage
MGMT="${MGMT:-$ADDR}"
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; }

SSH="ssh -o BatchMode=yes -o ConnectTimeout=10 $SSH_USER@$ADDR"
say() { printf '  %s\n' "$*"; }

echo "onboarding $NODE at $ADDR"
$SSH true || { echo "cannot reach $SSH_USER@$ADDR over ssh" >&2; exit 1; }

# 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 "$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.
say "installing the binary"
scp -q "$BINARY" "$SSH_USER@$ADDR:/tmp/llmgw"
$SSH 'sudo install -m 0755 /tmp/llmgw /usr/local/bin/llmgw && rm -f /tmp/llmgw'

say "writing the configuration"
FABRICS="fabrics: []"
if [ -n "$FABRIC_NAME" ]; then
  FABRICS="fabrics:
  - name: $FABRIC_NAME
    hcas: [$(echo "$FABRIC_HCAS" | sed 's/,/, /g')]"
fi
cat > "$TMP/agent.yaml" <<YAML
# $NODE
#
# Written by onboard-node.sh. Fabric membership is configuration and there is no
# way around it: OpenSM's default subnet prefix is the same on every fabric, so
# a matching prefix proves nothing and a wrong guess produces a replica that
# starts, passes health checks and then hangs on its first all-reduce.
node_id: $NODE
server: "$SERVER"

tls:
  ca: /etc/llmgw/ca.pem
  cert: /etc/llmgw/agent.pem
  key: /etc/llmgw/agent-key.pem

mgmt_addr: "$MGMT"
mgmt_interface: $MGMT_IF

heartbeat_interval: 5s
reconnect_base: 1s
reconnect_max: 30s

$FABRICS

discovery:
  prefer_nvml: false
  skip_gpus: false

supervisor:
  work_dir: /var/lib/llmgw-agent
  log_lines: 512
  stop_timeout: 60s
  health_interval: 2s
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/"

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
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."