#!/usr/bin/env bash
set -Eeuo pipefail

LTXVIDEO_REPO="https://github.com/Lightricks/ComfyUI-LTXVideo.git"
LTXVIDEO_REVISION="4f45fd6c222eb06eb3e46605da62e7c889e4be5c"
COMFYUI_CORE_REF="${COMFYUI_CORE_REF:-v0.27.0}"
WORKFLOW_REPO="FuzzPuppy/LTX-2.3-Foley-Workflow"
WORKFLOW_REVISION="${WORKFLOW_REVISION:-main}"
WORKFLOW_BASE_URL="https://huggingface.co/${WORKFLOW_REPO}/resolve/${WORKFLOW_REVISION}"

log() {
  printf '\n[%s] %s\n' "$(date +%H:%M:%S)" "$*"
}

die() {
  printf '\nERROR: %s\n' "$*" >&2
  exit 1
}

hf_token() {
  local token="${HF_TOKEN:-${HUGGINGFACE_TOKEN:-}}"
  printf '%s' "$token" | tr -d '[:space:]'
}

download_file() {
  local url="$1"
  local target="$2"
  local token
  local curl_args=(-fL --retry 5 --retry-delay 2)

  token="$(hf_token)"
  mkdir -p "$(dirname "$target")"
  if [[ -n "$token" ]]; then
    curl_args+=(-H "Authorization: Bearer $token")
  fi
  curl "${curl_args[@]}" "$url" -o "$target"
}

download_model_file() {
  local url="$1"
  local target_path="$2"
  local expected_sha256="$3"
  local partial_path="${target_path}.part"
  local token
  local curl_args=(-fL --retry 5 --retry-delay 2 -C -)
  local attempt

  token="$(hf_token)"
  if [[ -s "$target_path" ]]; then
    if printf '%s  %s\n' "$expected_sha256" "$target_path" | sha256sum -c - >/dev/null 2>&1; then
      printf 'verified: %s\n' "$target_path"
      return
    fi
    printf 'hash mismatch, redownloading: %s\n' "$target_path"
    rm -f "$target_path"
  fi

  mkdir -p "$(dirname "$target_path")"
  if [[ -n "$token" ]]; then
    curl_args+=(-H "Authorization: Bearer $token")
  fi

  for attempt in 1 2; do
    curl "${curl_args[@]}" "$url" -o "$partial_path"
    if printf '%s  %s\n' "$expected_sha256" "$partial_path" | sha256sum -c -; then
      mv "$partial_path" "$target_path"
      [[ -s "$target_path" ]] || die "Expected download missing: $target_path"
      return
    fi

    rm -f "$partial_path"
    curl_args=(-fL --retry 5 --retry-delay 2)
    if [[ -n "$token" ]]; then
      curl_args+=(-H "Authorization: Bearer $token")
    fi
    printf 'hash mismatch, retrying clean download: %s\n' "$target_path"
  done

  die "Downloaded file failed SHA-256 verification: $target_path"
}

restart_comfyui() {
  log "Restarting ComfyUI"
  # --cache-classic: the default RAM-pressure cache evicts node outputs
  # mid-prompt while the large AV models load, which forces the sliding-window
  # loop expansion to re-execute the entire upstream graph on every iteration.
  local comfyui_args=(--listen 0.0.0.0 --port 8188 --enable-cors-header --cache-classic)

  local pid
  while IFS= read -r pid; do
    [[ -z "$pid" ]] && continue
    kill "$pid" 2>/dev/null || true
  done < <(pgrep -f "python.*${COMFYUI_DIR}/main.py|python.*main.py.*--port 8188" || true)

  local attempt
  for attempt in {1..20}; do
    if ! pgrep -f "python.*${COMFYUI_DIR}/main.py|python.*main.py.*--port 8188" >/dev/null; then
      break
    fi
    sleep 1
  done

  if pgrep -f "python.*${COMFYUI_DIR}/main.py|python.*main.py.*--port 8188" >/dev/null; then
    die "Timed out waiting for the old ComfyUI process to stop."
  fi

  (
    cd "$COMFYUI_DIR"
    nohup "$PYTHON_BIN" main.py "${comfyui_args[@]}" > "$LOG_FILE" 2>&1 &
    printf '%s\n' "$!" > "$PID_FILE"
  )

  sleep 2
  if ! kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
    die "ComfyUI failed to start. Check $LOG_FILE"
  fi

  for attempt in {1..60}; do
    if curl -fsS http://127.0.0.1:8188/ >/dev/null 2>&1; then
      log "ComfyUI restarted with pid $(cat "$PID_FILE"). Logs: $LOG_FILE"
      return
    fi
    sleep 1
  done

  die "ComfyUI process started, but port 8188 did not become ready. Check $LOG_FILE"
}

