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

log()  { echo "[animatorlipsync] $*"; }
warn() { echo "[animatorlipsync][WARN] $*"; }

export DEBIAN_FRONTEND=noninteractive
export PIP_ROOT_USER_ACTION=ignore

COMFYUI_DIR="${COMFYUI_DIR:-/workspace/ComfyUI}"
CUSTOM_NODES_DIR="${COMFYUI_DIR}/custom_nodes"
MODELS_DIR="${COMFYUI_DIR}/models"
WORKFLOWS_DIR="${COMFYUI_DIR}/user/default/workflows"
INPUT_DIR="${COMFYUI_DIR}/input"
OUTPUT_DIR="${COMFYUI_DIR}/output"
TEMP_DIR="${COMFYUI_DIR}/temp"

HF_REPO_BASE="${HF_REPO_BASE:-https://huggingface.co/realgen4/OFMHUB-assets/resolve/main}"
WORKFLOW_JSON_URL="${WORKFLOW_JSON_URL:-${HF_REPO_BASE}/lipsync.json}"
WORKFLOW_JSON_NAME="${WORKFLOW_JSON_NAME:-lipsync.json}"
RIFE_FLOWNET_URL="${RIFE_FLOWNET_URL:-https://github.com/hzwer/Practical-RIFE/raw/main/train_log/flownet.pkl}"

log "Starting animatorlipsync provisioning"

if [ -f /venv/main/bin/activate ]; then
  # shellcheck disable=SC1091
  source /venv/main/bin/activate
  log "Activated /venv/main"
else
  warn "/venv/main/bin/activate not found, continuing without explicit venv activation"
fi

if [ ! -d "${COMFYUI_DIR}" ]; then
  warn "ComfyUI directory not found at ${COMFYUI_DIR}"
  exit 1
fi

mkdir -p "${CUSTOM_NODES_DIR}" "${WORKFLOWS_DIR}" "${INPUT_DIR}" "${OUTPUT_DIR}" "${TEMP_DIR}"
mkdir -p \
  "${MODELS_DIR}/diffusion_models" \
  "${MODELS_DIR}/vae" \
  "${MODELS_DIR}/text_encoders" \
  "${MODELS_DIR}/clip_vision" \
  "${MODELS_DIR}/controlnet" \
  "${MODELS_DIR}/loras" \
  "${MODELS_DIR}/detection" \
  "${MODELS_DIR}/sam2" \
  "${MODELS_DIR}/rife"

apt-get update -y || true
apt-get install -y \
  git wget curl rsync ffmpeg \
  libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 \
  ca-certificates || true

python3 -m pip install -U pip setuptools wheel || true
python3 -m pip install \
  accelerate \
  opencv-python-headless \
  imageio-ffmpeg \
  safetensors \
  einops \
  pillow \
  huggingface_hub || true

install_repo_node() {
  local repo_url="$1"
  local dir_name="$2"
  local target="${CUSTOM_NODES_DIR}/${dir_name}"

  if [ -d "${target}/.git" ]; then
    log "Updating ${dir_name}"
    git -C "${target}" pull --ff-only || warn "git pull failed for ${dir_name}"
  elif [ -d "${target}" ]; then
    log "${dir_name} already exists"
  else
    log "Cloning ${dir_name}"
    git clone --depth=1 "${repo_url}" "${target}"
  fi

  if [ -f "${target}/requirements.txt" ]; then
    log "Installing requirements for ${dir_name}"
    python3 -m pip install -r "${target}/requirements.txt" || warn "requirements failed for ${dir_name}"
  fi

  if [ -f "${target}/install.py" ]; then
    log "Running install.py for ${dir_name}"
    python3 "${target}/install.py" || warn "install.py failed for ${dir_name}"
  fi
}

download_if_missing() {
  local url="$1"
  local out="$2"

  if [ -f "${out}" ]; then
    log "Exists: ${out}"
    return 0
  fi

  mkdir -p "$(dirname "${out}")"
  log "Downloading $(basename "${out}")"
  wget -O "${out}" --tries=5 --timeout=120 "${url}" || {
    warn "Failed download: ${url}"
    rm -f "${out}"
    return 1
  }
}

