#!/usr/bin/env bash
# Purpose: drive an already-running exo cluster through its API: show which nodes it
#          found, list the placements it will accept for a model, create one, wait
#          until it is ready, and delete it afterwards. exo itself is started by hand
#          in its own terminal on each Mac, because it runs in the foreground and its
#          log is worth watching the first few times.
# Platform: mac (Track M). exo runs on Linux too, on CPU only as its README states,
#           which is why the other tracks use Part 19's llama.cpp RPC path instead.
# Minimum memory: 32 GB per Mac for the primary path; 24 GB for a single Mac with a
#                 model that fits on it.
# Assumes: exo running on this Mac and on the peer, curl and python3 on PATH, and an
#          .env copied from env-example.txt. Nothing here starts or stops exo itself.
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

: "${EXO_PORT:=52415}"
: "${EXO_DIR:=}"
: "${EXO_MODEL:=}"
: "${EXO_PLACEMENT_INDEX:=0}"

BASE="http://localhost:${EXO_PORT}"
MODE="${1:-status}"

usage() {
    cat <<'USAGE'
Usage: bash run-exo.sh [start-help|status|previews|place|delete <instance-id>]

  start-help  print the command that starts exo on each Mac, and stop
  status      the cluster exo currently sees, from GET /state
  previews    every placement exo will accept for EXO_MODEL, with its cost
  place       create the placement at EXO_PLACEMENT_INDEX and wait until ready
  delete      remove an instance and free the memory it holds
USAGE
}

require() {
    command -v "$1" >/dev/null 2>&1 || { echo "$1 is not on PATH." >&2; exit 1; }
}
require curl
require python3

exo_up() {
    curl -fsS --max-time 5 "${BASE}/node_id" >/dev/null 2>&1
}

if [ "$MODE" = "start-help" ] || [ "$MODE" = "-h" ] || [ "$MODE" = "--help" ]; then
    usage
    echo ""
    echo "Start exo on EACH Mac, in its own terminal, and leave it running:"
    if [ -n "$EXO_DIR" ]; then
        echo "  cd ${EXO_DIR} && uv run exo"
    else
        echo "  cd <your exo checkout> && uv run exo"
        echo "  (or launch the macOS application, which needs macOS 26.2 or later)"
    fi
    echo ""
    echo "Then come back here. The dashboard is at ${BASE}/"
    exit 0
fi

if ! exo_up; then
    echo "Nothing answers at ${BASE}. Start exo first:" >&2
    echo "  bash run-exo.sh start-help" >&2
    exit 1
fi

case "$MODE" in
    status)
        echo "==> node id"
        curl -fsS "${BASE}/node_id"
        echo ""
        echo "==> cluster state"
        curl -fsS "${BASE}/state" | python3 -m json.tool
        ;;

    previews)
        if [ -z "$EXO_MODEL" ]; then
            echo "EXO_MODEL is empty in ${ENV_FILE}." >&2
            exit 1
        fi
        echo "==> placements exo will accept for ${EXO_MODEL}"
        curl -fsS --get "${BASE}/instance/previews" \
            --data-urlencode "model_id=${EXO_MODEL}" \
            | python3 "${HERE}/exo-placement.py" --list
        echo ""
        echo "Set EXO_PLACEMENT_INDEX in ${ENV_FILE} to the number in brackets."
        ;;

    place)
        if [ -z "$EXO_MODEL" ]; then
            echo "EXO_MODEL is empty in ${ENV_FILE}." >&2
            exit 1
        fi
        echo "==> creating placement ${EXO_PLACEMENT_INDEX} for ${EXO_MODEL}"
        BODY="$(curl -fsS --get "${BASE}/instance/previews" \
            --data-urlencode "model_id=${EXO_MODEL}" \
            | python3 "${HERE}/exo-placement.py" --select "$EXO_PLACEMENT_INDEX")"
        curl -fsS -X POST "${BASE}/instance" \
            -H 'Content-Type: application/json' \
            -d "$BODY" | python3 -m json.tool
        echo ""
        echo "==> waiting until exo reports the instance ready"
        echo "    A first placement downloads weights, so this can take a while."
        curl -fsS -N --get "${BASE}/instance/await" \
            --data-urlencode "model_id=${EXO_MODEL}" \
            | sed -n '1,20p'
        echo ""
        echo "Ready. Measure it with:"
        echo "  python3 measure-pair.py --base-url ${BASE}/v1 \\"
        echo "      --model ${EXO_MODEL} --label exo --exo-bench-url ${BASE}"
        ;;

    delete)
        INSTANCE_ID="${2:-}"
        if [ -z "$INSTANCE_ID" ]; then
            echo "Usage: bash run-exo.sh delete <instance-id>" >&2
            echo "Find the id with: bash run-exo.sh status" >&2
            exit 1
        fi
        echo "==> deleting instance ${INSTANCE_ID}"
        curl -fsS -X DELETE "${BASE}/instance/${INSTANCE_ID}"
        echo ""
        echo "Memory freed on every node that held a shard."
        ;;

    *)
        usage >&2
        exit 1
        ;;
esac