#!/usr/bin/env bash
# setup_comfyui_wan22_fixed.sh
# Script corrigido para instalação do ComfyUI com Wan 2.2

set -euo pipefail

# -----------------------------
# Cores para output
# -----------------------------
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info()    { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[✓]${NC} $1"; }
log_warn()    { echo -e "${YELLOW}[!]${NC} $1"; }
log_error()   { echo -e "${RED}[✗]${NC} $1"; }

# -----------------------------
# Configuração
# -----------------------------
COMFY_DIR="/root/comfy/ComfyUI"
MODELS_DIR="$COMFY_DIR/models"
COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
COMFY_PORT="${COMFY_PORT:-8818}"
CIVITAI_TOKEN="${CIVITAI_TOKEN:-}"   # NÃO deixe token hardcoded. Defina via env var.
HF_TOKEN="${HF_TOKEN:-}"             # Token opcional do HuggingFace

# Configurações de performance
export MAX_JOBS=32
export NVCC_APPEND_FLAGS="--threads 8"
export UV_SYSTEM_PYTHON=1
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
export HF_HUB_ENABLE_HF_TRANSFER=1
export HF_TRANSFER_CONCURRENCY=16
export EXT_PARALLEL=4

# -----------------------------
# Listas de download
# -----------------------------
readonly DOWNLOAD_FILES=(
  # Wan 2.2 models
  "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors"
  "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors"
  "hf://Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors"
  "hf://Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors"
  "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors"
  "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors"
  "hf://Kijai/WanVideo_comfy/Wan2_1_VAE_fp32.safetensors|vae|Wan2_1_VAE_fp32.safetensors"
  "hf://Kijai/WanVideo_comfy/umt5-xxl-enc-bf16.safetensors|text_encoders|umt5-xxl-enc-bf16.safetensors"
  # Wan 2.1 models
  "hf://Kijai/WanVideo_comfy/Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors|loras|lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors"
  "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/vae/wan_2.1_vae.safetensors|vae|wan_2.1_vae.safetensors"
  "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/text_encoders/umt5_xxl_fp16.safetensors|text_encoders|umt5_xxl_fp16.safetensors"
  # Wan 2.1 clip vision
  "hf://Kijai/WanVideo_comfy/open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors|clip_vision|open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors"
  "hf://Comfy-Org/Wan_2.1_ComfyUI_repackaged/split_files/clip_vision/clip_vision_h.safetensors|clip_vision|clip_vision_h.safetensors"
  # Upscaler extra
  "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
)

readonly MEGA_LORAS=(
  "https://mega.nz/file/QFxREJpZ#0dge8TyRgb52lMFKdsk-Ac-rwcoG49hvBHvaqNsZkuk"
  "https://mega.nz/file/lIhQETjT#nKlQ01Euw6fUxo8P5bfwlQrxAkWEedXQQOqHv5u4yrM"
  "https://mega.nz/file/UIAESACY#Qlei1Pj5Nwno3Sz-wsTekvD-YgCkN1A5QbFKoRHXG9E"
  "https://mega.nz/file/cdAg1Q7D#tQVrer2sB-M7XdYiKAqr7kRRfJw4pXCdliOg2aPJUGE"
  "https://mega.nz/file/5RQFiAAJ#QCWrHkuJjXlFufR18g5R9YXJGqJ1pQp7xbqNRgSYUjc"
  "https://mega.nz/file/tABgxLTS#NNlPV8NPc7R52MIjABMFUb41EOFDxFoszBxOrPMN13Q"
  "https://mega.nz/file/4MxBgJ5Y#R8OSby0u-t-QFBnpnOtdjN3xuiVRvHa4_OEzTwWoGKQ"
  "https://mega.nz/file/dJJxiSRR#o3XsYfSvb04xg7gNVdfHb3XTVHgvx0ZyXe9N8VQ9Lfw"
  "https://mega.nz/file/wcAFhRJZ#ifZcckts0wPq1FG693tYLkpFrvjXvab1ZQeL6DuL1T8"
  "https://mega.nz/file/wNIgmRwb#vn6rKM3QuTDpWVTLEBDsKjIkNkaX2wLKeauKNeVkcFs"
  "https://mega.nz/file/hZhhhCQS#4je0Im2g1kEvTCuPotjVWkA-g6XICkqMb-mKxMt6R3Y"
)

readonly MEGA_UPSCALE=(
  "https://mega.nz/file/qdQwAJSC#rPY2Hg4x-_RbHiHUBFW-YU8xZUqPpq5Gpl1uPFiJvW4"
  "https://mega.nz/file/KIh2marA#CehVNXxQhGfp1bqUL0K1YPWNqrCUoLPCZGq2W6XqmWo"
)

# Declarado (vazio) para não quebrar com set -u; preencha se quiser baixar RIFE TensorRT
readonly RIFE_MODELS=()

# -----------------------------
# Funções auxiliares
# -----------------------------
command_exists() { command -v "$1" >/dev/null 2>&1; }

# APT helpers com tolerância a falhas
apt_update_safe() {
  if command_exists apt-get; then
    export DEBIAN_FRONTEND=noninteractive
    apt-get update -y || { log_warn "apt-get update falhou; tentando continuar mesmo assim."; return 0; }
  else
    log_warn "apt-get não disponível; pulando update."
  fi
}

apt_install_safe() {
  # uso: apt_install_safe pkg1 pkg2 ...
  if command_exists apt-get; then
    export DEBIAN_FRONTEND=noninteractive
    apt-get install -y "$@" || log_warn "Falha ao instalar: $* (continuando)."
  else
    log_warn "apt-get não disponível; não foi possível instalar: $*"
  fi
}

# Adicionar token do Civitai se necessário (usa ? ou & corretamente)
add_civitai_token() {
  local url="$1"
  if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]] && [[ -n "$CIVITAI_TOKEN" ]]; then
    if [[ "$url" == *"?"* ]]; then
      echo "${url}&token=${CIVITAI_TOKEN}"
    else
      echo "${url}?token=${CIVITAI_TOKEN}"
    fi
  else
    echo "$url"
  fi
}

