#!/bin/bash

# ============================================================================
# Instalador Otimizado para ComfyUI com Wan 2.2
# Versão 3.1 - Com instalação do PyTorch CUDA 12.8
# ============================================================================

set -e
set -o pipefail

# ============================================================================
# CONFIGURAÇÕES GLOBAIS
# ============================================================================

# Diretórios base
readonly INSTALL_BASE_DIR="$(pwd)"
readonly WORKSPACE_DIR="$INSTALL_BASE_DIR/ComfyUI"
readonly WORKFLOW_FILENAME="workflow.json"
readonly WORKFLOW_FILE="$INSTALL_BASE_DIR/$WORKFLOW_FILENAME"
readonly LOG_FILE="$INSTALL_BASE_DIR/install.log"

# URL do workflow
readonly WORKFLOW_URL="https://huggingface.co/adbrasi/akarris_trainer/resolve/main/Wan_GenerationBase.json"

# ============================================================================
# FUNÇÕES DE UTILIDADE
# ============================================================================

log_info() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] INFO: $1"
    echo "$msg" | tee -a "$LOG_FILE" >&2
}

log_warning() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] WARNING: $1"
    echo "$msg" | tee -a "$LOG_FILE" >&2
}

log_error() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1"
    echo "$msg" | tee -a "$LOG_FILE" >&2
}

command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# ============================================================================
# VERIFICAÇÃO DE REQUISITOS
# ============================================================================

