#!/usr/bin/env bash
# RunPod ComfyUI Krea 2 V2 by Aitrepreneur

set -euo pipefail

# ───────────────────── Config ─────────────────────

HF_BASE="${HF_BASE:-https://huggingface.co/Aitrepreneur/FLX/resolve/main}"

PYTHON_BIN="${PYTHON_BIN:-python3}"
VENV_DIR="${VENV_DIR:-venv}"

# RunPod CUDA 12.8 safe stack
TORCH_VERSION="${TORCH_VERSION:-2.8.0}"
TORCHVISION_VERSION="${TORCHVISION_VERSION:-0.23.0}"
TORCHAUDIO_VERSION="${TORCHAUDIO_VERSION:-2.8.0}"
CUDA_TAG="${CUDA_TAG:-cu128}"
TORCH_INDEX="${TORCH_INDEX:-https://download.pytorch.org/whl/${CUDA_TAG}}"

# Node install behavior
INSTALL_ALL_NODES="${INSTALL_ALL_NODES:-false}"

# This list must match the nodes actually cloned below.
REQUIRED_NODES="${REQUIRED_NODES:-"ComfyUI-Manager rgthree-comfy ComfyUI-KJNodes ComfyUI_essentials ComfyUI-RBG-SmartSeedVariance ComfyUI-ConditioningKrea2Rebalance"}"

# Fragile dependency behavior
MANAGER_ENABLE_MATRIX="${MANAGER_ENABLE_MATRIX:-false}"

# Pins
PIN_PILLOW_MIN="${PIN_PILLOW_MIN:-11.0.0}"
PIN_NUMPY="${PIN_NUMPY:-1.26.4}"

# Keep Transformers 4.x for compatibility with ComfyUI and custom nodes.
PIN_TRANSFORMERS="${PIN_TRANSFORMERS:-4.50.3}"
PIN_TOKENIZERS_RANGE="${PIN_TOKENIZERS_RANGE:->=0.21,<0.22}"
PIN_HF_HUB_RANGE="${PIN_HF_HUB_RANGE:->=0.26,<1.0}"

PIN_URLLIB3_1X="${PIN_URLLIB3_1X:-1.26.18}"
PIN_LIBROSA_VERSION="${PIN_LIBROSA_VERSION:-}"

export PIP_DISABLE_PIP_VERSION_CHECK=1
export PIP_ROOT_USER_ACTION=ignore
export PYTHONNOUSERSITE=1
unset PYTHONPATH || true

# ───────────────────── Helpers ─────────────────────

[[ "$(id -u)" -eq 0 ]] && SUDO="" || SUDO="sudo"

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

need_pkg() {
  if ! command -v "$1" >/dev/null 2>&1; then
    echo "[INFO] installing $1 ..."
    $SUDO apt-get update -y
    $SUDO apt-get install -y "$1"
  fi
}

grab() {
  local target="$1"
  local url="$2"

  if [[ -f "$target" ]]; then
    echo " • $(basename "$target") exists - skip"
    return 0
  fi

  echo " • downloading $(basename "$target")"
  mkdir -p "$(dirname "$target")"
  curl -L --fail --progress-bar --show-error -o "$target" "$url"
}

get_node() {
  local dir="$1"
  local url="$2"
  local 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
}

write_constraints() {
  CONSTRAINT_FILE="/tmp/aitrepreneur_constraints.txt"

  cat > "$CONSTRAINT_FILE" <<EOF
torch==${TORCH_VERSION}+${CUDA_TAG}
torchvision==${TORCHVISION_VERSION}+${CUDA_TAG}
torchaudio==${TORCHAUDIO_VERSION}+${CUDA_TAG}
numpy==${PIN_NUMPY}
transformers==${PIN_TRANSFORMERS}
tokenizers${PIN_TOKENIZERS_RANGE}
huggingface-hub${PIN_HF_HUB_RANGE}
Pillow>=${PIN_PILLOW_MIN}
EOF

  echo "──────── Using pip constraints ────────"
  cat "$CONSTRAINT_FILE"
}

install_torch_stack() {
  echo
  echo "──────── Installing locked Torch stack ────────"

  $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}"
}

install_core_pins() {
  echo
  echo "──────── Installing core pinned packages ────────"

  $PYTHON -m pip install --no-input --upgrade-strategy only-if-needed \
    "numpy==${PIN_NUMPY}" \
    "pillow>=${PIN_PILLOW_MIN}" \
    "transformers==${PIN_TRANSFORMERS}" \
    "tokenizers${PIN_TOKENIZERS_RANGE}" \
    "huggingface-hub${PIN_HF_HUB_RANGE}" \
    "requests>=2.32.3,<3" \
    "charset-normalizer>=2,<4" \
    "chardet<6" \
    "fsspec<=2025.3.0,>=2023.1.0"
}

