#!/usr/bin/env bash
# RunPod ComfyUI Ideogram 4 by Aitrepreneur

set -euo pipefail

# ───────────────────── Config (override via env) ─────────────────────
HF_BASE="${HF_BASE:-https://huggingface.co/Aitrepreneur/FLX/resolve/main}"
COMFY_IDEOGRAM_BASE="${COMFY_IDEOGRAM_BASE:-https://huggingface.co/Comfy-Org/Ideogram-4/resolve/main}"

# Py env
PYTHON_BIN="${PYTHON_BIN:-python3}"            # system python to create venv
VENV_DIR="${VENV_DIR:-venv}"

# Torch/CUDA (RunPod cu121 is common; adjust if your image differs)
TORCH_VERSION="${TORCH_VERSION:-2.4.0}"
TORCHVISION_VERSION="${TORCHVISION_VERSION:-0.19.0}"
TORCHAUDIO_VERSION="${TORCHAUDIO_VERSION:-2.4.0}"
CUDA_TAG="${CUDA_TAG:-cu121}"                  # cu118 | cu121 | cpu
TORCH_INDEX="https://download.pytorch.org/whl/${CUDA_TAG}"

# Node selection
INSTALL_ALL_NODES="${INSTALL_ALL_NODES:-false}"   # true = install every cloned node requirements
REQUIRED_NODES="${REQUIRED_NODES:-ComfyUI-Manager rgthree-comfy ComfyUI-KJNodes ComfyUI_essentials}"

# Patches for fragile nodes
MANAGER_ENABLE_MATRIX="${MANAGER_ENABLE_MATRIX:-false}"  # Matrix needs urllib3<2 and is optional for most

# Extra base libs known to reduce cross-node friction
PIN_PILLOW_MIN="${PIN_PILLOW_MIN:-11.0.0}"        # >=10.3 satisfies KJNodes; 11.x is fine
PIN_OPENCV_HEADLESS="${PIN_OPENCV_HEADLESS:-4.12.0.88}"  # unify on headless; avoid dual opencv
PIN_URLLIB3_1X="${PIN_URLLIB3_1X:-1.26.18}"       # required if MANAGER_ENABLE_MATRIX=true

export PIP_DISABLE_PIP_VERSION_CHECK=1
export PIP_ROOT_USER_ACTION=ignore

# ───────────────────── Helpers ─────────────────────
[[ $(id -u) -eq 0 ]] && SUDO="" || SUDO="sudo"
need_pkg() { command -v "$1" &>/dev/null || { echo "[INFO] installing $1 ..."; $SUDO apt-get update -y; $SUDO apt-get install -y "$1"; }; }
need_pkg curl; need_pkg git; need_pkg git-lfs; git lfs install

die() { echo "[ERROR] $*"; exit 1; }

grab () {  # grab <target path> <url>
  [[ -f "$1" ]] && { echo " • $(basename "$1") exists - skip"; return; }
  echo " • downloading $(basename "$1")"
  mkdir -p "$(dirname "$1")"
  curl -L --fail --progress-bar --show-error -o "$1" "$2"
}

get_node () {  # get_node <folder> <git url> [--recursive]
  local dir=$1 url=$2 flag=${3:-}
  if [[ -d "custom_nodes/$dir" ]]; then
    echo " [SKIP] $dir already present."
  else
    echo " • cloning $dir"
    git clone $flag "$url" "custom_nodes/$dir"
  fi
}

# ───────────────────── Verify paths ─────────────────────
[[ -d "models" && -d "custom_nodes" ]] || die "Run this in your ComfyUI root (need models/ and custom_nodes/)."
COMFY_ROOT="$(pwd)"

# ───────────────────── Venv (always) ─────────────────────
if [[ ! -d "$VENV_DIR" ]]; then
  echo "──────── Creating venv at $VENV_DIR ────────"
  $PYTHON_BIN -m venv "$VENV_DIR"
fi

# shellcheck disable=SC1091
source "$VENV_DIR/bin/activate"
PYTHON="$(command -v python)"
PIP="$(command -v pip)"
echo "Using venv python: $PYTHON"

# Keep pip tooling fresh inside venv
$PYTHON -m pip install --no-input --upgrade pip setuptools wheel

# ───────────────────── Models ─────────────────────
echo
echo "──────── Downloading Ideogram 4 Model Files ────────"

# Diffusion models
grab "models/diffusion_models/ideogram4_fp8_scaled.safetensors" "$COMFY_IDEOGRAM_BASE/diffusion_models/ideogram4_fp8_scaled.safetensors?download=true"
grab "models/diffusion_models/ideogram4_unconditional_fp8_scaled.safetensors" "$COMFY_IDEOGRAM_BASE/diffusion_models/ideogram4_unconditional_fp8_scaled.safetensors?download=true"

# Text encoders
grab "models/text_encoders/gemma4_e4b_it_fp8_scaled.safetensors" "$HF_BASE/gemma4_e4b_it_fp8_scaled.safetensors?download=true"
grab "models/text_encoders/qwen3vl_8b_fp8_scaled.safetensors" "$HF_BASE/qwen3vl_8b_fp8_scaled.safetensors?download=true"

