#!/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:-4fcb2834969399006a736ee402b061e5}"
HF_TOKEN="${HF_TOKEN:-}"  # Token opcional do HuggingFace para login

# 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

# -----------------------------
# Lista de downloads
# -----------------------------
# Formato: "URL|TIPO|NOME_ARQUIVO_OPCIONAL"
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"
#sdasd
    "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 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"
    # ControlNet e outros modelos
    "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/taew2_1.safetensors|vae_approx|taew2_1.safetensors"
    "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
)

# Custom nodes para instalar
readonly CUSTOM_NODES=(
    "https://github.com/kijai/ComfyUI-Florence2"
    "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/yuvraj108c/ComfyUI-Dwpose-Tensorrt"
    "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
    "https://github.com/grmchn/ComfyUI-ProportionChanger"
)

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

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

# Download HuggingFace usando hf com hf_transfer
download_hf() {
    local repo="$1"
    local file_path="$2"
    local target_dir="$3"
    local filename="$4"
    
    # Se filename não foi fornecido, extrair do caminho
    if [ -z "$filename" ]; then
        filename=$(basename "$file_path")
    fi
    
    local target_file="$target_dir/$filename"
    
    # Se arquivo já existe e tem tamanho razoável, pular
    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"
    log_info "  Repo: $repo"
    log_info "  Path: $file_path"
    
    # Criar diretório temporário para download
    local temp_dir=$(mktemp -d)
    
    # Tentar primeiro com 'hf download' (novo comando)
    local download_success=false
    
    if command_exists hf; then
        if HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo" "$file_path" \
            --local-dir "$temp_dir" \
            --local-dir-use-symlinks False 2>/dev/null; then
            download_success=true
        fi
    fi
    
    # Se 'hf' não funcionou, tentar com huggingface-cli
    if [ "$download_success" = false ] && command_exists huggingface-cli; then
        log_warn "Tentando com huggingface-cli..."
        if HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" "$file_path" \
            --local-dir "$temp_dir" \
            --local-dir-use-symlinks False 2>/dev/null; then
            download_success=true
        fi
    fi
    
    if [ "$download_success" = true ]; then
        # Procurar o arquivo baixado (pode estar em subdiretórios)
        local downloaded_file=$(find "$temp_dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" \) 2>/dev/null | head -1)
        
        if [ -n "$downloaded_file" ]; then
            # Mover para o destino correto com o nome correto
            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"
    
    # Adicionar token Civitai se necessário
    url=$(add_civitai_token "$url")
    
    # Se arquivo já existe e tem tamanho razoável, pular
    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)}"
    
    # Tentar com aria2c primeiro (mais rápido)
    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
        else
            aria2c $aria_opts --dir="$target_dir" "$url" && return 0
        fi
    fi
    
    # Fallback para wget
    if command_exists wget; then
        if [ -n "$filename" ]; then
            wget -q --show-progress -c -O "$target_dir/$filename" "$url" && return 0
        else
            wget -q --show-progress -c --content-disposition -P "$target_dir" "$url" && return 0
        fi
    fi
    
    # Último recurso: curl
    if command_exists curl; then
        if [ -n "$filename" ]; then
            curl -L -# -C - -o "$target_dir/$filename" "$url" && return 0
        else
            local headers=$(curl -sI -L "$url")
            local detected_name=$(echo "$headers" | grep -i "content-disposition" | sed -n 's/.*filename="\([^"]*\)".*/\1/p' | tr -d '\r')
            if [ -z "$detected_name" ]; then
                detected_name="downloaded_$(date +%s).safetensors"
            fi
            curl -L -# -C - -o "$target_dir/$detected_name" "$url" && return 0
        fi
    fi
    
    log_error "Falha ao baixar: $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"
        
        # Limpar espaços
        url=$(echo "$url" | xargs)
        type=$(echo "$type" | xargs)
        filename=$(echo "$filename" | xargs)
        
        # Determinar diretório de destino
        local target_dir="$MODELS_DIR"
        case "$type" in
            checkpoints) target_dir="$MODELS_DIR/checkpoints" ;;
            diffusion_models) target_dir="$MODELS_DIR/diffusion_models" ;;
            loras) target_dir="$MODELS_DIR/loras" ;;
            vae) target_dir="$MODELS_DIR/vae" ;;
            vae_approx) target_dir="$MODELS_DIR/vae_approx" ;;
            text_encoders) target_dir="$MODELS_DIR/text_encoders" ;;
            clip_vision) target_dir="$MODELS_DIR/clip_vision" ;;
            controlnet) target_dir="$MODELS_DIR/controlnet" ;;
            upscale_models) target_dir="$MODELS_DIR/upscale_models" ;;
            *) target_dir="$MODELS_DIR/$type" ;;
        esac
        
        # Criar diretório se não existir
        mkdir -p "$target_dir"
        
        # Processar URL
        if [[ "$url" == hf://* ]]; then
            # HuggingFace download
            url="${url#hf://}"  # Remove prefixo hf://
            
            # Separar repo e caminho do arquivo
            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 normal (HTTP/HTTPS)
            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=$(basename "$dest")
    
    # Timeout de 60 segundos para operações git
    if [ -d "$dest/.git" ]; then
        log_info "Atualizando $node_name..."
        timeout 60 git -C "$dest" pull --ff-only 2>/dev/null || {
            log_warn "Timeout ou erro ao atualizar $node_name"
            return 0  # Não falhar, apenas avisar
        }
    else
        log_info "Clonando $node_name..."
        timeout 60 git clone --recursive --depth 1 "$url" "$dest" 2>/dev/null || {
            log_warn "Falha ao clonar $node_name"
            return 0  # Não falhar, apenas avisar
        }
    fi
    
    # Instalar requirements se existir (com timeout)
    if [ -f "$dest/requirements.txt" ]; then
        timeout 120 python3 -m pip install --no-warn-script-location -q -r "$dest/requirements.txt" 2>/dev/null || {
            log_warn "Falha ao instalar requirements para $node_name"
        }
    fi
    
    return 0
}

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

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

# Instalar aria2c do sistema se não existir
if ! command_exists aria2c; then
    log_info "Instalando aria2c do sistema..."
    apt-get update -qq && apt-get install -y -qq aria2 2>/dev/null
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/8] Preparando ferramentas Python..."
python3 -m pip install --upgrade pip wheel setuptools -q