check_requirements() {
    log_info "Verificando requisitos do sistema..."
    
    local missing_deps=()
    
    if ! command_exists python3; then
        missing_deps+=("python3")
    fi
    
    if ! command_exists git; then
        missing_deps+=("git")
    fi
    
    if ! command_exists gcc; then
        missing_deps+=("gcc/build-essential")
    fi
    
    if ! command_exists nvcc; then
        log_warning "CUDA toolkit (nvcc) não encontrado. A compilação do SageAttention pode falhar."
    fi
    
    if [ ${#missing_deps[@]} -gt 0 ]; then
        log_error "Dependências faltando: ${missing_deps[*]}"
        log_error "Por favor, instale as dependências e execute o script novamente."
        exit 1
    fi
    
    # Verifica versão do Python
    local python_version=$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
    log_info "Python version: $python_version"
    
    if ! python3 -c 'import sys; exit(0 if sys.version_info >= (3,8) else 1)'; then
        log_error "Python 3.8 ou superior é necessário"
        exit 1
    fi
    
    log_info "Todos os requisitos básicos estão presentes."
}

# ============================================================================
# LIMPEZA DE INSTALAÇÃO EXISTENTE
# ============================================================================

cleanup_existing_installation() {
    log_info "Verificando instalações existentes em: $WORKSPACE_DIR"
    
    if [ -d "$WORKSPACE_DIR" ]; then
        log_warning "Diretório existente encontrado: $WORKSPACE_DIR"
        read -p "Deseja remover a instalação existente? (s/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Ss]$ ]]; then
            log_info "Removendo a instalação existente..."
            rm -rf "$WORKSPACE_DIR"
            log_info "Diretório removido."
        else
            log_error "Instalação cancelada pelo usuário."
            exit 1
        fi
    else
        log_info "Nenhuma instalação anterior encontrada."
    fi
}

# ============================================================================
# INSTALAÇÃO DO COMFY-CLI
# ============================================================================

install_comfy_cli() {
    log_info "Instalando/atualizando comfy-cli..."
    
    python3 -m pip install --upgrade pip
    python3 -m pip install --upgrade comfy-cli || {
        log_error "Falha ao instalar comfy-cli"
        return 1
    }
    
    # Adiciona ~/.local/bin ao PATH se necessário
    if [ -d "$HOME/.local/bin" ] && [[ ":$PATH:" != *":$HOME/.local/bin:"* ]]; then
        export PATH="$HOME/.local/bin:$PATH"
        log_info "Adicionado $HOME/.local/bin ao PATH."
    fi
    
    if ! command_exists comfy; then
        log_error "Comando 'comfy' não encontrado no PATH."
        return 1
    fi
    
    log_info "comfy-cli instalado: $(comfy --version 2>/dev/null || echo 'versão desconhecida')"
    return 0
}

# ============================================================================
# INSTALAÇÃO DO COMFYUI
# ============================================================================

install_comfyui() {
    log_info "Instalando ComfyUI em: $WORKSPACE_DIR"
    
    # Instala ComfyUI com suporte NVIDIA
    if ! comfy --workspace "$WORKSPACE_DIR" --skip-prompt install --nvidia; then
        log_error "Falha ao instalar ComfyUI"
        return 1
    fi
    
    log_info "ComfyUI instalado com sucesso."
    return 0
}

# ============================================================================
# INSTALAÇÃO DO HUGGINGFACE HUB
# ============================================================================

install_hf_tools() {
    log_info "Instalando HuggingFace Hub e hf_transfer..."
    
    python3 -m pip install -U "huggingface_hub[cli,hf_transfer]" || {
        log_error "Falha ao instalar ferramentas HuggingFace"
        return 1
    }
    
    # Habilita o hf_transfer para downloads mais rápidos
    export HF_HUB_ENABLE_HF_TRANSFER=1
    
    log_info "HuggingFace tools instalados com sucesso."
    return 0
}

# ============================================================================
# INSTALAÇÃO DO SAGEATTENTION 2.2.0
# ============================================================================

install_sageattention() {
    log_info "Instalando SageAttention 2.2.0 com 2++ (compilação otimizada)..."
    
    # =========================================================================
    # INSTALAÇÃO DO PYTORCH COM CUDA 12.8
    # =========================================================================
    log_info "Instalando PyTorch 2.7.0 com CUDA 12.8..."
    
    python3 -m pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128 || {
        log_error "Falha ao instalar PyTorch com CUDA 12.8"
        return 1
    }
    
    # Verifica instalação do PyTorch
    log_info "Verificando instalação do PyTorch..."
    python3 -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}'); print(f'CUDA version: {torch.version.cuda if torch.cuda.is_available() else \"N/A\"}')" 2>&1 | tee -a "$LOG_FILE"
    
    local sage_dir="$INSTALL_BASE_DIR/SageAttention"
    
    # Remove instalação anterior se existir
    if [ -d "$sage_dir" ]; then
        log_info "Removendo instalação anterior do SageAttention..."
        rm -rf "$sage_dir"
    fi
    
    # Instala dependências necessárias
    log_info "Instalando dependências do SageAttention..."
    python3 -m pip install 'triton>=3.0.0,<3.1.0' packaging ninja || {
        log_error "Falha ao instalar dependências do SageAttention"
        return 1
    }
    
    # Clona o repositório
    log_info "Clonando repositório SageAttention..."
    if ! git clone https://github.com/thu-ml/SageAttention.git "$sage_dir"; then
        log_error "Falha ao clonar repositório SageAttention"
        return 1
    fi
    
    cd "$sage_dir"
    
    # Configura variáveis de ambiente para compilação otimizada
    log_info "Configurando ambiente de compilação otimizada..."
    export EXT_PARALLEL=4
    export NVCC_APPEND_FLAGS="--threads 8"
    export MAX_JOBS=32
    export TORCH_CUDA_ARCH_LIST="8.9"  # Ada architecture
    
    # Detecta automaticamente outras arquiteturas CUDA se disponível
    if command_exists nvidia-smi; then
        local cuda_capability=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -n1 | tr -d '.')
        if [ -n "$cuda_capability" ]; then
            log_info "Detectada capacidade CUDA: $cuda_capability"
            if [ "$cuda_capability" != "89" ]; then
                export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;${cuda_capability:0:1}.${cuda_capability:1}"
                log_info "TORCH_CUDA_ARCH_LIST: $TORCH_CUDA_ARCH_LIST"
            fi
        fi
    fi
    
    # Compila e instala
    log_info "Compilando SageAttention (5-10 minutos)..."
    if ! python3 setup.py install >> "$LOG_FILE" 2>&1; then
        log_error "Falha ao compilar SageAttention. Verifique $LOG_FILE"
        cd "$INSTALL_BASE_DIR"
        return 1
    fi
    
    cd "$INSTALL_BASE_DIR"
    log_info "SageAttention 2.2.0 instalado com sucesso!"
    return 0
}

# ============================================================================
# DOWNLOAD DO WORKFLOW
# ============================================================================

download_workflow() {
    log_info "Baixando o workflow..."
    
    # Usa wget com user-agent para evitar bloqueios
    if command_exists wget; then
        wget -U "Mozilla/5.0" -O "$WORKFLOW_FILE" "$WORKFLOW_URL" 2>&1 | tee -a "$LOG_FILE"
    elif command_exists curl; then
        curl -L -H "User-Agent: Mozilla/5.0" -o "$WORKFLOW_FILE" "$WORKFLOW_URL" 2>&1 | tee -a "$LOG_FILE"
    else
        log_error "wget ou curl não encontrado"
        return 1
    fi
    
    if [ -f "$WORKFLOW_FILE" ] && [ -s "$WORKFLOW_FILE" ]; then
        log_info "Workflow baixado com sucesso."
        return 0
    else
        log_error "Falha ao baixar o workflow."
        return 1
    fi
}

# ============================================================================
# DOWNLOAD DE MODELOS WAN 2.2 (CORRIGIDO)
# ============================================================================

download_models() {
    log_info "Iniciando download dos modelos Wan 2.2..."
    
    # Criação de diretórios
    local dirs=(
        "$WORKSPACE_DIR/models/diffusion_models"
        "$WORKSPACE_DIR/models/loras"
        "$WORKSPACE_DIR/models/text_encoders"
        "$WORKSPACE_DIR/models/vae"
        "$WORKSPACE_DIR/models/clip_vision"
    )
    
    for dir in "${dirs[@]}"; do
        mkdir -p "$dir"
        log_info "Diretório criado/verificado: $dir"
    done
    
    # Configuração para usar hf
    export HF_HUB_ENABLE_HF_TRANSFER=1
    
    # Lista de modelos organizados por tipo
    log_info "Baixando modelos de difusão..."
    
    # Diffusion Models
    hf download "Comfy-Org/Wan_2.2_ComfyUI_Repackaged" \
        "split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors" \
        "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors" \
        --local-dir "$WORKSPACE_DIR/models" \
        --cache-dir "$WORKSPACE_DIR/models/.cache" 2>&1 | tee -a "$LOG_FILE"
    
    log_info "Baixando LoRAs..."
    
    # LoRAs
    hf download "Comfy-Org/Wan_2.2_ComfyUI_Repackaged" \
        "split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors" \
        "split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors" \
        "split_files/loras/wan2.2_t2v_lightx2v_4steps_lora_v1.1_high_noise.safetensors" \
        "split_files/loras/wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors" \
        --local-dir "$WORKSPACE_DIR/models" \
        --cache-dir "$WORKSPACE_DIR/models/.cache" 2>&1 | tee -a "$LOG_FILE"
    
    log_info "Baixando Text Encoders..."
    
    # Text Encoders
    hf download "Comfy-Org/Wan_2.2_ComfyUI_Repackaged" \
        "split_files/text_encoders/umt5_xxl_fp16.safetensors" \
        "split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" \
        --local-dir "$WORKSPACE_DIR/models" \
        --cache-dir "$WORKSPACE_DIR/models/.cache" 2>&1 | tee -a "$LOG_FILE"
    
    log_info "Baixando VAE..."
    
    # VAE
    hf download "Comfy-Org/Wan_2.1_ComfyUI_repackaged" \
        "split_files/vae/wan_2.1_vae.safetensors" \
        --local-dir "$WORKSPACE_DIR/models" \
        --cache-dir "$WORKSPACE_DIR/models/.cache" 2>&1 | tee -a "$LOG_FILE"
    
    log_info "Baixando CLIP Vision..."
    
    # CLIP Vision
    hf download "Comfy-Org/Wan_2.1_ComfyUI_repackaged" \
        "split_files/clip_vision/clip_vision_h.safetensors" \
        --local-dir "$WORKSPACE_DIR/models" \
        --cache-dir "$WORKSPACE_DIR/models/.cache" 2>&1 | tee -a "$LOG_FILE"
    
    # Move arquivos dos subdiretórios para os locais corretos
    log_info "Organizando arquivos de modelos..."
    
    # Move diffusion models
    if [ -d "$WORKSPACE_DIR/models/split_files/diffusion_models" ]; then
        mv "$WORKSPACE_DIR/models/split_files/diffusion_models/"*.safetensors "$WORKSPACE_DIR/models/diffusion_models/" 2>/dev/null || true
    fi
    
    # Move loras
    if [ -d "$WORKSPACE_DIR/models/split_files/loras" ]; then
        mv "$WORKSPACE_DIR/models/split_files/loras/"*.safetensors "$WORKSPACE_DIR/models/loras/" 2>/dev/null || true
    fi
    
    # Move text encoders
    if [ -d "$WORKSPACE_DIR/models/split_files/text_encoders" ]; then
        mv "$WORKSPACE_DIR/models/split_files/text_encoders/"*.safetensors "$WORKSPACE_DIR/models/text_encoders/" 2>/dev/null || true
    fi
    
    # Move vae
    if [ -d "$WORKSPACE_DIR/models/split_files/vae" ]; then
        mv "$WORKSPACE_DIR/models/split_files/vae/"*.safetensors "$WORKSPACE_DIR/models/vae/" 2>/dev/null || true
    fi
    
    # Move clip vision
    if [ -d "$WORKSPACE_DIR/models/split_files/clip_vision" ]; then
        mv "$WORKSPACE_DIR/models/split_files/clip_vision/"*.safetensors "$WORKSPACE_DIR/models/clip_vision/" 2>/dev/null || true
    fi
    
    # Remove diretório split_files vazio
    rm -rf "$WORKSPACE_DIR/models/split_files" 2>/dev/null || true
    
    # Verifica arquivos baixados
    log_info "Verificando modelos baixados:"
    for dir in "${dirs[@]}"; do
        if [ -d "$dir" ]; then
            local count=$(ls -1 "$dir"/*.safetensors 2>/dev/null | wc -l)
            log_info "  $(basename $(dirname "$dir"))/$(basename "$dir"): $count arquivo(s)"
            
            # Lista os arquivos
            for file in "$dir"/*.safetensors; do
                if [ -f "$file" ]; then
                    local size=$(du -h "$file" | cut -f1)
                    log_info "    - $(basename "$file") ($size)"
                fi
            done
        fi
    done
    
    return 0
}

# ============================================================================
# INSTALAÇÃO MANUAL DE CUSTOM NODES
# ============================================================================

install_custom_nodes_manually() {
    log_info "Instalando custom nodes manualmente..."
    
    local custom_nodes_dir="$WORKSPACE_DIR/custom_nodes"
    
    # Lista completa de custom nodes (sem duplicatas)
    local custom_nodes=(
        "https://github.com/kijai/ComfyUI-WanVideoWrapper"
        "https://github.com/kijai/ComfyUI-MelBandRoFormer"
        "https://github.com/kijai/ComfyUI-KJNodes"
        "https://github.com/kijai/ComfyUI-Florence2"
        "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"
        "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"
        "https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"
        "https://github.com/rdancer/ComfyUI_Florence2SAM2"
        "https://github.com/1038lab/ComfyUI-RMBG"
        "https://github.com/fpgaminer/joycaption_comfyui"
        "https://github.com/MieMieeeee/ComfyUI-CaptionThis"
        "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
        "https://github.com/LAOGOU-666/Comfyui-Memory_Cleanup"
        "https://github.com/chrisgoringe/cg-use-everywhere"
        "https://github.com/pythongosssss/ComfyUI-Custom-Scripts"
        "https://github.com/chibiace/ComfyUI-Chibi-Nodes"
        "https://github.com/ptmaster/comfyui-audio-speed"
        "https://github.com/filliptm/ComfyUI_Fill-Nodes"
        "https://github.com/justUmen/Bjornulf_custom_nodes"
        "https://github.com/raindrop313/ComfyUI-WanVideoStartEndFrames"
        "https://github.com/ShmuelRonen/ComfyUI_wav2lip"
        "https://github.com/crystian/ComfyUI-Crystools"
        "https://github.com/bash-j/mikey_nodes"
    )
    
    cd "$custom_nodes_dir"
    
    for node_url in "${custom_nodes[@]}"; do
        local node_name=$(basename "$node_url" .git)
        log_info "Instalando: $node_name"
        
        # Remove instalação anterior se existir
        if [ -d "$node_name" ]; then
            log_info "  Removendo versão anterior..."
            rm -rf "$node_name"
        fi
        
        # Clona o repositório
        if git clone --depth 1 "$node_url" "$node_name" 2>&1 | tee -a "$LOG_FILE"; then
            log_info "  Clone concluído: $node_name"
            
            # Instala requirements se existir
            if [ -f "$node_name/requirements.txt" ]; then
                log_info "  Instalando dependências de $node_name..."
                python3 -m pip install -r "$node_name/requirements.txt" 2>&1 | tee -a "$LOG_FILE" || {
                    log_warning "  Algumas dependências de $node_name falharam"
                }
            fi
            
            # Executa install.py se existir
            if [ -f "$node_name/install.py" ]; then
                log_info "  Executando install.py de $node_name..."
                cd "$node_name"
                python3 install.py 2>&1 | tee -a "$LOG_FILE" || {
                    log_warning "  install.py de $node_name retornou erro"
                }
                cd ..
            fi
        else
            log_warning "  Falha ao clonar: $node_name"
        fi
    done
    
    cd "$INSTALL_BASE_DIR"
    log_info "Instalação de custom nodes concluída!"
}

# ============================================================================
# VERIFICAÇÃO FINAL DA INSTALAÇÃO
# ============================================================================

verify_installation() {
    log_info "Verificando a instalação..."
    
    local issues=()
    
    # Verifica ComfyUI
    if [ ! -d "$WORKSPACE_DIR" ]; then
        issues+=("Diretório ComfyUI não encontrado")
    fi
    
    # Verifica workflow
    if [ ! -f "$WORKFLOW_FILE" ]; then
        issues+=("Arquivo de workflow não encontrado")
    fi
    
    # Verifica modelos em cada diretório
    log_info "Verificando modelos instalados:"
    
    local expected_models=(
        "diffusion_models:2"
        "loras:4"
        "text_encoders:2"
        "vae:1"
        "clip_vision:1"
    )
    
    for expectation in "${expected_models[@]}"; do
        IFS=':' read -r dir_name expected_count <<< "$expectation"
        local dir="$WORKSPACE_DIR/models/$dir_name"
        
        if [ -d "$dir" ]; then
            local count=$(ls -1 "$dir"/*.safetensors 2>/dev/null | wc -l)
            log_info "  $dir_name: $count arquivo(s) (esperado: $expected_count)"
            
            if [ $count -lt $expected_count ]; then
                issues+=("Faltam modelos em $dir_name ($count/$expected_count)")
            fi
        else
            issues+=("Diretório não existe: $dir_name")
        fi
    done
    
    # Verifica PyTorch com CUDA
    log_info "Verificando PyTorch e CUDA:"
    if python3 -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
        log_info "✓ PyTorch com CUDA instalado corretamente"
    else
        log_warning "PyTorch instalado mas CUDA não está disponível"
    fi
    
    # Verifica SageAttention
    if ! python3 -c "import sageattention" 2>/dev/null; then
        log_warning "SageAttention não está instalado (opcional)"
    else
        log_info "✓ SageAttention instalado corretamente"
    fi
    
    if [ ${#issues[@]} -gt 0 ]; then
        log_error "Problemas encontrados na instalação:"
        for issue in "${issues[@]}"; do
            log_error "  ✗ $issue"
        done
        return 1
    fi
    
    log_info "✓ Instalação verificada com sucesso!"
    return 0
}

# ============================================================================
# INICIAR O COMFYUI
# ============================================================================

launch_comfyui() {
    log_info "================================================="
    log_info "Instalação concluída com sucesso!"
    log_info "================================================="
    log_info ""
    log_info "Iniciando ComfyUI..."
    log_info "  Porta: 8818"
    log_info "  Workspace: $WORKSPACE_DIR"
    log_info "  Workflow: $WORKFLOW_FILE"
    log_info ""
    log_info "Para usar o workflow:"
    log_info "  1. Acesse http://127.0.0.1:8818"
    log_info "  2. Arraste o arquivo '$WORKFLOW_FILENAME' na interface"
    log_info ""
    log_info "Pressione Ctrl+C para parar o servidor"
    log_info "================================================="
    
    # Verifica se o SageAttention está instalado
    local sage_args=""
    if python3 -c "import sageattention" 2>/dev/null; then
        sage_args="--use-sage-attention"
        log_info "SageAttention habilitado!"
    fi
    
    # Inicia o ComfyUI
    exec comfy --workspace "$WORKSPACE_DIR" launch -- \
        --fast \
        --listen 0.0.0.0 \
        --port 8818 \
        $sage_args \
        --preview-method latent2rgb
}

# ============================================================================
# FUNÇÃO PRINCIPAL
# ============================================================================

main() {
    # Inicializa log
    echo "=== Instalação ComfyUI iniciada em $(date) ===" > "$LOG_FILE"
    
    log_info "================================================="
    log_info "Instalador ComfyUI com Wan 2.2 - v3.1"
    log_info "Com PyTorch 2.7.0 + CUDA 12.8"
    log_info "================================================="
    
    # Etapa 1: Verificar requisitos
    check_requirements || exit 1
    
    # Etapa 2: Limpar instalação existente
    cleanup_existing_installation
    
    # Etapa 3: Instalar comfy-cli
    install_comfy_cli || exit 1
    
    # Etapa 4: Instalar ferramentas HuggingFace
    install_hf_tools || exit 1
    
    # Etapa 5: Instalar ComfyUI
    install_comfyui || exit 1
    
    # Etapa 6: Baixar workflow
    download_workflow || log_warning "Workflow não foi baixado"
    
    # Etapa 7: Instalar custom nodes manualmente
    install_custom_nodes_manually
    
    # Etapa 8: PROCESSOS PARALELOS - SageAttention e Downloads
    log_info "================================================="
    log_info "Iniciando processos paralelos:"
    log_info "  1. Instalação do PyTorch + SageAttention (10-15 min)"
    log_info "  2. Download dos modelos Wan 2.2"
    log_info "================================================="
    
    # Inicia instalação do SageAttention em background
    install_sageattention &
    local sage_pid=$!
    
    # Inicia downloads dos modelos em background
    download_models &
    local models_pid=$!
    
    # Monitor de progresso
    log_info "Aguardando conclusão dos processos paralelos..."
    log_info "Monitore o progresso em: $LOG_FILE"
    
    local both_done=false
    local sage_done=false
    local models_done=false
    
    while [ "$both_done" = false ]; do
        # Verifica SageAttention
        if [ "$sage_done" = false ]; then
            if ! ps -p $sage_pid > /dev/null 2>&1; then
                wait $sage_pid
                local sage_exit=$?
                sage_done=true
                if [ $sage_exit -eq 0 ]; then
                    log_info "✓ PyTorch + SageAttention instalado com sucesso!"
                else
                    log_warning "⚠ PyTorch/SageAttention pode não ter sido instalado corretamente"
                fi
            fi
        fi
        
        # Verifica Downloads
        if [ "$models_done" = false ]; then
            if ! ps -p $models_pid > /dev/null 2>&1; then
                wait $models_pid
                local models_exit=$?
                models_done=true
                if [ $models_exit -eq 0 ]; then
                    log_info "✓ Downloads dos modelos concluídos!"
                else
                    log_warning "⚠ Alguns modelos podem não ter sido baixados"
                fi
            fi
        fi
        
        # Verifica se ambos terminaram
        if [ "$sage_done" = true ] && [ "$models_done" = true ]; then
            both_done=true
        else
            sleep 10
            if [ "$sage_done" = false ] && ps -p $sage_pid > /dev/null 2>&1; then
                echo -n "."
            fi
            if [ "$models_done" = false ] && ps -p $models_pid > /dev/null 2>&1; then
                echo -n "."
            fi
        fi
    done
    
    echo ""
    log_info "Todos os processos paralelos foram concluídos!"
    
    # Etapa 9: Verificar instalação
    verify_installation || {
        log_error "A instalação teve problemas. Verifique o log em: $LOG_FILE"
        log_error "Você pode tentar executar o script novamente."
        exit 1
    }
    
    # Etapa 10: Iniciar ComfyUI
    launch_comfyui
}

# ============================================================================
# TRATAMENTO DE SINAIS E EXECUÇÃO
# ============================================================================

cleanup_on_exit() {
    log_info "Interrompendo instalação..."
    # Mata processos filhos
    jobs -p | xargs -r kill 2>/dev/null
    exit 1
}

# Captura sinais de interrupção
trap cleanup_on_exit INT TERM

# Executa o script principal
main "$@"