patch_ltxvideo_kornia_import() {
  local target="$CUSTOM_NODES_DIR/ComfyUI-LTXVideo/pyramid_blending.py"
  [[ -f "$target" ]] || return

  "$PYTHON_BIN" - "$target" <<'PY'
from pathlib import Path
import sys

path = Path(sys.argv[1])
text = path.read_text(encoding="utf-8")
old = """from kornia.geometry.transform.pyramid import (
    PyrUp,
    build_laplacian_pyramid,
    build_pyramid,
    find_next_powerof_two,
    is_powerof_two,
    pad,
)
"""
new = """from kornia.geometry.transform.pyramid import (
    PyrUp,
    build_laplacian_pyramid,
    build_pyramid,
    find_next_powerof_two,
    is_powerof_two,
)
from torch.nn.functional import pad
"""
if old in text:
    path.write_text(text.replace(old, new), encoding="utf-8")
    print(f"patched kornia pad import: {path}")
elif "from torch.nn.functional import pad" in text:
    print(f"kornia pad import already patched: {path}")
else:
    print(f"warning: expected kornia import block not found in {path}", file=sys.stderr)
PY
}

ensure_pip() {
  if "$PYTHON_BIN" -m pip --version >/dev/null 2>&1; then
    return
  fi

  log "Bootstrapping pip for $PYTHON_BIN"
  "$PYTHON_BIN" -m ensurepip --upgrade || true
  if ! "$PYTHON_BIN" -m pip --version >/dev/null 2>&1; then
    die "Could not run pip with $PYTHON_BIN. The RunPod Python environment is missing pip/ensurepip."
  fi
}

pip_install() {
  ensure_pip
  "$PYTHON_BIN" -m pip install "$@"
}

install_comfyui_core_requirements() {
  local requirements="$COMFYUI_DIR/requirements.txt"

  [[ -f "$requirements" ]] || die "Could not find ComfyUI requirements file: $requirements"

  log "Installing ComfyUI core requirements for $COMFYUI_CORE_REF"
  pip_install --upgrade pip setuptools wheel
  pip_install --upgrade -r "$requirements"
}

install_comfyui_core() {
  log "Installing ComfyUI core at $COMFYUI_CORE_REF"
  if [[ "$COMFYUI_CORE_REF" == v* ]]; then
    git -C "$COMFYUI_DIR" fetch --force origin "refs/tags/$COMFYUI_CORE_REF:refs/tags/$COMFYUI_CORE_REF"
  else
    git -C "$COMFYUI_DIR" fetch --force origin "$COMFYUI_CORE_REF"
  fi
  if ! git -C "$COMFYUI_DIR" checkout -f "$COMFYUI_CORE_REF"; then
    die "Could not check out ComfyUI core ref: $COMFYUI_CORE_REF"
  fi

  log "ComfyUI core resolved to $(git -C "$COMFYUI_DIR" describe --tags --always --dirty)"
}

install_ltxvideo_nodes() {
  local target="$CUSTOM_NODES_DIR/ComfyUI-LTXVideo"

  log "Installing official LTXVideo custom nodes at $LTXVIDEO_REVISION"
  if [[ ! -d "$target/.git" ]]; then
    rm -rf "$target"
    git clone "$LTXVIDEO_REPO" "$target"
  fi

  if ! git -C "$target" fetch origin "$LTXVIDEO_REVISION"; then
    git -C "$target" fetch origin
  fi

  if ! git -C "$target" checkout -f "$LTXVIDEO_REVISION"; then
    log "Existing LTXVideo checkout could not switch revisions; re-cloning."
    rm -rf "$target"
    git clone "$LTXVIDEO_REPO" "$target"
    git -C "$target" fetch origin "$LTXVIDEO_REVISION" || git -C "$target" fetch origin
    git -C "$target" checkout -f "$LTXVIDEO_REVISION"
  fi
}

COMFYUI_DIR="${COMFYUI_DIR:-/workspace/runpod-slim/ComfyUI}"
[[ -f "$COMFYUI_DIR/main.py" ]] || die "Could not find ComfyUI at $COMFYUI_DIR. Use RunPod's official ComfyUI - CUDA 12.8 template or set COMFYUI_DIR=/path/to/ComfyUI."
CUSTOM_NODES_DIR="$COMFYUI_DIR/custom_nodes"
MODELS_DIR="$COMFYUI_DIR/models"
INPUT_DIR="$COMFYUI_DIR/input"
VENV_DIR="$COMFYUI_DIR/.venv-cu128"
LOG_FILE="/workspace/runpod-slim/comfyui-restart.log"
PID_FILE="/workspace/runpod-slim/comfyui-restart.pid"
PYTHON_BIN="$VENV_DIR/bin/python"

if [[ ! -x "$PYTHON_BIN" ]]; then
  PYTHON_BIN="$COMFYUI_DIR/venv/bin/python"
fi
if [[ ! -x "$PYTHON_BIN" ]]; then
  PYTHON_BIN="$(command -v python3.12 || command -v python3 || true)"
fi
[[ -x "$PYTHON_BIN" ]] || die "Could not find Python. Expected $COMFYUI_DIR/.venv-cu128/bin/python or python3 on PATH."

