#!/usr/bin/env bash
# Purpose: run the same workload against two endpoints, one on each side of the
#          comparison, at both the single-user and the batch concurrency, so that the
#          reality check has four measurements taken by one tool with one prompt set
#          rather than four numbers from four sources. Every run is labelled with the
#          side and the concurrency and appended to the lab notebook.
# Platform: all (spark, strix, mac, nvidia). The script itself only makes HTTP
#           requests, so it can run on a laptop while both machines serve.
# Minimum memory: none on the machine running this; the memory floors belong to the
#                 two machines being compared, and each side's own page states them.
# Assumes: two OpenAI-compatible endpoints already serving, python3 on PATH,
#          measure-pair.py beside this script, and an .env copied from
#          env-example.txt with SIDE_A_URL, SIDE_B_URL and their names filled in.
#          Both endpoints must be serving the model each machine can actually hold;
#          that they are different models is the point, not a mistake.
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

: "${SIDE_A_NAME:=side-a}"
: "${SIDE_B_NAME:=side-b}"
: "${SIDE_A_URL:=}"
: "${SIDE_B_URL:=}"
: "${SIDE_A_MODEL:=}"
: "${SIDE_B_MODEL:=}"
: "${BATCH_CONCURRENCY:=8}"
: "${REQUESTS:=20}"
: "${MAX_TOKENS:=128}"
: "${PREFILL_TOKENS:=2048}"
: "${LABBOOK:=${HERE}/labbook.md}"

if [ -z "$SIDE_A_URL" ] || [ -z "$SIDE_B_URL" ]; then
    echo "SIDE_A_URL and SIDE_B_URL must both be set in ${ENV_FILE}." >&2
    echo "Each is the /v1 base of one machine's OpenAI-compatible endpoint." >&2
    exit 1
fi

if ! command -v python3 >/dev/null 2>&1; then
    echo "python3 is not on PATH." >&2
    exit 1
fi

BATCH_REQUESTS=$((REQUESTS * 2))

run_side() {
    # run_side <name> <base-url> <model-or-empty> <concurrency> <requests> <suffix>
    local name="$1" url="$2" model="$3" conc="$4" reqs="$5" suffix="$6"
    local model_args=()
    if [ -n "$model" ]; then
        model_args=(--model "$model")
    fi
    echo ""
    echo "==> ${name}, concurrency ${conc}, ${reqs} requests"
    python3 "${HERE}/measure-pair.py" \
        --base-url "$url" \
        --concurrency "$conc" \
        --requests "$reqs" \
        --max-tokens "$MAX_TOKENS" \
        --prefill-tokens "$PREFILL_TOKENS" \
        --label "${name}-${suffix}" \
        --labbook "$LABBOOK" \
        "${model_args[@]+"${model_args[@]}"}"
}

echo "Comparing ${SIDE_A_NAME} against ${SIDE_B_NAME}."
echo "Single-user runs answer 'which feels faster to one person'."
echo "Batch runs answer 'which serves more people at once'. They can disagree."

run_side "$SIDE_A_NAME" "$SIDE_A_URL" "$SIDE_A_MODEL" 1 "$REQUESTS" "single-user"
run_side "$SIDE_B_NAME" "$SIDE_B_URL" "$SIDE_B_MODEL" 1 "$REQUESTS" "single-user"
run_side "$SIDE_A_NAME" "$SIDE_A_URL" "$SIDE_A_MODEL" \
    "$BATCH_CONCURRENCY" "$BATCH_REQUESTS" "batch"
run_side "$SIDE_B_NAME" "$SIDE_B_URL" "$SIDE_B_MODEL" \
    "$BATCH_CONCURRENCY" "$BATCH_REQUESTS" "batch"

echo ""
echo "Four measurements recorded in ${LABBOOK}. Read them back with:"
echo "  grep 'served-measurement' ${LABBOOK}"
echo ""
echo "Take aggregate_output_tokens_per_second from the batch runs and the median"
echo "time to first token from the single-user runs into the results table, then"
echo "feed the batch throughput into cost-per-million-tokens.py."