#!/usr/bin/env bash
# Purpose: launch one model across the cluster with mlx.launch, either as a one-shot
#          generation that records its own numbers or as an OpenAI-compatible server
#          that measure-pair.py can then drive. It builds the mlx.launch command from
#          .env so that the ring and jaccl backends, and the tensor and pipeline
#          splits, are one variable apart rather than one retyped command apart.
# Platform: mac (Track M). With PEER_SSH empty it launches both ranks on this Mac,
#           which is the single-machine path.
# Minimum memory: 32 GB per Mac for the primary path; 24 GB for the single-Mac path.
# Assumes: mlx-lm installed at the same python path on both Macs, this directory
#          present at the same absolute path on both, a host file written by
#          setup-thunderbolt-bridge.sh, and an .env copied from env-example.txt.
#          Serving mode runs in the foreground until you stop it with Ctrl-C.
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

: "${PEER_SSH:=}"
: "${MLX_BACKEND:=ring}"
: "${HOSTFILE:=hosts.json}"
: "${RANKS:=2}"
: "${CLUSTER_PYTHON:=}"
: "${CLUSTER_DIR:=$HERE}"
: "${MODEL_REPO:=}"
: "${BASELINE_REPO:=}"
: "${USE_PIPELINE:=}"
: "${SERVER_HOST:=127.0.0.1}"
: "${SERVER_PORT:=8080}"
: "${MAX_TOKENS:=256}"
: "${LABBOOK:=${HERE}/labbook.md}"

MODE="${1:-generate}"

usage() {
    cat <<'USAGE'
Usage: bash run-distributed.sh [generate|serve|baseline]

  generate   one distributed completion, numbers appended to the lab notebook
  serve      mlx_lm.server across the cluster, for measure-pair.py to drive
  baseline   one completion of BASELINE_REPO on this Mac alone, for the first row
USAGE
}

case "$MODE" in
    generate|serve|baseline) ;;
    -h|--help) usage; exit 0 ;;
    *) usage >&2; exit 1 ;;
esac

if ! command -v mlx.launch >/dev/null 2>&1; then
    echo "mlx.launch is not on PATH. Install mlx-lm in this python first." >&2
    exit 1
fi

if [ -z "$CLUSTER_PYTHON" ]; then
    CLUSTER_PYTHON="$(mlx.launch --print-python)"
    echo "==> using the python mlx.launch reports: ${CLUSTER_PYTHON}"
    echo "    Confirm the same path exists on the other Mac before going further."
fi

# ------------------------------------------------------- the launcher arguments
LAUNCH=(mlx.launch --verbose --backend "$MLX_BACKEND")

if [ -n "$PEER_SSH" ]; then
    if [ ! -f "${HERE}/${HOSTFILE}" ]; then
        echo "No ${HERE}/${HOSTFILE}. Run setup-thunderbolt-bridge.sh first." >&2
        exit 1
    fi
    LAUNCH+=(--hostfile "${HERE}/${HOSTFILE}")
else
    echo "==> PEER_SSH is empty: single-machine path, ${RANKS} ranks on this Mac."
    echo "    The mechanism is the same; the memory is not."
    LAUNCH+=(-n "$RANKS")
fi

LAUNCH+=(--cwd "$CLUSTER_DIR")

SPLIT_ARGS=()
if [ -n "$USE_PIPELINE" ]; then
    SPLIT_ARGS+=(--pipeline)
fi

LABEL="${MLX_BACKEND}-$([ -n "$USE_PIPELINE" ] && echo pipeline || echo tensor)"

# --------------------------------------------------------------------- baseline
if [ "$MODE" = "baseline" ]; then
    if [ -z "$BASELINE_REPO" ]; then
        echo "BASELINE_REPO is empty in ${ENV_FILE}." >&2
        exit 1
    fi
    echo "==> one Mac, ${BASELINE_REPO}, no distributed group"
    "$CLUSTER_PYTHON" "${CLUSTER_DIR}/sharded-generate.py" \
        --model "$BASELINE_REPO" \
        --max-tokens "$MAX_TOKENS" \
        --label "one-mac-baseline" \
        --labbook "$LABBOOK"
    exit 0
fi

if [ -z "$MODEL_REPO" ]; then
    echo "MODEL_REPO is empty in ${ENV_FILE}." >&2
    exit 1
fi

# --------------------------------------------------------------------- generate
if [ "$MODE" = "generate" ]; then
    echo "==> ${MODEL_REPO} across the group, backend ${MLX_BACKEND}, split ${LABEL}"
    "${LAUNCH[@]}" -- \
        "$CLUSTER_PYTHON" "${CLUSTER_DIR}/sharded-generate.py" \
        --model "$MODEL_REPO" \
        --max-tokens "$MAX_TOKENS" \
        --label "$LABEL" \
        --labbook "$LABBOOK" \
        "${SPLIT_ARGS[@]+"${SPLIT_ARGS[@]}"}"
    exit 0
fi

# ------------------------------------------------------------------------ serve
echo "==> serving ${MODEL_REPO} across the group on ${SERVER_HOST}:${SERVER_PORT}"
echo "    Rank 0 answers HTTP; the other ranks do their share and stay quiet."
echo "    Leave this running and drive it from another terminal with:"
echo "      python3 measure-pair.py --label ${LABEL}"
echo "    Stop it with Ctrl-C when the measurement is recorded."
"${LAUNCH[@]}" -- \
    "$CLUSTER_PYTHON" -m mlx_lm.server \
    --model "$MODEL_REPO" \
    --host "$SERVER_HOST" \
    --port "$SERVER_PORT" \
    "${SPLIT_ARGS[@]+"${SPLIT_ARGS[@]}"}"