#!/usr/bin/env bash
# Purpose: start one ggml-rpc-server on this machine, bound to the cluster link and nothing
#          else, with the local tensor cache enabled, and print exactly what it exposes
# Platform: all (CUDA, Vulkan, Metal or CPU builds; the backend is whatever this host built)
# Minimum memory: 8 GB
# Assumes: llama.cpp built with -DGGML_RPC=ON so that ggml-rpc-server exists in $LLAMA_BIN;
#          the interface named in CLUSTER_IFACE exists and has an IPv4 address; nothing else
#          is listening on RPC_PORT; the environment from env-example.txt has been sourced
#
# Usage: . ./rpc.env && bash start-rpc-server.sh
#        . ./rpc.env && bash start-rpc-server.sh --background
#
# Environment (all from env-example.txt):
#   CLUSTER_IFACE  interface carrying cluster traffic on this host   (required)
#   CLUSTER_ADDR   address to bind to, overriding the interface      (default: read from CLUSTER_IFACE)
#   RPC_PORT       port to listen on                                 (default: 50052)
#   RPC_DEVICE     devices to expose, e.g. CUDA0 or Vulkan0 or CPU   (default: every accelerator)
#   RPC_THREADS    CPU threads for the CPU device                    (default: the tool's own)
#   LLAMA_BIN      directory holding ggml-rpc-server                 (default: found on PATH)
#   RPC_NO_RDMA    set to 1 to force plain TCP for a comparison run  (default: unset)
#
# The RPC server has no authentication and no encryption, and its own README says so in
# capitals: never run it on an open network. This script therefore refuses to bind to a
# wildcard address, and prints the one address it did bind to so you can check it.

set -euo pipefail

BACKGROUND=0
[ "${1:-}" = "--background" ] && BACKGROUND=1

RPC_PORT="${RPC_PORT:-50052}"
CLUSTER_ADDR="${CLUSTER_ADDR:-}"
RPC_DEVICE="${RPC_DEVICE:-}"
RPC_THREADS="${RPC_THREADS:-}"

die() { echo "start-rpc-server: $*" >&2; exit 1; }

# --- 1. Find the binary ----------------------------------------------------------------
if [ -n "${LLAMA_BIN:-}" ]; then
  SERVER="$LLAMA_BIN/ggml-rpc-server"
else
  SERVER="$(command -v ggml-rpc-server || true)"
fi
[ -n "$SERVER" ] && [ -x "$SERVER" ] || die \
  "ggml-rpc-server not found. Set LLAMA_BIN, and check the build was configured with -DGGML_RPC=ON."

# --- 2. Work out which address to bind to ----------------------------------------------
if [ -z "$CLUSTER_ADDR" ]; then
  [ -n "${CLUSTER_IFACE:-}" ] || die "set CLUSTER_IFACE (or CLUSTER_ADDR); see env-example.txt"
  case "$(uname -s)" in
    Darwin)
      CLUSTER_ADDR="$(ipconfig getifaddr "$CLUSTER_IFACE" || true)"
      ;;
    *)
      CLUSTER_ADDR="$(ip -4 -o addr show dev "$CLUSTER_IFACE" 2>/dev/null \
        | awk '{ print $4 }' | cut -d/ -f1 | head -n 1)"
      ;;
  esac
fi
[ -n "$CLUSTER_ADDR" ] || die "$CLUSTER_IFACE has no IPv4 address; bring the link up first (Part 18)"

case "$CLUSTER_ADDR" in
  0.0.0.0|"*"|"")
    die "refusing to bind to every interface. Name the cluster link in CLUSTER_IFACE instead."
    ;;
esac

# --- 3. Say what is about to happen ------------------------------------------------------
echo "==> ggml-rpc-server"
echo "    binary     : $SERVER"
echo "    binding to : $CLUSTER_ADDR port $RPC_PORT (interface ${CLUSTER_IFACE:-set explicitly})"
echo "    cache      : on (-c); large tensors are kept on local disk instead of re-sent"
[ -n "$RPC_DEVICE" ] && echo "    devices    : $RPC_DEVICE" || echo "    devices    : every accelerator this build can see"
[ -n "$RPC_THREADS" ] && echo "    threads    : $RPC_THREADS"
if [ -n "${RPC_NO_RDMA:-}" ]; then
  echo "    transport  : TCP forced (GGML_RPC_NO_RDMA is set); unset it to allow RDMA"
else
  echo "    transport  : negotiated at handshake; RDMA if both peers can, TCP otherwise"
fi

# --- 4. Build the argument list ----------------------------------------------------------
ARGS=(-H "$CLUSTER_ADDR" -p "$RPC_PORT" -c)
[ -n "$RPC_DEVICE" ] && ARGS+=(-d "$RPC_DEVICE")
[ -n "$RPC_THREADS" ] && ARGS+=(-t "$RPC_THREADS")

if [ -n "${RPC_NO_RDMA:-}" ]; then
  export GGML_RPC_NO_RDMA=1
fi

# --- 5. Run it ---------------------------------------------------------------------------
if [ "$BACKGROUND" = "1" ]; then
  LOG="rpc-server-$(hostname -s)-${RPC_PORT}.log"
  nohup "$SERVER" "${ARGS[@]}" > "$LOG" 2>&1 &
  echo "    pid        : $! (log: $LOG)"
  echo "    stop it with: kill $!"
else
  echo "    press Ctrl-C to stop it"
  exec "$SERVER" "${ARGS[@]}"
fi