log "Using ComfyUI at $COMFYUI_DIR"
log "Using Python at $PYTHON_BIN"
log "Requested ComfyUI core ref: $COMFYUI_CORE_REF"
log "Requested workflow revision: $WORKFLOW_REVISION"

mkdir -p "$CUSTOM_NODES_DIR" "$MODELS_DIR/checkpoints" "$MODELS_DIR/loras" "$MODELS_DIR/text_encoders" "$INPUT_DIR"

command -v curl >/dev/null 2>&1 || die "curl is required. RunPod's official ComfyUI - CUDA 12.8 template includes it."
command -v git >/dev/null 2>&1 || die "git is required. RunPod's official ComfyUI - CUDA 12.8 template includes it."
command -v ffmpeg >/dev/null 2>&1 || die "ffmpeg is required. RunPod's official ComfyUI - CUDA 12.8 template includes it."
command -v sha256sum >/dev/null 2>&1 || die "sha256sum is required. RunPod's official ComfyUI - CUDA 12.8 template includes it."

install_comfyui_core
install_comfyui_core_requirements

install_ltxvideo_nodes

if [[ -f "$CUSTOM_NODES_DIR/ComfyUI-LTXVideo/requirements.txt" ]]; then
  pip_install --upgrade -r "$CUSTOM_NODES_DIR/ComfyUI-LTXVideo/requirements.txt"
fi
if [[ -f "$CUSTOM_NODES_DIR/ComfyUI-LTXVideo/install.py" ]]; then
  "$PYTHON_BIN" "$CUSTOM_NODES_DIR/ComfyUI-LTXVideo/install.py" || log "LTXVideo install.py returned non-zero; continuing after requirements install."
fi
patch_ltxvideo_kornia_import

log "Installing LTX Foley helper node and workflow from $WORKFLOW_REPO"
NODE_DEST="$CUSTOM_NODES_DIR/ltx_foley_v2a"
WORKFLOW_DEST="$COMFYUI_DIR/user/default/workflows"
rm -rf "$NODE_DEST"
mkdir -p "$NODE_DEST" "$WORKFLOW_DEST"
rm -f \
  "$WORKFLOW_DEST/ltx_23_foley_v2a_v3_loop.json" \
  "$WORKFLOW_DEST/ltx_23_foley_v2a_legacy_audio_vae.json" \
  "$WORKFLOW_DEST/ltx_23_foley_v2a_v1_wav_debug.json"
download_file "$WORKFLOW_BASE_URL/ltx_foley_v2a/__init__.py" "$NODE_DEST/__init__.py"
download_file "$WORKFLOW_BASE_URL/ltx_foley_v2a/nodes.py" "$NODE_DEST/nodes.py"
download_file "$WORKFLOW_BASE_URL/ltx_23_foley_v2a.json" "$WORKFLOW_DEST/ltx_23_foley_v2a.json"
download_file "$WORKFLOW_BASE_URL/foley-sliding-window.json" "$WORKFLOW_DEST/foley-sliding-window.json"
download_file "$WORKFLOW_BASE_URL/tennis-no-sound.mp4" "$INPUT_DIR/input.mp4"

log "Downloading LTX 2.3 Foley model files"
download_model_file \
  "https://huggingface.co/Lightricks/LTX-2.3-fp8/resolve/main/ltx-2.3-22b-dev-fp8.safetensors" \
  "$MODELS_DIR/checkpoints/ltx-2.3-22b-dev-fp8.safetensors" \
  "28606c5b5a06ce56f896d4dfcb20f212739e07a68fbe48e53638188449d26450"
download_model_file \
  "https://huggingface.co/Comfy-Org/ltx-2/resolve/main/split_files/text_encoders/gemma_3_12B_it_fp8_scaled.safetensors" \
  "$MODELS_DIR/text_encoders/gemma_3_12B_it_fp8_scaled.safetensors" \
  "60216ce97c01c3a8753c2dfd0a89fc76e16fbe446d1a32ef8f1b528ac8bae466"
download_model_file \
  "https://huggingface.co/FuzzPuppy/LTX-2.3-Foley-LoRA/resolve/main/ltx-2.3-foley-400-steps.safetensors" \
  "$MODELS_DIR/loras/ltx-2.3-foley-400-steps.safetensors" \
  "12b37ffff8444ac1348a8c9fb24843a6bd70e77cf4202ac39b1d056b81fa41c9"

log "Done"
restart_comfyui

cat <<EOF

Next steps:
1. Open the RunPod HTTP service on port 8188.
2. Load workflow: user/default/workflows/foley-sliding-window.json
   Short-clip workflow: user/default/workflows/ltx_23_foley_v2a.json
3. Queue the workflow. The tennis test video is already installed as input.mp4 and the prompt is prefilled.

This run used:
  COMFYUI_CORE_REF=$COMFYUI_CORE_REF
  WORKFLOW_REVISION=$WORKFLOW_REVISION

If model download fails with 401/403, accept the model licenses on Hugging Face and rerun with:
  export HF_TOKEN=hf_...
  bash $0
EOF