install_comfyui_requirements() {
  echo
  echo "──────── Installing ComfyUI requirements ────────"

  if [[ ! -f "$COMFY_ROOT/requirements.txt" ]]; then
    die "Could not find $COMFY_ROOT/requirements.txt"
  fi

  $PYTHON -m pip install --no-input --prefer-binary \
    --upgrade-strategy only-if-needed \
    -c "$CONSTRAINT_FILE" \
    -r "$COMFY_ROOT/requirements.txt"
}

sanitize_requirements_file() {
  local input="$1"
  local output="$2"

  # Never allow custom-node requirements to replace the locked GPU stack.
  grep -Eiv '^[[:space:]]*(torch|torchvision|torchaudio|xformers|transformers|tokenizers|huggingface-hub|numpy|nvidia-|cuda-toolkit|cuda-bindings|triton)([<>=!~ ].*)?$' "$input" > "$output" || true
}

install_node_requirements_safe() {
  local req="$1"
  local req_dir
  local tmp_req

  req_dir="$(dirname "$req")"
  tmp_req="$(mktemp)"

  sanitize_requirements_file "$req" "$tmp_req"

  echo "   • $req"

  if [[ ! -s "$tmp_req" ]]; then
    echo "     ↳ [SKIP] no safe packages left after filtering"
    rm -f "$tmp_req"
    return 0
  fi

  pushd "$req_dir" >/dev/null

  if ! $PYTHON -m pip install --no-input --prefer-binary --no-build-isolation \
    --upgrade-strategy only-if-needed \
    -c "$CONSTRAINT_FILE" \
    -r "$tmp_req"; then

    echo "     ↳ [WARN] First attempt failed; retrying with build isolation..."

    $PYTHON -m pip install --no-input --prefer-binary \
      --upgrade-strategy only-if-needed \
      -c "$CONSTRAINT_FILE" \
      -r "$tmp_req" || {
        echo "     ↳ [ERROR] Failed: $req - continuing"
      }
  fi

  popd >/dev/null
  rm -f "$tmp_req"
}

protect_comfyui_manager_pins() {
  echo
  echo "──────── Protecting pinned packages from ComfyUI-Manager ────────"

  local manager_dir="$COMFY_ROOT/user/__manager"
  mkdir -p "$manager_dir"

  cat > "$manager_dir/pip_blacklist.list" <<EOF
torch
torchvision
torchaudio
xformers
transformers
tokenizers
huggingface-hub
numpy
nvidia-
cuda-toolkit
cuda-bindings
triton
torchcodec
EOF

  cat > "$manager_dir/pip_auto_fix.list" <<EOF
torch==${TORCH_VERSION}+${CUDA_TAG}
torchvision==${TORCHVISION_VERSION}+${CUDA_TAG}
torchaudio==${TORCHAUDIO_VERSION}+${CUDA_TAG}
transformers==${PIN_TRANSFORMERS}
tokenizers${PIN_TOKENIZERS_RANGE}
huggingface-hub${PIN_HF_HUB_RANGE}
numpy==${PIN_NUMPY}
EOF

  if [[ -f "$manager_dir/config.ini" ]]; then
    sed -i 's/^always_lazy_install *= *.*/always_lazy_install = False/' "$manager_dir/config.ini" || true
    sed -i 's/^use_uv *= *.*/use_uv = False/' "$manager_dir/config.ini" || true
    sed -i 's/^network_mode *= *.*/network_mode = public/' "$manager_dir/config.ini" || true
  else
    cat > "$manager_dir/config.ini" <<EOF
[default]
always_lazy_install = False
use_uv = False
network_mode = public
security_level = normal
EOF
  fi
}

