#!/bin/bash
set -euo pipefail

# ================== GCP INSTALLATION CONFIG ==================
INSTALL_BASE="${HOME}/workspace"
COMFYUI_ROOT="${INSTALL_BASE}/comfyui"
COMFYUI_DIR="${INSTALL_BASE}/comfyui/ComfyUI"
COMFY_PORT=8188

LOGFILE="/tmp/cloudflared_8188.log"
HELPERS_LOGFILE="/tmp/helpers.log"
SETUP_LOG="/tmp/gcp_comfysetup_setup.log"
COMFYUI_LOG="/tmp/comfyui.log"
# ============================================================
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Color formatting
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}"
: > "${COMFYUI_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%-22s%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
}

# ------------------ 1. CREATE WORKSPACE DIRECTORY ------------------
section "1. Preparing GCP Workspace Directory"
mkdir -p "${INSTALL_BASE}"
ok "Workspace directory ready: ${INSTALL_BASE}"

section "gcp_comfysetup Dashboard"
kv "Install base"    "${INSTALL_BASE}"
kv "ComfyUI root"    "${COMFYUI_ROOT}"
kv "ComfyUI Dir"     "${COMFYUI_DIR}"
kv "Port"            "${COMFY_PORT}"
kv "Log file"        "${SETUP_LOG}"
kv "ComfyUI Log"     "${COMFYUI_LOG}"

# ------------------ 2. SYSTEM DEPENDENCIES & TOOLING ------------------
section "2. Installing System Dependencies & Tools"

SUDO_CMD=""
if command -v sudo >/dev/null 2>&1; then
    SUDO_CMD="sudo"
fi

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

# ---- Install uv (fast Python package manager) ----
if ! command -v uv >/dev/null 2>&1; then
    run_logged "install uv" "${SETUP_LOG}" bash -lc 'curl -Ls https://astral.sh/uv/install.sh | sh'
fi

export PATH="${HOME}/.local/bin:${HOME}/.cargo/bin:/root/.cargo/bin:$PATH"
if [ -f "${HOME}/.local/bin/env" ]; then
  source "${HOME}/.local/bin/env" || true
fi

# ---- Install cloudflared ----
if ! command -v cloudflared >/dev/null 2>&1; then
    run_logged "download cloudflare key" "${SETUP_LOG}" ${SUDO_CMD} curl -fsSL https://pkg.cloudflare.com/cloudflare-main.gpg -o /usr/share/keyrings/cloudflare-archive-keyring.gpg || curl -fsSL https://pkg.cloudflare.com/cloudflare-main.gpg -o /usr/share/keyrings/cloudflare-archive-keyring.gpg
    
    ${SUDO_CMD} bash -c "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" || true
    
    run_logged "apt update cloudflared" "${SETUP_LOG}" ${SUDO_CMD} apt-get update -y || apt-get update -y
    run_logged "apt install cloudflared" "${SETUP_LOG}" ${SUDO_CMD} apt-get install -y cloudflared || apt-get install -y cloudflared
fi

# ---- Install rclone ----
if ! command -v rclone >/dev/null 2>&1; then
    run_logged "install rclone" "${SETUP_LOG}" bash -c "curl -fsSL https://rclone.org/install.sh | ${SUDO_CMD} bash"
fi

ok "System dependencies & tools installed"

# ------------------ 3. COMFYUI ENVIRONMENT SETUP & LAUNCH ------------------
section "3. Setting up ComfyUI Screen Session & Dependencies"

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

  # ---- Python Virtual Environment ----
  if [ ! -d venv ]; then
      python3 -m venv venv
  fi
  source venv/bin/activate

  # Upgrade pip
  pip install --upgrade pip

  # Determine installer (uv if available, else pip)
  INSTALLER=\"pip install\"
  if command -v uv >/dev/null 2>&1; then
      INSTALLER=\"uv pip install\"
  elif [ -f \"\${HOME}/.local/bin/uv\" ]; then
      INSTALLER=\"\${HOME}/.local/bin/uv pip install\"
  fi

  # ---- Install PyTorch 2.11.0 (CUDA 12.8) ----
  echo '=== Installing PyTorch 2.11.0 (CUDA 12.8) ==='
  \${INSTALLER} torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu128

  # ---- Install xFormers ----
  echo '=== Installing xFormers ==='
  \${INSTALLER} xformers --no-deps || \${INSTALLER} xformers || true

  # ---- Install 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 ] && \${INSTALLER} -r /tmp/merged_reqs.txt --extra-index-url https://download.pytorch.org/whl/cu128 || true
  rm -f /tmp/merged_reqs.txt

  # ---- Extra Dependencies ----
  \${INSTALLER} insightface || true

  # ---- Launch ComfyUI ----
  echo '=== Launching ComfyUI Server ==='
  exec python main.py \
    --listen 0.0.0.0 \
    --port ${COMFY_PORT}
) 2>&1 | tee ${COMFYUI_LOG}
exec bash
"

