#!/usr/bin/env bash
# Purpose: measure throughput in both directions and round-trip time to every named
#          cluster peer, then append one JSON line per link to the lab notebook so the
#          numbers are dated, attributed to a machine and never retyped by hand
# Platform: all (spark, strix, mac, nvidia); each peer must already be running
#           `iperf3 --server` on the port in .env, which task 4 of the lab starts
# Minimum memory: 8 GB
# Assumes: iperf3 and python3 on PATH, an .env beside this script copied from
#          env-example.txt with MACHINE_NAME and CLUSTER_PEERS filled in, and every
#          peer reachable by the name written in CLUSTER_PEERS. Raw iperf3 output is
#          kept under link-results/ so a surprising number can be re-read later.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${ENV_FILE:-${HERE}/.env}"

if [ -f "$ENV_FILE" ]; then
    set -a
    # shellcheck source=/dev/null
    . "$ENV_FILE"
    set +a
else
    echo "No ${ENV_FILE}. Copy env-example.txt to .env and fill it in." >&2
    exit 1
fi

: "${MACHINE_NAME:=}"
: "${CLUSTER_PEERS:=}"
: "${IPERF_PORT:=5201}"
: "${IPERF_SECONDS:=20}"
: "${IPERF_STREAMS:=4}"
: "${PING_COUNT:=50}"
: "${LABBOOK:=${HERE}/labbook.md}"
: "${CLUSTER_IFACE:=}"

RESULTS_DIR="${RESULTS_DIR:-${HERE}/link-results}"

if [ -z "$MACHINE_NAME" ]; then
    echo "MACHINE_NAME is empty in ${ENV_FILE}. Give this machine a short name first." >&2
    exit 1
fi

if [ -z "$CLUSTER_PEERS" ]; then
    echo "CLUSTER_PEERS is empty in ${ENV_FILE}. Nothing to measure." >&2
    echo "On the single-machine path, set it to the container name of the iperf3 server." >&2
    exit 1
fi

for tool in iperf3 ping python3; do
    if ! command -v "$tool" >/dev/null 2>&1; then
        echo "${tool} is not on PATH." >&2
        exit 1
    fi
done

mkdir -p "$RESULTS_DIR"

read -r -a PEERS <<< "$CLUSTER_PEERS"

# Turns one peer's raw output into a single JSON line for the notebook. Kept in Python
# because iperf3's JSON is nested and awk would be a poor place to learn that.
summarise() {
    # summarise <peer> <forward-json> <reverse-json> <ping-text>
    python3 - "$1" "$2" "$3" "$4" <<'PY'
import json
import os
import re
import sys
from datetime import date, datetime, timezone
from pathlib import Path

peer, forward_path, reverse_path, ping_path = sys.argv[1:5]


def read_iperf(path):
    """Sender and receiver bit rates from one iperf3 --json run, or None."""
    try:
        data = json.loads(Path(path).read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    if data.get("error"):
        return {"error": str(data["error"])}
    end = data.get("end", {})
    sent = end.get("sum_sent", {})
    received = end.get("sum_received", {})
    out = {}
    if "bits_per_second" in sent:
        out["sent_bits_per_second"] = round(float(sent["bits_per_second"]))
    if "bits_per_second" in received:
        out["received_bits_per_second"] = round(float(received["bits_per_second"]))
    if sent.get("retransmits") is not None:
        out["retransmits"] = sent["retransmits"]
    return out or None


def read_ping(path):
    """min/avg/max round trip in milliseconds, from either the Linux or the BSD summary."""
    try:
        text = Path(path).read_text(encoding="utf-8")
    except OSError:
        return None
    m = re.search(r"=\s*([\d.]+)/([\d.]+)/([\d.]+)", text)
    loss = re.search(r"([\d.]+)%\s*packet loss", text)
    out = {}
    if m:
        out["rtt_min_ms"] = float(m.group(1))
        out["rtt_avg_ms"] = float(m.group(2))
        out["rtt_max_ms"] = float(m.group(3))
    if loss:
        out["packet_loss_percent"] = float(loss.group(1))
    return out or None


record = {
    "lab": "part-18/build-and-measure-your-cluster-network",
    "record": "link",
    "date": date.today().isoformat(),
    "recorded_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
    "from_machine": os.environ.get("MACHINE_NAME", ""),
    "to_peer": peer,
    "interface": os.environ.get("CLUSTER_IFACE", "") or None,
    "iperf3_seconds": int(os.environ.get("IPERF_SECONDS", "0") or 0),
    "iperf3_streams": int(os.environ.get("IPERF_STREAMS", "0") or 0),
    "forward": read_iperf(forward_path),
    "reverse": read_iperf(reverse_path),
    "ping": read_ping(ping_path),
}
print(json.dumps(record, sort_keys=True))
PY
}

echo "==> measuring ${#PEERS[@]} link(s) from ${MACHINE_NAME}"
echo "    results kept in ${RESULTS_DIR}, notebook lines appended to ${LABBOOK}"
echo ""

FAILURES=0

for peer in "${PEERS[@]}"; do
    safe="${peer//[^A-Za-z0-9._-]/_}"
    forward="${RESULTS_DIR}/${safe}-forward.json"
    reverse="${RESULTS_DIR}/${safe}-reverse.json"
    pingout="${RESULTS_DIR}/${safe}-ping.txt"

    echo "--> ${peer}: sending"
    if ! iperf3 --client "$peer" --port "$IPERF_PORT" --time "$IPERF_SECONDS" \
        --parallel "$IPERF_STREAMS" --json > "$forward"; then
        echo "    iperf3 failed sending to ${peer}; is 'iperf3 --server' running there?" >&2
        FAILURES=$((FAILURES + 1))
    fi

    echo "--> ${peer}: receiving"
    if ! iperf3 --client "$peer" --port "$IPERF_PORT" --time "$IPERF_SECONDS" \
        --parallel "$IPERF_STREAMS" --reverse --json > "$reverse"; then
        echo "    iperf3 failed receiving from ${peer}." >&2
        FAILURES=$((FAILURES + 1))
    fi

    echo "--> ${peer}: round trip"
    if ! ping -c "$PING_COUNT" -q "$peer" > "$pingout" 2>&1; then
        echo "    ping failed for ${peer}; the name may not resolve." >&2
        FAILURES=$((FAILURES + 1))
    fi

    summarise "$peer" "$forward" "$reverse" "$pingout" | tee -a "$LABBOOK"
    echo ""
done

if [ "$FAILURES" -gt 0 ]; then
    echo "${FAILURES} step(s) failed. The notebook lines above show which fields are missing." >&2
    exit 1
fi

echo "All links measured. Read the lines back with:"
echo "  grep '\"record\": \"link\"' ${LABBOOK}"