final_check() {
  echo
  echo "──────── Final dependency check ────────"

  $PYTHON - <<PY
import torch
print("Torch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())

if not torch.__version__.startswith("${TORCH_VERSION}"):
    raise RuntimeError(
        f"Wrong Torch version: {torch.__version__}. Expected ${TORCH_VERSION}."
    )

if not torch.cuda.is_available():
    raise RuntimeError("CUDA is not available. Torch/CUDA install is broken.")

import torchvision
print("TorchVision:", torchvision.__version__)

import torchaudio
print("TorchAudio:", torchaudio.__version__)

import numpy
print("NumPy:", numpy.__version__)

import transformers
import tokenizers
import huggingface_hub
print("Transformers:", transformers.__version__)
print("Tokenizers:", tokenizers.__version__)
print("HuggingFace Hub:", huggingface_hub.__version__)

from transformers import AutoProcessor, AutoTokenizer, AutoModel, WhisperProcessor
print("Transformers import check: OK")
PY

  if [[ ! -f "$COMFY_ROOT/comfy/text_encoders/krea2.py" ]]; then
    die "This ComfyUI installation is too old and does not contain native Krea 2 support."
  fi

  echo "Native Krea 2 support check: OK"
}

# ───────────────────── System packages ─────────────────────

echo
echo "──────── Checking prerequisites ────────"
need_pkg curl
need_pkg git
need_pkg git-lfs
need_pkg ffmpeg
need_pkg python3-venv
need_pkg python3-pip
need_pkg build-essential

git lfs install || true

# ───────────────────── Verify paths ─────────────────────

[[ -d "models" && -d "custom_nodes" ]] || die "Run this inside your ComfyUI root folder. Need models/ and custom_nodes/."

COMFY_ROOT="$(pwd)"
echo "ComfyUI root: $COMFY_ROOT"

# ───────────────────── Update ComfyUI ─────────────────────

if [[ -d "$COMFY_ROOT/.git" ]]; then
  echo
  echo "──────── Updating ComfyUI ────────"

  if ! git -C "$COMFY_ROOT" pull --ff-only; then
    echo "[WARN] ComfyUI update failed. Continuing with the current installation."
  fi
else
  echo
  echo "[WARN] This ComfyUI folder is not a Git repository, so it cannot be updated automatically."
fi

# ───────────────────── Venv ─────────────────────

if [[ ! -d "$VENV_DIR" ]]; then
  echo "──────── Creating venv at $VENV_DIR ────────"
  $PYTHON_BIN -m venv "$VENV_DIR"
fi

if [[ -f "$VENV_DIR/pyvenv.cfg" ]]; then
  sed -i 's/^include-system-site-packages = .*/include-system-site-packages = false/' "$VENV_DIR/pyvenv.cfg" || true
fi

# shellcheck disable=SC1091
source "$VENV_DIR/bin/activate"

PYTHON="$(command -v python)"
PIP="$(command -v pip)"

echo "Using venv python: $PYTHON"

$PYTHON -m pip install --no-input --upgrade pip setuptools wheel

write_constraints

# ───────────────────── Baseline ─────────────────────

install_torch_stack
install_core_pins
install_comfyui_requirements

# ───────────────────── Models ─────────────────────

echo
echo "──────── Downloading Krea 2 Model Files ────────"

grab "models/text_encoders/qwen3vl_4b_fp8_scaled.safetensors" "$HF_BASE/qwen3vl_4b_fp8_scaled.safetensors?download=true"
grab "models/vae/qwen_image_vae.safetensors" "$HF_BASE/qwen_image_vae.safetensors?download=true"
grab "models/diffusion_models/krea2_turbo_fp8.safetensors" "$HF_BASE/krea2_turbo_fp8.safetensors?download=true"
grab "models/loras/krea2_turbo_lora_rank_64_bf16.safetensors" "$HF_BASE/krea2_turbo_lora_rank_64_bf16.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.git"
get_node "ComfyUI-KJNodes"                    "https://github.com/kijai/ComfyUI-KJNodes.git"
get_node "ComfyUI_essentials"                 "https://github.com/cubiq/ComfyUI_essentials.git"
get_node "ComfyUI-RBG-SmartSeedVariance"      "https://github.com/RamonGuthrie/ComfyUI-RBG-SmartSeedVariance.git"
get_node "ComfyUI-Krea2T-Enhancer" "https://github.com/capitan01R/ComfyUI-Krea2T-Enhancer.git"

# ───────────────────── Patch fragile requirements ─────────────────────

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/' \
      custom_nodes/ComfyUI-Manager/requirements.txt || true
  fi
fi

if [[ "$MANAGER_ENABLE_MATRIX" == "true" ]]; then
  $PYTHON -m pip install --no-input "urllib3==${PIN_URLLIB3_1X}"
fi

# ───────────────────── 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 safely ────────"

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

for req in "${REQ_FILES[@]}"; do
  install_node_requirements_safe "$req"
done

# ───────────────────── Base extras ─────────────────────

echo
echo "──────── Installing base extras ────────"
$PYTHON -m pip install --no-input -c "$CONSTRAINT_FILE" piexif lark || true

if [[ -n "$PIN_LIBROSA_VERSION" ]]; then
  echo "──────── Installing librosa==${PIN_LIBROSA_VERSION} ────────"
  $PYTHON -m pip install --no-input -c "$CONSTRAINT_FILE" "librosa==${PIN_LIBROSA_VERSION}" || true
else
  echo "──────── Installing latest librosa ────────"
  $PYTHON -m pip install --no-input -c "$CONSTRAINT_FILE" librosa || true
fi

# ───────────────────── Protect Manager and verify ─────────────────────

install_core_pins
protect_comfyui_manager_pins
final_check

echo
echo "✅ All Krea 2 models and custom nodes are ready"
echo