download_force() {
  local url="$1"
  local out="$2"

  mkdir -p "$(dirname "${out}")"
  rm -f "${out}"
  log "Downloading (force) $(basename "${out}")"
  wget -O "${out}" --tries=5 --timeout=120 "${url}" || {
    warn "Failed forced download: ${url}"
    rm -f "${out}"
    return 1
  }
}

copy_if_missing() {
  local src="$1"
  local dst="$2"

  if [ ! -f "${src}" ]; then
    warn "copy_if_missing source missing: ${src}"
    return 1
  fi

  if [ -f "${dst}" ]; then
    log "Exists: ${dst}"
    return 0
  fi

  mkdir -p "$(dirname "${dst}")"
  cp -f "${src}" "${dst}"
  log "Copied alias $(basename "${src}") -> ${dst}"
}

install_ts_node_from_hf() {
  local target_dir="${CUSTOM_NODES_DIR}/NEW-UTILS"

  rm -rf "${target_dir}"
  mkdir -p "${target_dir}/nodes"
  mkdir -p "${target_dir}/web/js"

  download_if_missing "${HF_REPO_BASE}/__init__.py" "${target_dir}/__init__.py" || true
  download_if_missing "${HF_REPO_BASE}/requirements.txt" "${target_dir}/requirements.txt" || true

  download_if_missing "${HF_REPO_BASE}/nodes/color_match.py" "${target_dir}/nodes/color_match.py" || true
  download_if_missing "${HF_REPO_BASE}/nodes/openpose_smoother.py" "${target_dir}/nodes/openpose_smoother.py" || true
  download_if_missing "${HF_REPO_BASE}/nodes/preview_image_metadata.py" "${target_dir}/nodes/preview_image_metadata.py" || true
  download_if_missing "${HF_REPO_BASE}/nodes/video_combine_metadata.py" "${target_dir}/nodes/video_combine_metadata.py" || true
  download_if_missing "${HF_REPO_BASE}/nodes/load_video_batch.py" "${target_dir}/nodes/load_video_batch.py" || true
  download_if_missing "${HF_REPO_BASE}/nodes/rename_files.py" "${target_dir}/nodes/rename_files.py" || true
  download_if_missing "${HF_REPO_BASE}/nodes/save_load_pose.py" "${target_dir}/nodes/save_load_pose.py" || true

  download_if_missing "${HF_REPO_BASE}/web/js/ts_video_preview.js" "${target_dir}/web/js/ts_video_preview.js" || true

  if [ -f "${target_dir}/requirements.txt" ]; then
    log "Installing requirements for NEW-UTILS"
    python3 -m pip install -r "${target_dir}/requirements.txt" || warn "requirements failed for NEW-UTILS"
  fi
}

log "Installing custom nodes"
install_repo_node "https://github.com/kijai/ComfyUI-WanVideoWrapper" "ComfyUI-WanVideoWrapper"
install_repo_node "https://github.com/kijai/ComfyUI-WanAnimatePreprocess" "ComfyUI-WanAnimatePreprocess"
install_repo_node "https://github.com/kijai/ComfyUI-segment-anything-2" "ComfyUI-segment-anything-2"
install_repo_node "https://github.com/kijai/ComfyUI-KJNodes" "ComfyUI-KJNodes"
install_repo_node "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite" "ComfyUI-VideoHelperSuite"
install_repo_node "https://github.com/rgthree/rgthree-comfy" "rgthree-comfy"
install_repo_node "https://github.com/PGCRT/CRT-Nodes" "CRT-Nodes"
install_repo_node "https://github.com/evanspearman/ComfyMath" "ComfyMath"
install_repo_node "https://github.com/GACLove/ComfyUI-VFI" "ComfyUI-VFI"
install_repo_node "https://github.com/Fannovel16/comfyui_controlnet_aux" "comfyui_controlnet_aux"
install_ts_node_from_hf

if [ -n "${WORKFLOW_JSON_URL}" ]; then
  log "Downloading workflow"
  download_if_missing "${WORKFLOW_JSON_URL}" "${WORKFLOWS_DIR}/${WORKFLOW_JSON_NAME}" || warn "Workflow download failed"