# Download HuggingFace com hf/huggingface-cli + hf_transfer
download_hf() {
  local repo="$1"; local file_path="$2"; local target_dir="$3"; local filename="$4"
  if [ -z "$filename" ]; then filename=$(basename "$file_path"); fi
  local target_file="$target_dir/$filename"

  if [ -f "$target_file" ] && [ "$(stat -c%s "$target_file" 2>/dev/null || echo 0)" -gt 1000000 ]; then
    log_success "Arquivo já existe: $filename"; return 0
  fi

  log_info "Baixando de HF: $filename"
  local temp_dir; temp_dir="$(mktemp -d)"
  local download_success=false

  if command_exists hf; then
    HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo" "$file_path" \
      --local-dir "$temp_dir" --local-dir-use-symlinks False && download_success=true || true
  fi
  if [ "$download_success" = false ] && command_exists huggingface-cli; then
    log_warn "Tentando com huggingface-cli..."
    HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" "$file_path" \
      --local-dir "$temp_dir" --local-dir-use-symlinks False && download_success=true || true
  fi

  if [ "$download_success" = true ]; then
    local downloaded_file
    downloaded_file="$(find "$temp_dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" \) | head -1)"
    if [ -n "$downloaded_file" ]; then
      mv "$downloaded_file" "$target_file"
      rm -rf "$temp_dir"
      log_success "Download concluído: $filename"
      return 0
    else
      log_error "Arquivo baixado mas não encontrado em $temp_dir"
    fi
  fi

  rm -rf "$temp_dir"
  log_error "Falha ao baixar: $filename"
  return 1
}

