#!/bin/bash
set -euo pipefail

# ================== INSTALLATION CONFIG ==================
INSTALL_BASE="/workspace/swarmui"
echo "Installing in /workspace"

# ================== CONFIG ==================
COMFYUI_ROOT="${INSTALL_BASE}/comfyui"
COMFYUI_DIR="${COMFYUI_ROOT}/ComfyUI"
COMFY_PORT=8188
LOGFILE="/cloudflared_8188.log"
HELPERS_LOGFILE="/helpers.log"
SETUP_LOG="/tmp/comfysetup_setup.log"
# ============================================
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

if [ -t 1 ]; then
  RESET=$'\033[0m'
  BOLD=$'\033[1m'
  RED=$'\033[31m'
  GREEN=$'\033[32m'
  YELLOW=$'\033[33m'
  BLUE=$'\033[34m'
  CYAN=$'\033[36m'
else
  RESET=""
  BOLD=""
  RED=""
  GREEN=""
  YELLOW=""
  BLUE=""
  CYAN=""
fi

: > "${SETUP_LOG}"

log() {
  printf "%b[%s]%b %s\n" "${BLUE}" "INFO" "${RESET}" "$1"
}

err() {
  printf "%b[%s]%b %s\n" "${RED}" "ERROR" "${RESET}" "$1" >&2
    exit 1
}

ok() {
  printf "%b[%s]%b %s\n" "${GREEN}" "OK" "${RESET}" "$1"
}

section() {
  printf "\n%b%s%b\n" "${BOLD}${CYAN}" "$1" "${RESET}"
}

kv() {
  printf "  %b%-18s%b %s\n" "${YELLOW}" "$1" "${RESET}" "$2"
}

run_logged() {
  local label="$1"
  local logfile="$2"
  shift 2

  printf "  %b[%s]%b %s\n" "${BLUE}" ".." "${RESET}" "$label"
  if "$@" >>"${logfile}" 2>&1; then
    printf "  %b[%s]%b %s\n" "${GREEN}" "OK" "${RESET}" "$label"
  else
    printf "  %b[%s]%b %s\n" "${RED}" "FAIL" "${RESET}" "$label" >&2
    printf "%s\n" "--- last 40 lines from ${logfile} ---" >&2
    tail -n 40 "${logfile}" >&2 || true
    exit 1
  fi
}

get_cloudflare_url() {
  local timeout="${1:-60}"
  local waited=0

  while [ "$waited" -lt "$timeout" ]; do
    if [ -f "${LOGFILE}" ]; then
      local url
      url=$(grep -oE 'https://[a-zA-Z0-9.-]+\.trycloudflare\.com' "${LOGFILE}" | head -n 1 || true)
      if [ -n "$url" ]; then
        echo "$url"
        return 0
      fi
    fi

    sleep 1
    waited=$((waited + 1))
  done

  return 1
}

get_gradio_url() {
  local timeout="${1:-60}"
  local waited=0

  while [ "$waited" -lt "$timeout" ]; do
    if [ -f "${HELPERS_LOGFILE}" ]; then
      local url
      url=$(grep -oE 'https://[a-zA-Z0-9.-]+\.gradio\.live' "${HELPERS_LOGFILE}" | head -n 1 || true)
      if [ -n "$url" ]; then
        echo "$url"
        return 0
      fi
    fi

    sleep 1
    waited=$((waited + 1))
  done

  return 1
}

# ------------------ ROOT CHECK ------------------
[ "$EUID" -eq 0 ] || err "Run this script as root or with sudo"

# ------------------ CREATE DIRECTORY ------------------
log "Creating installation directory"

# Create the installation directory
mkdir -p "${INSTALL_BASE}"

ok "Installation directory ready: ${INSTALL_BASE}"

section "ComfySetup dashboard"
kv "Install base" "${INSTALL_BASE}"
kv "ComfyUI root" "${COMFYUI_ROOT}"
kv "Port" "${COMFY_PORT}"
kv "Log file" "${SETUP_LOG}"

# ------------------ SYSTEM DEPS ------------------
log "Installing system dependencies"

run_logged "apt update" "${SETUP_LOG}" apt update -y
run_logged "apt install core tools" "${SETUP_LOG}" apt install -y \
  curl git screen python3 python3-venv lsb-release \
  zip aria2 nano

# ---- uv ----
if ! command -v uv >/dev/null; then
    run_logged "install uv" "${SETUP_LOG}" bash -lc 'curl -Ls https://astral.sh/uv/install.sh | sh'
fi
export PATH="/root/.cargo/bin:$PATH"

# ---- cloudflared ----
run_logged "download cloudflare key" "${SETUP_LOG}" curl -fsSL https://pkg.cloudflare.com/cloudflare-main.gpg -o /usr/share/keyrings/cloudflare-archive-keyring.gpg

printf 'deb [signed-by=/usr/share/keyrings/cloudflare-archive-keyring.gpg] https://pkg.cloudflare.com/cloudflared/ %s main\n' "$(lsb_release -cs)" \
  > /etc/apt/sources.list.d/cloudflared.list

run_logged "apt update cloudflared" "${SETUP_LOG}" apt update -y
run_logged "install cloudflared" "${SETUP_LOG}" apt install -y cloudflared

ok "System dependencies installed"

# ------------------ COMFYUI ------------------
log "Starting ComfyUI screen"

screen -dmS comfyui bash -lc "
set -euo pipefail

mkdir -p ${COMFYUI_ROOT}
cd ${COMFYUI_ROOT}

if [ ! -d ComfyUI ]; then
    git clone --depth 1 https://github.com/comfyanonymous/ComfyUI
fi

cd ComfyUI
git reset --hard || true
git clean -fd || true
git pull --force || true