else
  warn "WORKFLOW_JSON_URL is empty, skipping workflow download. Put your JSON manually into ${WORKFLOWS_DIR}"
fi

log "Downloading models used by the workflow"
# Wan model stack

download_if_missing \
  "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/resolve/main/Wan22Animate/Wan2_2-Animate-14B_fp8_scaled_e4m3fn_KJ_v2.safetensors" \
  "${MODELS_DIR}/diffusion_models/Wan2_2-Animate-14B_fp8_scaled_e4m3fn_KJ_v2.safetensors" || true

download_if_missing \
  "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Wan2_1_VAE_bf16.safetensors" \
  "${MODELS_DIR}/vae/wan_2.1_vae.safetensors" || true

download_if_missing \
  "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" \
  "${MODELS_DIR}/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" || true

download_if_missing \
  "https://huggingface.co/calcuis/wan-gguf/resolve/f52f5a1f0ba441d50277fb7cdd7c1b36611837f9/clip_vision_h.safetensors" \
  "${MODELS_DIR}/clip_vision/clip_vision_h.safetensors" || true

download_if_missing \
  "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Wan21_Uni3C_controlnet_fp16.safetensors" \
  "${MODELS_DIR}/controlnet/Wan21_Uni3C_controlnet_fp16.safetensors" || true

download_if_missing \
  "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors" \
  "${MODELS_DIR}/loras/lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors" || true

download_if_missing \
  "https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/resolve/main/split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors" \
  "${MODELS_DIR}/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors" || true

download_if_missing \
  "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Pusa/Wan21_PusaV1_LoRA_14B_rank512_bf16.safetensors" \
  "${MODELS_DIR}/loras/Wan21_PusaV1_LoRA_14B_rank512_bf16.safetensors" || true

download_if_missing \
  "https://huggingface.co/rahul7star/wan2.2Lora/resolve/main/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors" \
  "${MODELS_DIR}/loras/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors" || true

# WanAnimatePreprocess models -> ComfyUI/models/detection

download_if_missing \
  "https://huggingface.co/JunkyByte/easy_ViTPose/resolve/main/onnx/wholebody/vitpose-l-wholebody.onnx" \
  "${MODELS_DIR}/detection/vitpose-l-wholebody.onnx" || true

download_if_missing \
  "https://huggingface.co/onnx-community/yolov10m/resolve/main/onnx/model.onnx" \
  "${MODELS_DIR}/detection/yolov10m.onnx" || true

# SAM2

download_if_missing \
  "https://huggingface.co/Kijai/sam2-safetensors/resolve/main/sam2.1_hiera_base_plus.safetensors" \
  "${MODELS_DIR}/sam2/sam2.1_hiera_base_plus.safetensors" || true

# DWPose models for DWPreprocessor
# yolox_l.onnx comes from yzd-v/DWPose
# dw-ll_ucoco.onnx must come from hr16/UnJIT-DWPose (do not alias the 384 model here)
DWPOSE_DIR_YZD="${CUSTOM_NODES_DIR}/comfyui_controlnet_aux/ckpts/yzd-v/DWPose"
DWPOSE_DIR_HR16="${CUSTOM_NODES_DIR}/comfyui_controlnet_aux/ckpts/hr16/UnJIT-DWPose"
mkdir -p "${DWPOSE_DIR_YZD}" "${DWPOSE_DIR_HR16}"

download_if_missing \
  "https://huggingface.co/yzd-v/DWPose/resolve/main/yolox_l.onnx" \
  "${DWPOSE_DIR_YZD}/yolox_l.onnx" || true

# Keep the 384 ONNX available too; some other workflows may use it explicitly.
download_if_missing \
  "https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.onnx" \
  "${DWPOSE_DIR_YZD}/dw-ll_ucoco_384.onnx" || true

# Force-replace because an older broken template may have copied the 384 model here under the wrong name.
download_force \
  "https://huggingface.co/hr16/UnJIT-DWPose/resolve/main/dw-ll_ucoco.onnx" \
  "${DWPOSE_DIR_HR16}/dw-ll_ucoco.onnx" || true

