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

COMFY="${COMFYUI_PATH:-/app/ComfyUI}"

# ------------------------------------------------------------
# Runtime directories
# ------------------------------------------------------------
# This is intentionally done before ComfyUI starts.
# ComfyUI 0.35.x creates /app/ComfyUI/user itself.
mkdir -p \
  "$COMFY/user" \
  "$COMFY/input" \
  "$COMFY/output" \
  "$COMFY/temp" \
  "$COMFY/workflows" \
  "$COMFY/models/unet" \
  "$COMFY/models/checkpoints" \
  "$COMFY/models/diffusion_models" \
  "$COMFY/models/audio_encoders" \
  "$COMFY/models/text_encoders" \
  "$COMFY/models/vae"

# ------------------------------------------------------------
# ComfyUI Manager Security Setup
# ------------------------------------------------------------
MANAGER_CONFIG_DIR="$COMFY/user/default/ComfyUI-Manager"
# mkdir -p "$MANAGER_CONFIG_DIR"
# cat << 'EOF' > "$MANAGER_CONFIG_DIR/config.ini"
# [ComfyUI-Manager]
# security_level = weak
# privacy_mode = false
# EOF
# echo "[INFO] ComfyUI Manager security level set to weak."

# ------------------------------------------------------------
# Persistent storage
# ------------------------------------------------------------
# If HF Storage is mounted at /data, models survive Space restarts.
# Without persistent storage, downloaded models live in the container
# filesystem and may have to be downloaded again after a rebuild/restart.
PERSIST_ROOT="/data/comfyui"

if [ -d "/data" ] && [ -w "/data" ]; then
  mkdir -p "$PERSIST_ROOT/models"
else
  PERSIST_ROOT="$COMFY"
fi

MODEL_ROOT="$PERSIST_ROOT/models"

# ------------------------------------------------------------
# Model configuration
# ------------------------------------------------------------
# Define models to download.
# Format: "repo_id|filename|dest_subdir"
# dest_subdir should be one of: unet, text_encoders, vae, loras, controlnet, workflows, etc.
# 
# You can override this list by setting the MODELS_TO_DOWNLOAD environment variable.
# Example: export MODELS_TO_DOWNLOAD="repo1|file1|unet;repo2|file2|vae"

if [ -n "${MODELS_TO_DOWNLOAD:-}" ]; then
  IFS=';' read -r -a MODEL_LIST <<< "$MODELS_TO_DOWNLOAD"
else
  MODEL_LIST=(
    "vantagewithai/Krea-2-Turbo-GGUF|krea2_turbo-Q5_K_M.gguf|unet"
    "unsloth/FLUX.2-klein-4B-GGUF|flux-2-klein-4b-Q5_K_M.gguf|unet"
    "InsecureErasure/Krea2-Turbo-mixed-NVFP4|krea2_turbo_mixed_nvfp4.safetensors|diffusion_models"
    "Comfy-Org/Qwen3-VL|text_encoders/qwen3vl_4b_fp8_scaled.safetensors|text_encoders"
    "Comfy-Org/Qwen-Image_ComfyUI|split_files/vae/qwen_image_vae.safetensors|vae"
    "realrebelai/KREA-2_GGUFs|Rebels KREA-2-TURBO.json|workflows"
    "Comfy-Org/YuE2|checkpoints/yue2_3b_int8_convrot.safetensors|checkpoints"
    "Comfy-Org/YuE2|audio_encoders/sheetsage2_bf16.safetensors|audio_encoders"
    #"Sporer/flux2-klein-4b-nvfp4|flux-2-klein-4b-nvfp4.safetensors|diffusion_models"
    "unsloth/FLUX.2-VAE|split_files/vae/flux2-vae.safetensors|vae"
    # เพิ่มโมเดลอื่นๆ ได้ที่นี่ เช่น:
    # "some-repo/some-model|path/to/model.safetensors|loras"audio_encoders
  )
fi

# ------------------------------------------------------------
# Model directory links
# ------------------------------------------------------------
# Only replace an empty ComfyUI model directory.
# This avoids deleting any models supplied in the Space image.
link_model_dir() {
  local name="$1"
  local target="$COMFY/models/$name"
  local source="$MODEL_ROOT/$name"

  mkdir -p "$source"

  if [ "$source" = "$target" ]; then
    return
  fi

  if [ -L "$target" ]; then
    return
  fi

  if [ -d "$target" ]; then
    if [ -z "$(find "$target" -mindepth 1 -maxdepth 1 -print -quit 2>/dev/null)" ]; then
      rmdir "$target"
      ln -s "$source" "$target"
    else
      echo "[INFO] Keeping non-empty model directory: $target"
    fi
  fi
}

# Dynamically link model directories based on MODEL_LIST
UNIQUE_SUBDIRS=$(printf "%s\n" "${MODEL_LIST[@]}" | awk -F'|' '$3 != "" {print $3}' | grep -v "^workflows$" | sort -u)
for subdir in $UNIQUE_SUBDIRS; do
  link_model_dir "$subdir"
done

echo "============================================================"
echo "ComfyUI Krea2 - Hugging Face CPU"
echo "============================================================"

python - <<'PY'
import os
import torch

print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("CPU threads:", os.environ.get("OMP_NUM_THREADS"))
print("RAM limit is controlled by the Hugging Face Space hardware.")
PY

# ------------------------------------------------------------
# Hugging Face downloader
# ------------------------------------------------------------
download_hf() {
  local repo="$1"
  local file="$2"
  local dest="$3"

  if [ -s "$dest" ]; then
    echo "[OK] Existing: $dest"
    return 0
  fi

  echo "[DOWNLOAD] $repo :: $file"

  python - "$repo" "$file" "$dest" <<'PY'
import os
import sys
from huggingface_hub import hf_hub_download

repo, filename, dest = sys.argv[1:4]
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")

path = hf_hub_download(
    repo_id=repo,
    filename=filename,
    token=token,
    local_dir=os.path.dirname(dest),
)

if os.path.abspath(path) != os.path.abspath(dest):
    if os.path.exists(dest):
        os.remove(dest)
    os.replace(path, dest)

print("[OK] downloaded:", dest)
PY
}

# ------------------------------------------------------------
# Download models
# ------------------------------------------------------------
if [ "${KREA2_DOWNLOAD_MODELS:-1}" = "1" ]; then
  for model_entry in "${MODEL_LIST[@]}"; do
    [ -z "$model_entry" ] && continue
    IFS='|' read -r repo file subdir <<< "$model_entry"
    
    if [ "$subdir" = "workflows" ]; then
      dest_dir="$COMFY/workflows"
    else
      dest_dir="$COMFY/models/$subdir"
    fi
    
    mkdir -p "$dest_dir"
    dest_path="$dest_dir/$(basename "$file")"
    
    download_hf "$repo" "$file" "$dest_path"
  done
else
  echo "[INFO] KREA2_DOWNLOAD_MODELS=0; model download skipped."
fi

# ------------------------------------------------------------
# CPU optimization
# ------------------------------------------------------------
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-2}"
export MKL_NUM_THREADS="${MKL_NUM_THREADS:-2}"
export OPENBLAS_NUM_THREADS="${OPENBLAS_NUM_THREADS:-2}"
export NUMEXPR_NUM_THREADS="${NUMEXPR_NUM_THREADS:-2}"
export TOKENIZERS_PARALLELISM=false

cd "$COMFY"

echo "============================================================"
echo "Starting ComfyUI"
echo "============================================================"

exec python -u main.py \
  --cpu \
  --listen 0.0.0.0 \
  --port "${PORT:-7860}" \
  --output-directory "$COMFY/output"