#!/usr/bin/env bash
# Join the Ray cluster as an external worker.
#
# The Ray head's GCS port must be reachable from this machine.
# Set RAY_HEAD_ADDRESS or pass it as the first argument.
#
# Usage:
#   ./scripts/ray-join.sh                          # uses RAY_HEAD_ADDRESS env
#   ./scripts/ray-join.sh 192.168.100.50:6379      # explicit address
#   RAY_HEAD_ADDRESS=192.168.100.50:6379 ./scripts/ray-join.sh
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
COMFYUI_DIR="$PROJECT_DIR/ComfyUI"

RAY_HEAD="${1:-${RAY_HEAD_ADDRESS:-}}"

if [[ -z "$RAY_HEAD" ]]; then
    echo "ERROR: Ray head address required."
    echo "Usage: $0 <ray-head-ip>:<port>"
    echo "   or: RAY_HEAD_ADDRESS=<ip>:<port> $0"
    exit 1
fi

# Activate ComfyUI venv (where Ray is installed)
# shellcheck disable=SC1091
source "$COMFYUI_DIR/.venv/bin/activate"

# Stop any existing Ray worker on this machine
ray stop --force 2>/dev/null || true

echo "Joining Ray cluster at $RAY_HEAD..."

ray start \
    --address="$RAY_HEAD" \
    --num-cpus=16 \
    --num-gpus=1 \
    --resources='{"3d_gen": 1, "rtx4070": 1}' \
    --node-name=desktop

echo ""
echo "Desktop joined Ray cluster. Verify with: ray status"
echo "To disconnect: ray stop"