#!/usr/bin/env bash
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
ENV_FILE="$SCRIPT_DIR/.env.hy3"
COMPOSE_FILE="$SCRIPT_DIR/docker-compose.hy3.yml"
API_URL="http://127.0.0.1:8888/v1/models"
WAIT_ATTEMPTS="${WAIT_ATTEMPTS:-120}"
WAIT_SECONDS="${WAIT_SECONDS:-15}"

set -a; source "$ENV_FILE"; set +a
cd "$SCRIPT_DIR"
WORKER_DIR="${WORKER_DIR:-$SCRIPT_DIR}"
REMOTE="cd $WORKER_DIR && env -u MASTER_ADDR -u MASTER_PORT -u NODE_RANK -u HEADLESS COMPOSE_DISABLE_ENV_FILE=1"

echo "Syncing Hy3 files to ${WORKER_HOST}:${WORKER_DIR}"
ssh "$WORKER_HOST" "mkdir -p $WORKER_DIR"
scp -q "$COMPOSE_FILE" "${WORKER_HOST}:${WORKER_DIR}/docker-compose.hy3.yml"
scp -q "$ENV_FILE" "${WORKER_HOST}:${WORKER_DIR}/.env.hy3"

echo "Starting Hy3 worker (rank 1) on ${WORKER_HOST}..."
ssh "$WORKER_HOST" "$REMOTE NODE_RANK=1 HEADLESS=1 VLLM_HOST_IP='$WORKER_VLLM_HOST_IP' docker compose --env-file .env.hy3 -f docker-compose.hy3.yml up -d"

echo "Starting Hy3 head (rank 0)..."
COMPOSE_DISABLE_ENV_FILE=1 NODE_RANK=0 HEADLESS= docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" up -d

echo "Waiting for Hy3 vLLM API (up to $((WAIT_ATTEMPTS*WAIT_SECONDS))s)..."
for _ in $(seq 1 "$WAIT_ATTEMPTS"); do
  if curl -fsS --max-time 5 "$API_URL" >/dev/null 2>&1; then
    echo "Hy3 vLLM is UP: $API_URL"; exit 0
  fi
  sleep "$WAIT_SECONDS"
done
echo "TIMEOUT. head logs:" >&2
COMPOSE_DISABLE_ENV_FILE=1 docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" logs --tail=80 vllm-hy3 2>&1 | tail -80
exit 1