#!/usr/bin/env bash
# Purpose: check that two Macs are ready to be clustered over Thunderbolt, then ask
#          mlx.distributed_config to work out the topology and write the host file.
#          It changes nothing by itself: every command that would touch a network
#          interface is printed for you to read and run, because taking the
#          Thunderbolt bridge down is a change to how the machine networks.
# Platform: mac (Track M) only. Needs macOS 26 or later on both machines, and macOS
#           26.2 or later plus Thunderbolt 5 for the jaccl backend.
# Minimum memory: 32 GB per Mac for the primary path; 24 GB for the single-Mac path.
# Assumes: mlx and mlx-lm installed in the python on PATH, an .env beside this script
#          copied from env-example.txt, password-less ssh to PEER_SSH already working
#          (Part 18's lab sets that up), and `dot` from Graphviz if you want the
#          topology drawing. Writes the host file named by HOSTFILE.
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:=}"
: "${PEER_SSH:=}"
: "${SELF_SSH:=}"
: "${MLX_BACKEND:=ring}"
: "${HOSTFILE:=hosts.json}"

if [ "$(uname -s)" != "Darwin" ]; then
    echo "This script is Track M only: it reads macOS-specific tools." >&2
    echo "On another track, the lab's fallback section points at Part 19." >&2
    exit 1
fi

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

echo "==> 1. This machine"
sw_vers
sysctl -n machdep.cpu.brand_string || true
echo ""

OS_VERSION="$(sw_vers -productVersion)"
OS_MAJOR="${OS_VERSION%%.*}"
OS_REST="${OS_VERSION#*.}"
OS_MINOR="${OS_REST%%.*}"
if [ "$OS_MINOR" = "$OS_VERSION" ]; then
    OS_MINOR=0
fi

echo "==> 2. Thunderbolt ports and the interfaces they map to"
echo "    system_profiler names the ports; networksetup names the interfaces."
system_profiler SPThunderboltDataType | sed -n '1,60p'
echo ""
networksetup -listallhardwareports | grep -A2 -i thunderbolt || \
    echo "    No hardware port with Thunderbolt in its name. Check the cable."
echo ""

echo "==> 3. The Thunderbolt bridge, which the cluster does not want"
echo "    macOS presents Thunderbolt cables as one bridged service. MLX takes them"
echo "    apart into per-cable subnets, and says the bridge must be down even for"
echo "    the RDMA backend, which does not use TCP/IP for model traffic at all."
ifconfig bridge0 2>/dev/null | sed -n '1,4p' || echo "    No bridge0 on this machine."
echo ""

if [ "$MLX_BACKEND" = "jaccl" ]; then
    echo "==> 4. RDMA devices"
    if [ "$OS_MAJOR" -lt 26 ] || { [ "$OS_MAJOR" -eq 26 ] && [ "$OS_MINOR" -lt 2 ]; }; then
        echo "    macOS ${OS_VERSION} is older than 26.2, which is the documented" >&2
        echo "    floor for RDMA over Thunderbolt. Use MLX_BACKEND=ring." >&2
        exit 1
    fi
    if ! command -v ibv_devices >/dev/null 2>&1; then
        echo "    ibv_devices is not on PATH; RDMA has not been enabled here." >&2
        echo "    Enable it once from macOS Recovery with 'rdma_ctl enable', reboot," >&2
        echo "    and run this script again. It cannot be done over ssh." >&2
        exit 1
    fi
    ibv_devices
    echo ""
else
    echo "==> 4. RDMA devices: skipped, MLX_BACKEND is ${MLX_BACKEND}"
    echo ""
fi

if [ -z "$PEER_SSH" ]; then
    echo "==> 5. PEER_SSH is empty: treating this as the single-Mac path."
    echo "    No host file is needed. Launch two ranks on this machine with:"
    echo "      mlx.launch -n 2 -- python check-group.py"
    exit 0
fi

echo "==> 5. The peer"
SSH_TARGET="$PEER_SSH"
if [ -n "${SSH_USER:-}" ]; then
    SSH_TARGET="${SSH_USER}@${PEER_SSH}"
fi
if ! ssh -o BatchMode=yes -o ConnectTimeout=10 "$SSH_TARGET" true; then
    echo "    Cannot ssh to ${PEER_SSH} without a password." >&2
    echo "    mlx.launch starts every rank over ssh, so fix this first." >&2
    exit 1
fi
echo "    ssh to ${PEER_SSH} works with no password."
ssh "$SSH_TARGET" sw_vers -productVersion | sed 's/^/    peer macOS: /'
echo ""

echo "==> 6. Both Macs, as mlx.distributed_config sees them"
HOSTS="${SELF_SSH:-localhost},${PEER_SSH}"
echo "    hosts: ${HOSTS}"
mlx.distributed_config --verbose --hosts "$HOSTS" --over thunderbolt --dot \
    > "${HERE}/topology.dot"
echo "    Topology written to ${HERE}/topology.dot."
if command -v dot >/dev/null 2>&1; then
    echo "    Render it with:  dot -Tpng topology.dot -o topology.png"
else
    echo "    Install Graphviz to render it: brew install graphviz"
fi
echo ""

echo "==> 7. The configuration commands, printed rather than run"
echo "    Read every line before you run any of them. They bring the bridge"
echo "    interface down and give each cable's interface its own tiny subnet."
echo "    Answer the helper's prompt only once you have run them on that machine."
mlx.distributed_config --verbose --hosts "$HOSTS" --over thunderbolt \
    --backend "$MLX_BACKEND" --output-hostfile "${HERE}/${HOSTFILE}"

echo ""
echo "Host file written to ${HERE}/${HOSTFILE}."
echo "Check it against hosts-example.json, then run:"
echo "  bash run-distributed.sh"