# Download genérico com aria2c, wget ou curl
download_file() {
  local url="$1"; local target_dir="$2"; local filename="$3"
  url=$(add_civitai_token "$url")

  if [ -n "$filename" ] && [ -f "$target_dir/$filename" ] && [ "$(stat -c%s "$target_dir/$filename" 2>/dev/null || echo 0)" -gt 1000000 ]; then
    log_success "Arquivo já existe: $filename"; return 0
  fi

  log_info "Baixando: ${filename:-$(basename "$url" | cut -d'?' -f1)}"

  if command_exists aria2c; then
    local aria_opts="-c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10"
    if [ -n "$filename" ]; then aria2c $aria_opts --dir="$target_dir" --out="$filename" "$url" && return 0 || true
    else aria2c $aria_opts --dir="$target_dir" "$url" && return 0 || true
    fi
  fi
  if command_exists wget; then
    if [ -n "$filename" ]; then wget -q --show-progress -c -O "$target_dir/$filename" "$url" && return 0 || true
    else wget -q --show-progress -c --content-disposition -P "$target_dir" "$url" && return 0 || true
    fi
  fi
  if command_exists curl; then
    if [ -n "$filename" ]; then
      curl -L -# -C - -o "$target_dir/$filename" "$url" && return 0 || true
    else
      local headers detected_name
      headers="$(curl -sI -L "$url")"
      detected_name="$(echo "$headers" | grep -i "content-disposition" | sed -n 's/.*filename="\([^"]*\)".*/\1/p' | tr -d '\r')"
      [ -z "$detected_name" ] && detected_name="downloaded_$(date +%s).bin"
      curl -L -# -C - -o "$target_dir/$detected_name" "$url" && return 0 || true
    fi
  fi

  log_error "Falha ao baixar: $url"
  return 1
}

# Download de Mega.nz usando megatools (megadl)
download_mega() {
  local url="$1"; local target_dir="$2"
  if ! command_exists megadl; then
    log_warn "megadl não está instalado. Tentando instalar megatools..."
    apt_update_safe
    apt_install_safe megatools
  fi
  if ! command_exists megadl; then
    log_warn "megadl indisponível. Pulando download de Mega.nz: $url"
    return 1
  fi

  log_info "Baixando de Mega.nz: $url"
  ( cd "$target_dir" && megadl "$url" ) && { log_success "Download Mega.nz concluído"; return 0; } || {
    log_warn "Falha ao baixar de Mega.nz: $url"; return 1;
  }
}

# Processar lista de downloads
process_downloads() {
  local failed_count=0
  for entry in "${DOWNLOAD_FILES[@]}"; do
    IFS='|' read -r url type filename <<< "$entry"
    url="$(echo "$url" | xargs)"; type="$(echo "$type" | xargs)"; filename="$(echo "$filename" | xargs)"

    local target_dir="$MODELS_DIR/$type"
    case "$type" in
      checkpoints|diffusion_models|loras|vae|vae_approx|text_encoders|clip_vision|controlnet|upscale_models|tensorrt|tensorrt_rife) ;;
      *) target_dir="$MODELS_DIR/$type" ;;
    esac
    mkdir -p "$target_dir"

    if [[ "$url" == hf://* ]]; then
      url="${url#hf://}"
      local repo="$(echo "$url" | cut -d'/' -f1-2)"
      local file_path="$(echo "$url" | cut -d'/' -f3-)"
      download_hf "$repo" "$file_path" "$target_dir" "$filename" || ((failed_count++))
    else
      download_file "$url" "$target_dir" "$filename" || ((failed_count++))
    fi
  done
  return $failed_count
}

# Clonar ou atualizar repositório git
clone_or_update() {
  local url="$1"; local dest="$2"; local node_name
  node_name="$(basename "$dest")"

  if [ -d "$dest/.git" ]; then
    log_info "Atualizando $node_name..."
    timeout 60 git -C "$dest" pull --ff-only || log_warn "Timeout/erro ao atualizar $node_name"
  else
    log_info "Clonando $node_name..."
    timeout 60 git clone --recursive --depth 1 "$url" "$dest" || log_warn "Falha ao clonar $node_name"
  fi

  if [ -f "$dest/requirements.txt" ]; then
    timeout 180 python3 -m pip install --no-warn-script-location -q -r "$dest/requirements.txt" || log_warn "Falha ao instalar requirements para $node_name"
  fi
}

# -----------------------------
# Instalação principal
# -----------------------------
echo ""
log_info "========================================="
log_info " ComfyUI + Wan 2.2 Setup (Fixed v3)"
log_info "========================================="
echo ""

# [1] Verificar e instalar dependências
log_info "[1/10] Verificando e instalando dependências..."

if ! command_exists aria2c || ! command_exists megadl; then
  log_info "Instalando ferramentas de download..."
  apt_update_safe
  command_exists aria2c || apt_install_safe aria2
  command_exists megadl || apt_install_safe megatools
fi

for cmd in python3 git wget curl; do
  if ! command_exists "$cmd"; then
    log_error "Dependência faltando: $cmd"
    exit 1
  fi
done
log_success "Dependências OK"

# [2] Atualizar pip e instalar ferramentas
log_info "[2/10] Preparando ferramentas Python..."
python3 -m pip install --upgrade pip wheel setuptools -q
python3 -m pip install --upgrade "huggingface_hub[cli,hf_transfer]>=0.26.0" comfy-cli -q

if [ -n "$HF_TOKEN" ]; then
  if command_exists huggingface-cli; then
    huggingface-cli login --token "$HF_TOKEN" --add-to-git-credential || log_warn "Falha no login HF (seguindo sem auth)"
  fi
fi

if command_exists hf; then
  log_success "Comando 'hf' disponível (recomendado)"
elif command_exists huggingface-cli; then
  log_warn "Usando 'huggingface-cli' (funcional)"
else
  log_warn "Nenhum comando HF disponível — downloads HF podem falhar."
fi
log_success "Ferramentas instaladas"

# [3] PyTorch CUDA 12.8
log_info "[3/10] Instalando PyTorch com CUDA 12.8..."
python3 -m pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
log_success "PyTorch com CUDA 12.8 instalado"

# [4] Instalar ComfyUI
log_info "[4/10] Instalando ComfyUI..."
if [ -f "$COMFY_DIR/main.py" ]; then
  log_warn "ComfyUI já existe, pulando instalação base"
else
  comfy --skip-prompt install --fast-deps --nvidia
  if [ ! -f "$COMFY_DIR/main.py" ]; then
    log_error "ComfyUI não foi instalado corretamente!"
    exit 1
  fi
fi
log_success "ComfyUI instalado em: $COMFY_DIR"

# [5] SageAttention (wheel pré-compilada)
log_info "[5/10] Instalando SageAttention..."
if python3 - <<'PY' 2>/dev/null | grep -q True; then
import torch; print(torch.cuda.is_available())
PY
  python3 -m pip install "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" -q || log_warn "Falha ao instalar SageAttention"
  python3 - <<'PY' >/dev/null 2>&1 || log_warn "SageAttention não pôde ser importado (seguindo sem)"
import sageattention
PY
  log_success "SageAttention etapa concluída"
else
  log_warn "CUDA não detectado, pulando SageAttention"
fi

# [6] Baixar modelos principais
log_info "[6/10] Baixando modelos principais..."
mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,vae_approx,text_encoders,clip_vision,controlnet,upscale_models,checkpoints}
if command_exists aria2c; then
  log_success "Usando aria2c para downloads (rápido)"
else
  log_warn "aria2c não encontrado, usando wget/curl (mais lento)"
fi
if process_downloads; then
  log_success "Todos os modelos baixados (ou já presentes)"
else
  log_warn "Alguns downloads falharam; você pode reexecutar depois."
fi

# [7] Modelos RIFE (opcional)
log_info "[7/10] Baixando modelos RIFE (opcional)..."
RIFE_DIR="$MODELS_DIR/tensorrt/rife"; mkdir -p "$RIFE_DIR"
if [ ${#RIFE_MODELS[@]} -gt 0 ]; then
  for url in "${RIFE_MODELS[@]}"; do
    filename="$(basename "$url")"
    if [ -f "$RIFE_DIR/$filename" ] && [ "$(stat -c%s "$RIFE_DIR/$filename" 2>/dev/null || echo 0)" -gt 100000 ]; then
      log_success "RIFE já existe: $filename"
    else
      log_info "Baixando RIFE: $filename"
      download_file "$url" "$RIFE_DIR" "$filename" || log_warn "Falha ao baixar RIFE: $filename"
    fi
  done
  log_success "Modelos RIFE instalados"
else
  log_warn "RIFE_MODELS vazio — pulando etapa (ok)."
fi

# [8] Mega.nz LoRAs e Upscale
log_info "[8/10] Baixando LoRAs/UpScales do Mega.nz..."
LORA_DIR="$MODELS_DIR/loras"; mkdir -p "$LORA_DIR"
for mega_url in "${MEGA_LORAS[@]}"; do
  download_mega "$mega_url" "$LORA_DIR" || log_warn "Falha ao baixar: $mega_url"
done
UPSCALE_DIR="$MODELS_DIR/upscale_models"; mkdir -p "$UPSCALE_DIR"
for mega_url in "${MEGA_UPSCALE[@]}"; do
  download_mega "$mega_url" "$UPSCALE_DIR" || log_warn "Falha ao baixar: $mega_url"
done
log_success "Downloads do Mega.nz concluídos (quando possível)."

# [9] Custom nodes
log_info "[9/10] Instalando custom nodes..."
CN_DIR="$COMFY_DIR/custom_nodes"; mkdir -p "$CN_DIR"
readonly CUSTOM_NODES=(
  "https://github.com/kijai/ComfyUI-Florence2"
  "https://github.com/adbrasi/groqrouter"
  "https://github.com/kijai/ComfyUI-WanVideoWrapper"
  "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"
  "https://github.com/kijai/ComfyUI-GIMM-VFI"
  "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"
  "https://github.com/kijai/ComfyUI-KJNodes"
  "https://github.com/Fannovel16/comfyui_controlnet_aux"
  "https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"
  "https://github.com/shiimizu/ComfyUI_smZNodes"
  "https://github.com/CoreyCorza/ComfyUI-CRZnodes"
  "https://github.com/pythongosssss/ComfyUI-Custom-Scripts"
  "https://github.com/pythongosssss/ComfyUI-WD14-Tagger"
  "https://github.com/WASasquatch/was-node-suite-comfyui"
  "https://github.com/justUmen/Bjornulf_custom_nodes"
  "https://github.com/mingsky-ai/ComfyUI-MingNodes"
  "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler"
  "https://github.com/1038lab/ComfyUI-RMBG"
  "https://github.com/wallish77/wlsh_nodes"
  "https://github.com/yuvraj108c/ComfyUI-Upscaler-Tensorrt"
  "https://github.com/raindrop313/ComfyUI-WanVideoStartEndFrames"
  "https://github.com/MushroomFleet/DJZ-Nodes"
  "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
  "https://github.com/grmchn/ComfyUI-ProportionChanger"
)
for repo_url in "${CUSTOM_NODES[@]}"; do
  node_name="$(basename "$repo_url")"
  clone_or_update "$repo_url" "$CN_DIR/$node_name"
done
log_success "Custom nodes instalados"

# [10] Verificação final
echo ""
log_info "[10/10] Verificando instalação..."
check_critical_files() {
  local files=(
    "$COMFY_DIR/main.py"
    "$MODELS_DIR/vae/wan_2.1_vae.safetensors"
    "$MODELS_DIR/vae/Wan2_1_VAE_fp32.safetensors"
    "$MODELS_DIR/text_encoders/umt5_xxl_fp16.safetensors"
    "$MODELS_DIR/text_encoders/umt5-xxl-enc-bf16.safetensors"
    "$MODELS_DIR/clip_vision/clip_vision_h.safetensors"
  )
  for file in "${files[@]}"; do
    if [ -f "$file" ]; then
      log_success "✓ $(basename "$file")"
    else
      log_warn "✗ $(basename "$file") não encontrado (pode ser opcional)"
    fi
  done

  echo ""
  log_info "Modelos disponíveis:"
  for dir in diffusion_models loras vae vae_approx text_encoders clip_vision controlnet upscale_models checkpoints; do
    if [ -d "$MODELS_DIR/$dir" ]; then
      local count
      count=$(find "$MODELS_DIR/$dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" \) 2>/dev/null | wc -l)
      if [ "$count" -gt 0 ]; then log_success "  $dir: $count arquivo(s)"; fi
    fi
  done
  return 0
}

if check_critical_files; then
  log_success "Instalação concluída com sucesso!"
  echo ""
  log_info "========================================="
  log_info "Iniciando ComfyUI..."
  log_info "URL: http://localhost:$COMFY_PORT"
  log_info "========================================="

  SAGE_FLAG=""
  python3 - <<'PY' >/dev/null 2>&1 && SAGE_FLAG="--use-sage-attention" || true
import sageattention
PY

  cd "$COMFY_DIR"
  exec comfy launch -- ${SAGE_FLAG:+$SAGE_FLAG} --listen "$COMFY_HOST" --port "$COMFY_PORT"
else
  log_warn "Alguns arquivos não foram encontrados; iniciando ComfyUI mesmo assim..."
  cd "$COMFY_DIR"
  exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
fi