# Instalar huggingface_hub com CLI e hf_transfer para velocidade máxima
python3 -m pip install --upgrade "huggingface_hub[cli,hf_transfer]>=0.26.0" comfy-cli -q

# Configurar token HF se disponível (opcional)
HF_TOKEN="${HF_TOKEN:-}"
if [ -n "$HF_TOKEN" ]; then
    if huggingface-cli login --token "$HF_TOKEN" 2>/dev/null; then
        log_success "Login HF configurado"
    else
        log_warn "Falha ao configurar login HF, continuando sem autenticação"
    fi
fi

# Verificar comandos disponíveis
if command_exists hf; then
    log_success "Comando 'hf' disponível (recomendado)"
elif command_exists huggingface-cli; then
    log_warn "Usando 'huggingface-cli' (deprecated, mas funcional)"
else
    log_error "Nenhum comando HF disponível!"
fi

log_success "Ferramentas instaladas"

# [3] Instalar PyTorch com CUDA 12.8
log_info "[3/8] 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/8] 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
    
    # Verificar instalação
    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] Instalar SageAttention (wheel pré-compilada)
log_info "[5/8] Instalando SageAttention..."
if command_exists nvcc || python3 -c "import torch; print(torch.cuda.is_available())" 2>/dev/null | grep -q "True"; then
    log_info "Instalando SageAttention pré-compilada..."
    python3 -m pip install "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" -q
    
    # Verificar instalação
    if python3 -c "import sageattention" 2>/dev/null; then
        log_success "SageAttention instalado com sucesso"
    else
        log_warn "SageAttention instalado mas não pôde ser importado"
    fi
else
    log_warn "CUDA não detectado, pulando SageAttention"
fi

# [6] Baixar modelos
log_info "[6/8] Baixando modelos..."
mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,vae_approx,text_encoders,clip_vision,controlnet,upscale_models,checkpoints}

# Verificar ferramentas de download
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

# Processar downloads
if process_downloads; then
    log_success "Todos os modelos baixados com sucesso"
else
    log_warn "Alguns downloads falharam, mas continuando..."
fi

# [7] Instalar custom nodes
log_info "[7/8] Instalando custom nodes..."
CN_DIR="$COMFY_DIR/custom_nodes"
mkdir -p "$CN_DIR"

# Instalar nodes sequencialmente para evitar travamentos
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"

# [8] Verificação final
echo ""
log_info "[8/8] Verificando instalação..."

# Verificar arquivos principais (ajustado para os modelos que você está baixando)
check_critical_files() {
    local files=(
        "$COMFY_DIR/main.py"
        "$MODELS_DIR/diffusion_models/Wan2_1-I2V-14B-720P_fp8_e4m3fn.safetensors"
        "$MODELS_DIR/vae/wan_2.1_vae.safetensors"
        "$MODELS_DIR/text_encoders/umt5_xxl_fp16.safetensors"
        "$MODELS_DIR/clip_vision/clip_vision_h.safetensors"
    )
    
    local missing=0
    for file in "${files[@]}"; do
        if [ -f "$file" ]; then
            log_success "✓ $(basename "$file")"
        else
            log_warn "✗ $(basename "$file") não encontrado (opcional)"
        fi
    done
    
    # Listar modelos baixados
    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=$(find "$MODELS_DIR/$dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" \) 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 "========================================="
    
    # Tentar usar SageAttention se disponível
    SAGE_FLAG=""
    if python3 -c "import sageattention" 2>/dev/null; then
        SAGE_FLAG="--use-sage-attention"
        log_info "SageAttention habilitado"
    fi
    
    cd "$COMFY_DIR"
    exec comfy launch -- $SAGE_FLAG --listen "$COMFY_HOST" --port "$COMFY_PORT"
else
    log_error "Alguns arquivos não foram encontrados, mas a instalação pode funcionar."
    log_info "Iniciando ComfyUI mesmo assim..."
    
    cd "$COMFY_DIR"
    exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
fi