log "ComfyUI screen session launched (screen -r comfyui)"

# ------------------ 4. CLOUDFLARED TUNNEL ------------------
section "4. Starting Cloudflared Tunnel"

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 -r tunnel8188)"
log "Waiting for Cloudflared public URL..."
CLOUDFLARE_URL="$(get_cloudflare_url 60 || true)"

# ------------------ 5. HELPERS GRADIO UI ------------------
section "5. Starting Helper Tools UI"

rm -f "${HELPERS_LOGFILE}"
HELPERS_DIR="${INSTALL_BASE}/helpers"
mkdir -p "${HELPERS_DIR}"

if [ -f "${SCRIPT_DIR}/helpers.py" ]; then
    cp "${SCRIPT_DIR}/helpers.py" "${HELPERS_DIR}/helpers.py"
else
    curl -fsSL -o "${HELPERS_DIR}/helpers.py" https://huggingface.co/gtaayush010/gcp_comfysetup/resolve/main/helpers.py
fi

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

cd ${HELPERS_DIR}
export PATH=\"\${HOME}/.local/bin:\${HOME}/.cargo/bin:\$PATH\"

if [ ! -d venv ]; then
    python3 -m venv venv
fi
source venv/bin/activate

if command -v uv >/dev/null 2>&1; then
    uv pip install gradio huggingface_hub
else
    pip install --upgrade pip
    pip install gradio huggingface_hub
fi

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

log "Helpers UI launched (screen -r helpers)"

# ------------------ 6. REBOOT RESTART SCRIPT ------------------
section "6. Generating Reboot Recovery Script"

RESTART_SCRIPT="${INSTALL_BASE}/restart_after_reboot.sh"

cat <<'EOF' > "${RESTART_SCRIPT}"
#!/bin/bash
set -euo pipefail

INSTALL_BASE="${HOME}/workspace"
COMFYUI_ROOT="${INSTALL_BASE}/comfyui"
COMFYUI_DIR="${COMFYUI_ROOT}/ComfyUI"
COMFY_PORT=8188
LOGFILE="/tmp/cloudflared_8188.log"
HELPERS_LOGFILE="/tmp/helpers.log"

export PATH="${HOME}/.local/bin:${HOME}/.cargo/bin:$PATH"

echo "=== Restarting ComfyUI Services on GCP ==="

# 1. ComfyUI Screen
if ! screen -list | grep -q "comfyui"; then
  echo "Launching ComfyUI screen..."
  screen -dmS comfyui bash -lc "
    cd ${COMFYUI_DIR}
    source venv/bin/activate
    exec python main.py --listen 0.0.0.0 --port ${COMFY_PORT}
  "
else
  echo "ComfyUI screen session is already running."
fi

# 2. Cloudflared Screen
if ! screen -list | grep -q "tunnel8188"; then
  echo "Launching Cloudflared tunnel screen..."
  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
  "
else
  echo "Cloudflared screen session is already running."
fi

# 3. Helpers Screen
if ! screen -list | grep -q "helpers"; then
  echo "Launching Helpers screen..."
  rm -f ${HELPERS_LOGFILE}
  screen -dmS helpers bash -lc "
    cd ${INSTALL_BASE}/helpers
    source venv/bin/activate
    export PYTHONUNBUFFERED=1
    python helpers.py 2>&1 | tee ${HELPERS_LOGFILE} || true
    exec bash
  "
else
  echo "Helpers screen session is already running."
fi

echo "=== All services restarted successfully ==="
EOF

chmod +x "${RESTART_SCRIPT}"
ok "Reboot recovery script created: ${RESTART_SCRIPT}"

# ------------------ SUMMARY ------------------
section "Setup Complete!"
kv "Installation Path" "${INSTALL_BASE}"
kv "ComfyUI Directory" "${COMFYUI_DIR}"
kv "PyTorch Version"    "2.11.0 + xFormers (CUDA 12.8)"
if [ -n "${CLOUDFLARE_URL}" ]; then
  kv "Cloudflare URL"   "${CLOUDFLARE_URL}"
else
  kv "Cloudflare URL"   "Initializing... (check screen -r tunnel8188)"
fi

section "Screen Attach Commands"
kv "ComfyUI"           "screen -r comfyui"
kv "Tunnel"            "screen -r tunnel8188"
kv "Helpers"           "screen -r helpers"

section "GCP Reboot Recovery"
kv "Restart Command"   "bash ${RESTART_SCRIPT}"

section "Waiting for Gradio Helper URL..."
GRADIO_URL="$(get_gradio_url 3600 || true)"

if [ -n "${GRADIO_URL}" ]; then
  section "Gradio Helper UI"
  kv "Public URL"      "${GRADIO_URL}"
else
  section "Gradio Helper UI"
  kv "Status"          "Timed out waiting or check screen -r helpers"
fi

ok "All setup operations completed successfully!"