# ---- venv (FIX: standard venv, NOT uv) ----
python3 -m venv venv
source venv/bin/activate

# ---- core deps (unchanged) ----
pip install --upgrade pip
uv pip install \
  torch torchvision torchaudio \
  --index-url https://download.pytorch.org/whl/cu124

# ---- custom nodes ----
mkdir -p custom_nodes
cd custom_nodes

repos=(
  https://github.com/ltdrdata/ComfyUI-Manager
  https://github.com/city96/ComfyUI-GGUF
  https://github.com/ClownsharkBatwing/RES4LYF
  https://github.com/rgthree/rgthree-comfy
  https://github.com/crystian/ComfyUI-Crystools
  https://github.com/MoonGoblinDev/Civicomfy
  https://github.com/1038lab/ComfyUI-QwenVL
  https://github.com/1038lab/ComfyUI-RMBG
  https://github.com/Fannovel16/ComfyUI-Frame-Interpolation
  https://github.com/kijai/ComfyUI-KJNodes
  https://github.com/liusida/ComfyUI-Login
  https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler
  https://github.com/LAOGOU-666/Comfyui-Memory_Cleanup
)

rm -f /tmp/merged_reqs.txt

for repo in \"\${repos[@]}\"; do
    name=\$(basename \"\$repo\" .git)
    [ -d \"\$name\" ] || git clone --depth 1 \"\$repo\"
    cd \"\$name\" || continue
    git pull --force || true
    [ -f requirements.txt ] && echo \"\" >> /tmp/merged_reqs.txt && cat requirements.txt >> /tmp/merged_reqs.txt
    cd ..
done

cd ..
[ -f requirements.txt ] && echo \"\" >> /tmp/merged_reqs.txt && cat requirements.txt >> /tmp/merged_reqs.txt
[ -f /tmp/merged_reqs.txt ] && uv pip install -r /tmp/merged_reqs.txt --extra-index-url https://download.pytorch.org/whl/cu124 || true
rm -f /tmp/merged_reqs.txt

# -------- Optional / Extra Perf Extensions --------

uv pip install \
  \"https://huggingface.co/gtaayush010/comfysetup/resolve/main/Wheels/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl?download=true\" \
  || true

uv pip install \
  \"https://huggingface.co/gtaayush010/comfysetup/resolve/main/Wheels/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl?download=true\" \
  || true

uv pip install \
  insightface \
  || true

# ---- run (unchanged) ----
exec python main.py \
  --listen 0.0.0.0 \
  --port ${COMFY_PORT} \
  --use-sage-attention \
"

log "ComfyUI launched (screen: comfyui)"

# ------------------ CLOUDFLARED ------------------
log "Starting Cloudflared (no waiting)"

rm -f ${LOGFILE}

screen -dmS tunnel8188 bash -lc "
cloudflared tunnel \
  --no-autoupdate \
  --url http://127.0.0.1:${COMFY_PORT} \
  > ${LOGFILE} 2>&1
exec bash
"

log "Cloudflared launched (screen: tunnel8188)"

log "Waiting for Cloudflared public URL"
CLOUDFLARE_URL="$(get_cloudflare_url 60 || true)"

# ------------------ HELPERS ------------------
log "Starting Helpers screen"

rm -f ${HELPERS_LOGFILE}

screen -dmS helpers bash -lc "
set -euo pipefail

cd ${INSTALL_BASE}
mkdir -p helpers
cd helpers

curl -fsSL -o helpers.py https://huggingface.co/gtaayush010/comfysetup/resolve/main/helpers.py

# ---- venv ----
export PATH=\"/root/.cargo/bin:\$PATH\"
uv venv venv
source venv/bin/activate
uv pip install gradio huggingface_hub

export PYTHONUNBUFFERED=1
python helpers.py 2>&1 | tee ${HELPERS_LOGFILE} || true
exec bash
"

log "Helpers launched (screen: helpers)"

# ------------------ DONE ------------------
section "Setup complete"
kv "Installation" "${INSTALL_BASE}"
kv "ComfyUI" "${INSTALL_BASE}/comfyui/ComfyUI"
if [ -n "${CLOUDFLARE_URL}" ]; then
  kv "Cloudflare URL" "${CLOUDFLARE_URL}"
else
  kv "Cloudflare URL" "not found yet"
  kv "Tunnel log" "${LOGFILE}"
fi
section "Attach"
kv "ComfyUI" "screen -r comfyui"
kv "Tunnel" "screen -r tunnel8188"
kv "Helpers" "screen -r helpers"

section "RunPod restart"
ok "After restart, run: bash ${INSTALL_BASE}/restart_after_reboot.sh"

run_logged "download restart script" "${SETUP_LOG}" curl -fsSL \
  -o "${INSTALL_BASE}/restart_after_reboot.sh" \
  "https://huggingface.co/gtaayush010/comfysetup/resolve/main/restart_after_reboot.sh"
chmod +x "${INSTALL_BASE}/restart_after_reboot.sh"

section "Restart script saved"
kv "Command" "bash ${INSTALL_BASE}/restart_after_reboot.sh"
kv "Path" "${INSTALL_BASE}/restart_after_reboot.sh"

section "Waiting for Gradio"
kv "Status" "waiting for public URL"
kv "Note" "This may take a few minutes while dependencies install"

GRADIO_URL="$(get_gradio_url 3600 || true)"

if [ -n "${GRADIO_URL}" ]; then
  section "Gradio URL"
  kv "URL" "${GRADIO_URL}"
else
  section "Gradio URL"
  kv "Status" "timed out waiting"
  kv "Helpers log" "${HELPERS_LOGFILE}"
fi