# RIFE model for RIFEInterpolation -> ComfyUI/models/rife/flownet.pkl
# If this fails, ComfyUI-VFI can auto-download it on first use.
download_if_missing \
  "${RIFE_FLOWNET_URL}" \
  "${MODELS_DIR}/rife/flownet.pkl" || true

log "Validation"
missing=0
for path in \
  "${CUSTOM_NODES_DIR}/ComfyUI-WanVideoWrapper" \
  "${CUSTOM_NODES_DIR}/ComfyUI-WanAnimatePreprocess" \
  "${CUSTOM_NODES_DIR}/ComfyUI-segment-anything-2" \
  "${CUSTOM_NODES_DIR}/ComfyUI-KJNodes" \
  "${CUSTOM_NODES_DIR}/ComfyUI-VideoHelperSuite" \
  "${CUSTOM_NODES_DIR}/rgthree-comfy" \
  "${CUSTOM_NODES_DIR}/CRT-Nodes" \
  "${CUSTOM_NODES_DIR}/ComfyMath" \
  "${CUSTOM_NODES_DIR}/ComfyUI-VFI" \
  "${CUSTOM_NODES_DIR}/comfyui_controlnet_aux" \
  "${CUSTOM_NODES_DIR}/NEW-UTILS/__init__.py" \
  "${CUSTOM_NODES_DIR}/NEW-UTILS/nodes/color_match.py" \
  "${CUSTOM_NODES_DIR}/NEW-UTILS/nodes/openpose_smoother.py" \
  "${CUSTOM_NODES_DIR}/NEW-UTILS/nodes/preview_image_metadata.py" \
  "${CUSTOM_NODES_DIR}/NEW-UTILS/nodes/video_combine_metadata.py" \
  "${MODELS_DIR}/diffusion_models/Wan2_2-Animate-14B_fp8_scaled_e4m3fn_KJ_v2.safetensors" \
  "${MODELS_DIR}/vae/wan_2.1_vae.safetensors" \
  "${MODELS_DIR}/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" \
  "${MODELS_DIR}/clip_vision/clip_vision_h.safetensors" \
  "${MODELS_DIR}/controlnet/Wan21_Uni3C_controlnet_fp16.safetensors" \
  "${MODELS_DIR}/loras/lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors" \
  "${MODELS_DIR}/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors" \
  "${MODELS_DIR}/loras/Wan21_PusaV1_LoRA_14B_rank512_bf16.safetensors" \
  "${MODELS_DIR}/loras/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors" \
  "${MODELS_DIR}/detection/vitpose-l-wholebody.onnx" \
  "${MODELS_DIR}/detection/yolov10m.onnx" \
  "${MODELS_DIR}/sam2/sam2.1_hiera_base_plus.safetensors" \
  "${DWPOSE_DIR_YZD}/yolox_l.onnx" \
  "${DWPOSE_DIR_YZD}/dw-ll_ucoco_384.onnx" \
  "${DWPOSE_DIR_HR16}/dw-ll_ucoco.onnx"
do
  if [ -e "${path}" ]; then
    log "OK: ${path}"
  else
    warn "MISSING: ${path}"
    missing=1
  fi
done

if [ -n "${WORKFLOW_JSON_URL}" ]; then
  if [ -e "${WORKFLOWS_DIR}/${WORKFLOW_JSON_NAME}" ]; then
    log "OK: ${WORKFLOWS_DIR}/${WORKFLOW_JSON_NAME}"
  else
    warn "MISSING: ${WORKFLOWS_DIR}/${WORKFLOW_JSON_NAME}"
    missing=1
  fi
else
  warn "Workflow file was not validated because WORKFLOW_JSON_URL is empty"
fi

if [ -e "${MODELS_DIR}/rife/flownet.pkl" ]; then
  log "OK: ${MODELS_DIR}/rife/flownet.pkl"
else
  warn "RIFE flownet.pkl not preloaded. ComfyUI-VFI should auto-download it on first use."
fi

if [ "${missing}" -eq 1 ]; then
  warn "Provisioning completed with missing files"
else
  log "Provisioning completed successfully"
fi

log "Provisioning finished"