#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STOP_SGLANG_FOR_COMFY_POOL="${STOP_SGLANG_FOR_COMFY_POOL:-0}"
WORKERS="${COMFY_WORKER_COUNT:-4}"

if [[ ! -d /opt/dlami/nvme/ComfyUI ]]; then
  echo "Missing ComfyUI at /opt/dlami/nvme/ComfyUI" >&2
  exit 1
fi

if ! mountpoint -q /opt/dlami/nvme; then
  echo "/opt/dlami/nvme is not mounted; refusing to run model workers on root disk" >&2
  exit 1
fi

if [[ "$STOP_SGLANG_FOR_COMFY_POOL" == "1" ]]; then
  sudo systemctl stop desineuron-sglang.service || true
fi

sudo systemctl stop comfyui.service || true
sudo systemctl disable comfyui.service || true

sudo install -m 0755 "$SCRIPT_DIR/desineuron-start-comfy-worker" /usr/local/bin/desineuron-start-comfy-worker
sudo install -m 0644 "$SCRIPT_DIR/comfyui-worker@.service" /etc/systemd/system/comfyui-worker@.service
sudo systemctl daemon-reload

for index in $(seq 0 "$((WORKERS - 1))"); do
  sudo systemctl enable --now "comfyui-worker@${index}.service"
  sudo systemctl restart "comfyui-worker@${index}.service"
done

sleep 5
for index in $(seq 0 "$((WORKERS - 1))"); do
  port=$((8188 + index))
  echo "worker ${index} http://127.0.0.1:${port}"
  curl -fsS "http://127.0.0.1:${port}/models/checkpoints" | head -c 500
  echo
done