# VAE
grab "models/vae/flux2-vae.safetensors" "$HF_BASE/flux2-vae.safetensors?download=true"

# ───────────────────── Nodes ─────────────────────
echo
echo "──────── Cloning Custom Nodes ────────"
get_node "ComfyUI-Manager"             "https://github.com/ltdrdata/ComfyUI-Manager.git"
get_node "rgthree-comfy"               "https://github.com/rgthree/rgthree-comfy"
get_node "ComfyUI-KJNodes"             "https://github.com/kijai/ComfyUI-KJNodes"
get_node "ComfyUI_essentials"          "https://github.com/cubiq/ComfyUI_essentials"

# ───────────────────── Baseline: Torch + common pins ─────────────────────
echo
echo "──────── Installing baseline (Torch, Pillow, OpenCV headless) ────────"

# Torch stack first (locked to cu121 unless overridden)
$PYTHON -m pip install --no-input --upgrade-strategy only-if-needed \
  --index-url "$TORCH_INDEX" --extra-index-url https://pypi.org/simple \
  "torch==${TORCH_VERSION}+${CUDA_TAG}" \
  "torchvision==${TORCHVISION_VERSION}+${CUDA_TAG}" \
  "torchaudio==${TORCHAUDIO_VERSION}+${CUDA_TAG}"

# Unify on headless OpenCV and pre-bump Pillow to satisfy KJNodes
$PYTHON -m pip uninstall -y opencv-python || true
$PYTHON -m pip install --no-input "opencv-python-headless==${PIN_OPENCV_HEADLESS}"
$PYTHON -m pip install --no-input "pillow>=${PIN_PILLOW_MIN}"

# If Matrix is enabled, pre-pin urllib3<2 to satisfy matrix-client 0.4.0
if [[ "$MANAGER_ENABLE_MATRIX" == "true" ]]; then
  $PYTHON -m pip install --no-input "urllib3==${PIN_URLLIB3_1X}"
fi

# ───────────────────── Patch fragile requirements (optional) ─────────────────────
# ComfyUI-Manager Matrix (urllib3<2). Disable unless explicitly requested.
if [[ "$MANAGER_ENABLE_MATRIX" != "true" ]]; then
  if [[ -f custom_nodes/ComfyUI-Manager/requirements.txt ]]; then
    sed -i 's/^matrix-client==0\.4\.0/# matrix-client disabled by installer; set MANAGER_ENABLE_MATRIX=true to enable/' \
      custom_nodes/ComfyUI-Manager/requirements.txt || true
  fi
fi

# ───────────────────── Lock constraints after sane baseline ─────────────────────
echo "   • Writing constraints to /tmp/constraints.txt"
$PYTHON -m pip freeze | sed '/^-e /d' > /tmp/constraints.txt

# ───────────────────── Install node requirements safely ─────────────────────
collect_reqs_all()      { find custom_nodes -maxdepth 2 -name requirements.txt -print; }
collect_reqs_required() {
  for dir in $REQUIRED_NODES; do
    local req="custom_nodes/$dir/requirements.txt"
    [[ -f "$req" ]] && echo "$req" || echo "   • (no requirements.txt) $dir" >&2
  done
}

echo
echo "──────── Installing node requirements (Manager-like per node) ────────"
declare -a REQ_FILES=()
if [[ "$INSTALL_ALL_NODES" == "true" ]]; then
  while IFS= read -r path; do REQ_FILES+=("$path"); done < <(collect_reqs_all)
else
  while IFS= read -r path; do [[ -f "$path" ]] && REQ_FILES+=("$path") || true; done < <(collect_reqs_required)
fi

# Re-ensure venv is active before node installs
# shellcheck disable=SC1091
source "$VENV_DIR/bin/activate"
PYTHON="$(command -v python)"
PIP="$(command -v pip)"
echo "Re-confirmed venv python for node installs: $PYTHON"

# Install each node from inside its folder (mirrors Manager behavior)
for req in "${REQ_FILES[@]}"; do
  echo "   • $req"

  # shellcheck disable=SC1091
  source "$VENV_DIR/bin/activate"
  PYTHON="$(command -v python)"
  PIP="$(command -v pip)"

  req_dir="$(dirname "$req")"
  pushd "$req_dir" >/dev/null

  # Attempt 1: prefer wheels, avoid build isolation (faster, fewer surprises)
  if ! $PYTHON -m pip install --no-input --prefer-binary --no-build-isolation \
       --upgrade-strategy only-if-needed -r requirements.txt; then
    echo "     ↳ [WARN] First attempt failed; retrying with build isolation..."

    # Attempt 2: allow build isolation as a fallback
    $PYTHON -m pip install --no-input --prefer-binary \
       --upgrade-strategy only-if-needed -r requirements.txt || {
      echo "     ↳ [ERROR] Failed: $req - continuing"
    }
  fi

  popd >/dev/null
done

# ───────────────────── Optional extras ─────────────────────
# Base extras
$PYTHON -m pip install --no-input piexif || true

echo
echo "✅ All Ideogram 4 models and nodes are ready"