#!/usr/bin/env bash

set -euo pipefail

COMFYUI_HOST="${COMFYUI_HOST:-127.0.0.1}"
COMFYUI_PORT="${COMFYUI_PORT:-8188}"
COMFYUI_URL="${COMFYUI_URL:-http://${COMFYUI_HOST}:${COMFYUI_PORT}}"
COMFYUI_BASE_DIR="${COMFYUI_BASE_DIR:-/Users/syntaxbullet/Documents/ComfyUI}"
COMFYUI_SOURCE_DIR="${COMFYUI_SOURCE_DIR:-/Users/syntaxbullet/ComfyUI-Installs/ComfyUI/ComfyUI}"
COMFYUI_PYTHON="${COMFYUI_PYTHON:-${COMFYUI_BASE_DIR}/.venv/bin/python}"
COMFYUI_MODEL_PATHS_CONFIG="${COMFYUI_MODEL_PATHS_CONFIG:-/Users/syntaxbullet/Library/Application Support/Comfy Desktop/shared_model_paths.yaml}"
COMFYUI_READY_TIMEOUT_SECONDS="${COMFYUI_READY_TIMEOUT_SECONDS:-120}"

comfy_pid=""

cleanup() {
  if [[ -n "${comfy_pid}" ]] && kill -0 "${comfy_pid}" 2>/dev/null; then
    echo "Stopping ComfyUI (PID ${comfy_pid})..."
    kill "${comfy_pid}" 2>/dev/null || true
    wait "${comfy_pid}" 2>/dev/null || true
  fi
}

trap cleanup EXIT INT TERM

comfy_is_ready() {
  curl --fail --silent --show-error --max-time 2 "${COMFYUI_URL}/object_info" >/dev/null 2>&1
}

if comfy_is_ready; then
  echo "Using existing ComfyUI server at ${COMFYUI_URL}."
else
  if [[ ! -x "${COMFYUI_PYTHON}" ]]; then
    echo "ComfyUI Python executable not found: ${COMFYUI_PYTHON}" >&2
    exit 1
  fi

  if [[ ! -f "${COMFYUI_SOURCE_DIR}/main.py" ]]; then
    echo "ComfyUI entrypoint not found: ${COMFYUI_SOURCE_DIR}/main.py" >&2
    exit 1
  fi

  echo "Starting ComfyUI at ${COMFYUI_URL}..."
  "${COMFYUI_PYTHON}" "${COMFYUI_SOURCE_DIR}/main.py" \
    --base-directory "${COMFYUI_BASE_DIR}" \
    --user-directory "${COMFYUI_BASE_DIR}/user" \
    --input-directory "${COMFYUI_BASE_DIR}/input" \
    --output-directory "${COMFYUI_BASE_DIR}/output" \
    --extra-model-paths-config "${COMFYUI_MODEL_PATHS_CONFIG}" \
    --listen "${COMFYUI_HOST}" \
    --port "${COMFYUI_PORT}" &
  comfy_pid=$!

  deadline=$((SECONDS + COMFYUI_READY_TIMEOUT_SECONDS))
  until comfy_is_ready; do
    if ! kill -0 "${comfy_pid}" 2>/dev/null; then
      wait "${comfy_pid}" || true
      echo "ComfyUI exited before becoming ready." >&2
      exit 1
    fi

    if (( SECONDS >= deadline )); then
      echo "Timed out waiting for ComfyUI after ${COMFYUI_READY_TIMEOUT_SECONDS} seconds." >&2
      exit 1
    fi

    sleep 1
  done

  echo "ComfyUI is ready."
fi

echo "Starting Image Studio..."
COMFYUI_URL="${COMFYUI_URL}" bun --hot index.ts