#!/usr/bin/env bash
# Purpose: run Part 9's load generator against one endpoint with the settings from .env,
#          reading this machine's cluster-interface byte counters either side of the run,
#          and append one notebook line saying how many bytes crossed the link to produce
#          those tokens
# Platform: all (Linux reads /proc/net/dev; macOS reads netstat -ib)
# Minimum memory: 1 GB on the machine running the generator; it may be a third machine
# Assumes: python3 on PATH; LOAD_TEST in .env pointing at Part 9's load-test.py; a server
#          already answering at the URL given; CLUSTER_IFACE naming an interface on THIS
#          machine. Counters cover the whole interface, so run it on a quiet cluster and
#          read the byte figures as an order of magnitude, not an exact accounting.
#
# Usage: bash run-load.sh disagg   http://127.0.0.1:8000/v1
#        bash run-load.sh baseline http://127.0.0.1:8200/v1
set -euo pipefail

HERE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [ -f "${HERE}/.env" ]; then
    set -a
    # shellcheck disable=SC1091  # written by the reader from env-example.txt
    . "${HERE}/.env"
    set +a
fi

LABEL="${1:-}"
BASE_URL="${2:-}"

fail() { printf '%s\n' "$*" >&2; exit 1; }

[ -n "$LABEL" ] || fail "Usage: bash run-load.sh <label> <base-url ending in /v1>"
[ -n "$BASE_URL" ] || fail "Usage: bash run-load.sh <label> <base-url ending in /v1>"

LOAD_TEST="${LOAD_TEST:-../part-09-vllm-and-sglang/load-test.py}"
case "$LOAD_TEST" in
    /*) ;;
    *) LOAD_TEST="${HERE}/${LOAD_TEST}" ;;
esac
[ -f "$LOAD_TEST" ] || fail "LOAD_TEST does not point at a file: ${LOAD_TEST}. Set it in .env to Part 9's load-test.py."

command -v python3 >/dev/null 2>&1 || fail "python3 is not on PATH."

SERVED_NAME="${SERVED_NAME:-local-chat}"
CONCURRENCY="${CONCURRENCY:-1,4,8}"
REQUESTS="${REQUESTS:-24}"
MAX_TOKENS="${MAX_TOKENS:-128}"
PROMPT_SET="${PROMPT_SET:-shared-prefix}"
LABBOOK="${LABBOOK:-labbook.md}"
CLUSTER_IFACE="${CLUSTER_IFACE:-}"
CONNECTOR="${CONNECTOR:-shared}"
CTX="${CTX:-8192}"

# --- interface byte counters, before and after ----------------------------------------
# Prints "<received> <transmitted>" for one interface, or nothing if it cannot be read.
counters() {
    local iface="$1"
    [ -n "$iface" ] || return 0
    case "$(uname -s)" in
        Linux)
            awk -v want="${iface}:" '$1 == want { print $2, $10 }' /proc/net/dev
            ;;
        Darwin)
            netstat -ib | awk -v want="$iface" '$1 == want && $4 !~ /:/ { print $7, $10; exit }'
            ;;
        *)
            ;;
    esac
}

read -r RX_BEFORE TX_BEFORE <<<"$(counters "$CLUSTER_IFACE")"
START_EPOCH="$(date +%s)"

printf '==> %s against %s\n' "$LABEL" "$BASE_URL"
if [ -n "$CLUSTER_IFACE" ] && [ -n "${RX_BEFORE:-}" ]; then
    printf '    reading byte counters on %s\n' "$CLUSTER_IFACE"
else
    printf '    no interface counters (CLUSTER_IFACE unset or unreadable on this system)\n'
fi
printf '\n'

python3 "$LOAD_TEST" \
    --base-url "$BASE_URL" \
    --model "$SERVED_NAME" \
    --concurrency "$CONCURRENCY" \
    --requests "$REQUESTS" \
    --max-tokens "$MAX_TOKENS" \
    --prompt-set "$PROMPT_SET" \
    --label "$LABEL" \
    --engine vllm \
    --labbook "$LABBOOK"

read -r RX_AFTER TX_AFTER <<<"$(counters "$CLUSTER_IFACE")"
END_EPOCH="$(date +%s)"

RX_DELTA=""
TX_DELTA=""
if [ -n "${RX_BEFORE:-}" ] && [ -n "${RX_AFTER:-}" ]; then
    RX_DELTA=$(( RX_AFTER - RX_BEFORE ))
    TX_DELTA=$(( TX_AFTER - TX_BEFORE ))
fi

# --- one notebook line about the link, beside the generator's own lines ----------------
{
    printf '{"lab": "part-22/lab-two-machine-prefill-decode", "record": "link", '
    printf '"label": "%s", "base_url": "%s", "connector": "%s", ' "$LABEL" "$BASE_URL" "$CONNECTOR"
    printf '"context_length": %s, "prompt_set": "%s", "concurrency": "%s", "requests": %s, ' \
        "$CTX" "$PROMPT_SET" "$CONCURRENCY" "$REQUESTS"
    printf '"iface": "%s", "wall_s": %s, ' "$CLUSTER_IFACE" "$(( END_EPOCH - START_EPOCH ))"
    if [ -n "$RX_DELTA" ]; then
        printf '"rx_bytes": %s, "tx_bytes": %s, ' "$RX_DELTA" "$TX_DELTA"
    else
        printf '"rx_bytes": null, "tx_bytes": null, '
    fi
    printf '"recorded_at": "%s"}\n' "$(date +%Y-%m-%dT%H:%M:%S%z)"
} >>"$LABBOOK"

printf '\n'
if [ -n "$RX_DELTA" ]; then
    printf '==> Link during this run, on %s\n' "$CLUSTER_IFACE"
    printf '    received     %s bytes\n' "$RX_DELTA"
    printf '    transmitted  %s bytes\n' "$TX_DELTA"
    printf '    Divide the larger of these by the number of requests and compare against\n'
    printf '    the per-request key-value cache size the lesson had you compute. If it is\n'
    printf '    a thousand times smaller, no cache is crossing the link.\n'
else
    printf '==> No byte counters recorded. Set CLUSTER_IFACE in .env to measure the link.\n'
fi
printf '    Appended one link line to %s\n' "$LABBOOK"