#!/bin/bash

# ============================================================================
# Instalador Dinâmico para ComfyUI
# Instala o ComfyUI e suas dependências no diretório atual.
# ============================================================================

set -e
set -o pipefail

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

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

# --- Caminhos Dinâmicos ---
# Define o diretório de instalação como o diretório atual onde o script é executado
readonly INSTALL_BASE_DIR="$(pwd)"
# O ComfyUI será instalado em uma subpasta "ComfyUI"
readonly WORKSPACE_DIR="$INSTALL_BASE_DIR/ComfyUI"
# O arquivo de workflow será salvo na raiz do diretório de instalação
readonly WORKFLOW_FILENAME="workflow.json"
readonly WORKFLOW_FILE="$INSTALL_BASE_DIR/$WORKFLOW_FILENAME"

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

log_info() {
    echo "INFO: $1" >&2
}

log_warning() {
    echo "WARNING: $1" >&2
}

log_error() {
    echo "ERROR: $1" >&2
}

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

# ============================================================================
# FUNÇÃO DE LIMPEZA
# ============================================================================

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"
        log_info "Removendo a instalação existente para evitar conflitos..."
        rm -rf "$WORKSPACE_DIR"
        log_info "Diretório removido."
    else
        log_info "Nenhuma instalação anterior encontrada no local. Continuando..."
    fi
}

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

install_comfy_cli() {
    log_info "Instalando/atualizando comfy-cli..."
    
    python3 -m pip install --upgrade comfy-cli || {
        log_error "Falha ao instalar comfy-cli"
        exit 1
    }
    
    # Verifica se o 'comfy' está no PATH
    if ! command_exists comfy; then
        if [ -d "$HOME/.local/bin" ] && [[ ":$PATH:" != *":$HOME/.local/bin:"* ]]; then
            export PATH="$HOME/.local/bin:$PATH"
            log_info "Adicionado $HOME/.local/bin ao PATH para esta sessão."
        fi
        
        if ! command_exists comfy; then
            log_error "Comando 'comfy' não encontrado no PATH após a instalação."
            log_error "Verifique se '$HOME/.local/bin' está em seu PATH."
            exit 1
        fi
    fi
    
    log_info "comfy-cli version: $(comfy --version 2>/dev/null || echo 'desconhecida')"
}

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

install_comfyui() {
    log_info "Instalando ComfyUI em: $WORKSPACE_DIR"
    
    # Instala ComfyUI com dependências rápidas
    comfy --workspace "$WORKSPACE_DIR" --skip-prompt install --nvidia || {
        log_error "Falha ao instalar ComfyUI"
        exit 1
    }
    
    log_info "ComfyUI instalado com sucesso."
}

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

install_hf_tools() {
    log_info "Instalando HuggingFace Hub CLI e hf_transfer..."
    
    python3 -m pip install -U "huggingface_hub[cli]" || {
        log_error "Falha ao instalar huggingface_hub[cli]"
        exit 1
    }
    
    python3 -m pip install hf_transfer || {
        log_error "Falha ao instalar hf_transfer"
        exit 1
    }
    
    log_info "HuggingFace tools instalados com sucesso."
}

# ============================================================================
# FUNÇÃO DE DOWNLOAD DE ARQUIVO
# ============================================================================

download_file() {
    local url="$1"
    local dest_dir="$2"
    local filename="${3:-}" # Opcional: nome do arquivo de saída

    log_info "Baixando: $url -> $dest_dir"

    # Garante que o diretório de destino exista
    mkdir -p "$dest_dir"

    local output_opts=()
    local curl_opts=()
    if [ -n "$filename" ]; then
        output_opts=("-O" "$filename")
        curl_opts=("-o" "$filename")
    else
        # Usa --content-disposition para obter o nome do arquivo do cabeçalho
        output_opts=("--content-disposition")
        curl_opts=("-O" "-J")
    fi

    if command_exists aria2c; then
        # aria2c lida com o nome do arquivo automaticamente com -o
        local aria_output_opt=()
        if [ -n "$filename" ]; then
            aria_output_opt=("-o" "$filename")
        fi
        aria2c -x 16 -s 16 -k 1M --dir="$dest_dir" "${aria_output_opt[@]}" "$url" || {
            log_warning "aria2c falhou, tentando com wget..."
            if command_exists wget; then
                wget -P "$dest_dir" "${output_opts[@]}" "$url" || {
                    log_warning "wget falhou, tentando com curl..."
                    (cd "$dest_dir" && curl -L "${curl_opts[@]}" "$url") || {
                        log_error "Falha ao baixar $url com todas as ferramentas."
                        return 1
                    }
                }
            else
                (cd "$dest_dir" && curl -L "${curl_opts[@]}" "$url") || {
                    log_error "Falha ao baixar $url com curl."
                    return 1
                }
            fi
        }
    elif command_exists wget; then
        wget -P "$dest_dir" "${output_opts[@]}" "$url" || {
            log_warning "wget falhou, tentando com curl..."
            (cd "$dest_dir" && curl -L "${curl_opts[@]}" "$url") || {
                log_error "Falha ao baixar $url com curl."
                return 1
            }
        }
    elif command_exists curl; then
        (cd "$dest_dir" && curl -L "${curl_opts[@]}" "$url") || {
            log_error "Falha ao baixar $url com curl."
            return 1
        }
    else
        log_error "Nenhuma ferramenta de download (aria2c, wget, curl) encontrada."
        return 1
    fi

    log_info "Download concluído: $url"
    return 0
}

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

download_workflow() {
    log_info "Baixando o workflow de: $WORKFLOW_URL"
    download_file "$WORKFLOW_URL" "$INSTALL_BASE_DIR" "$WORKFLOW_FILENAME" || {
        log_error "Falha ao baixar o arquivo do workflow."
        exit 1
    }
    
    if [ ! -f "$WORKFLOW_FILE" ]; then
        log_error "O arquivo do workflow não foi baixado com sucesso."
        exit 1
    fi
    
    log_info "Workflow baixado com sucesso."
}

# ============================================================================
# DEPENDÊNCIAS DO WORKFLOW E CUSTOM NODES
# ============================================================================

install_custom_nodes_and_deps() {
    log_info "Instalando custom nodes e dependências..."
    
    # Instala dependências do workflow primeiro
    if [ -f "$WORKFLOW_FILE" ]; then
        log_info "Instalando dependências do workflow..."
        comfy --workspace "$WORKSPACE_DIR" node install-deps --workflow="$WORKFLOW_FILE" || {
            log_warning "Algumas dependências do workflow podem ter falhado ao instalar."
        }
    else
        log_warning "Arquivo de workflow não encontrado para instalar dependências."
    fi

    # Lista de custom nodes para instalar
    local custom_nodes=(
        "https://github.com/LAOGOU-666/Comfyui-Memory_Cleanup"
        "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
        "https://github.com/kijai/ComfyUI-KJNodes"
        "https://github.com/kijai/ComfyUI-WanVideoWrapper.git"
    )
    
    for node_url in "${custom_nodes[@]}"; do
        log_info "Instalando custom node: $node_url"
        comfy --workspace "$WORKSPACE_DIR" node install "$node_url" || {
            log_warning "Falha ao instalar custom node: $node_url"
        }
    done
    
    log_info "Instalação de custom nodes e dependências concluída."
}

# ============================================================================
# INSTALAÇÃO DO SAGEATTENTION (OTIMIZADO PARA EXECUTAR PRIMEIRO EM BACKGROUND)
# ============================================================================

install_sageattention() {
    log_info "Iniciando instalação do SageAttention em background (otimizado)..."
    
    (
        log_info "Clonando repositório SageAttention..."
        if [ -d "SageAttention" ]; then
            rm -rf SageAttention
        fi
        
        git clone https://github.com/thu-ml/SageAttention.git
        cd SageAttention
        
        log_info "Instalando SageAttention (isso pode demorar 5+ minutos)..."
        python3 setup.py install
        
        log_info "SageAttention instalado com sucesso!"
        
        # Criar arquivo de flag para indicar que a instalação terminou
        touch "$INSTALL_BASE_DIR/.sageattention_ready"
        
    ) &
    
    local sage_pid=$!
    log_info "Instalação do SageAttention iniciada em background. PID: $sage_pid"
    
    return 0
}

# ============================================================================
# DOWNLOAD DE MODELOS
# ============================================================================

download_models() {
    log_info "Iniciando downloads de modelos em paralelo..."
    
    # --- Criação de Diretórios ---
    mkdir -p "$WORKSPACE_DIR/models/diffusion_models"
    mkdir -p "$WORKSPACE_DIR/models/clip_vision"
    mkdir -p "$WORKSPACE_DIR/models/vae"
    mkdir -p "$WORKSPACE_DIR/models/text_encoders"
    
    local pids=()
    
    # --- Lista de Modelos ---
    local models=(
        "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Wan2_1-Anisora-I2V-480P-14B_fp16.safetensors:$WORKSPACE_DIR/models/diffusion_models"
        "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/vae/wan_2.1_vae.safetensors:$WORKSPACE_DIR/models/vae"
        "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/text_encoders/umt5_xxl_fp16.safetensors:$WORKSPACE_DIR/models/text_encoders"
        "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors:$WORKSPACE_DIR/models/text_encoders"
        "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/clip_vision/clip_vision_h.safetensors:$WORKSPACE_DIR/models/clip_vision"
        "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/diffusion_models/wan2.1_i2v_480p_14B_fp8_e4m3fn.safetensors:$WORKSPACE_DIR/models/diffusion_models"
        "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Wan2_1-Anisora-I2V-480P-14B_fp8_e4m3fn.safetensors:$WORKSPACE_DIR/models/diffusion_models"
    )
    
    # --- Inicia os Downloads em Background ---
    for model_info in "${models[@]}"; do
        IFS=':' read -r url dest_dir <<< "$model_info"
        (
            download_file "$url" "$dest_dir"
        ) &
        pids+=($!)
    done
    
    log_info "Downloads de modelos iniciados em paralelo. PIDs: ${pids[*]}"
    
    # --- Aguarda a Conclusão ---
    local has_errors=0
    log_info "Aguardando conclusão dos downloads de modelos..."
    for pid in "${pids[@]}"; do
        if ! wait "$pid"; then
            log_warning "Download em background (PID: $pid) falhou."
            has_errors=1
        fi
    done
    
    if [ "$has_errors" -eq 1 ]; then
        log_warning "Alguns downloads de modelos falharam. Verifique os logs."
    else
        log_info "Todos os downloads de modelos foram concluídos com sucesso."
    fi
    
    return "$has_errors"
}

# ============================================================================
# AGUARDAR SAGEATTENTION
# ============================================================================

wait_for_sageattention() {
    log_info "Verificando instalação do SageAttention..."
    
    # Aguarda até que o arquivo de flag seja criado
    while [ ! -f "$INSTALL_BASE_DIR/.sageattention_ready" ]; do
        log_info "SageAttention ainda instalando... aguardando 30 segundos"
        sleep 30
    done
    
    # Remove o arquivo de flag
    rm -f "$INSTALL_BASE_DIR/.sageattention_ready"
    
    log_info "SageAttention está pronto!"
}

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

launch_comfyui() {
    log_info "Iniciando ComfyUI na porta 8818 com SageAttention e preview latent2rgb..."
    log_info "Workspace do ComfyUI: $WORKSPACE_DIR"
    log_info "Arquivo de workflow disponível em: $WORKFLOW_FILE"
    log_info "Para usar o workflow, arraste e solte o arquivo '$WORKFLOW_FILENAME' na interface do ComfyUI."
    log_info "Acesse o ComfyUI em: http://127.0.0.1:8818 (ou o IP da sua máquina)"
    
    # Inicia o ComfyUI com SageAttention e preview latent2rgb
    exec comfy --workspace "$WORKSPACE_DIR" launch -- --fast --listen 0.0.0.0 --port 8818 --use-sage-attention --preview-method latent2rgb
}

# ============================================================================
# EXECUÇÃO PRINCIPAL
# ============================================================================

main() {

    # Passo 8: OTIMIZAÇÃO - Iniciar SageAttention PRIMEIRO (processo mais longo)
    log_info "================================================="
    log_info "Iniciando SageAttention em background (processo mais longo - ~5+ minutos)..."
    log_info "Enquanto isso, continuaremos com as outras instalações..."
    log_info "================================================="
    install_sageattention

    
    log_info "Iniciando instalação dinâmica do ComfyUI..."
    log_info "Diretório de instalação base: $INSTALL_BASE_DIR"
    log_info "================================================="
    
    # Passo 1: Limpar instalações existentes no diretório alvo
    cleanup_existing_installation
    
    # Passo 2: Instalar comfy-cli
    install_comfy_cli
    
    # Passo 3: Instalar HuggingFace tools
    install_hf_tools
    
    # Passo 4: Instalar ComfyUI
    install_comfyui
    
    # Passo 5: Baixar o workflow
    download_workflow
    
    # Passo 6: Instalar custom nodes e dependências do workflow
    install_custom_nodes_and_deps
    
    # Passo 7: Iniciar downloads de modelos (processo longo)
    log_info "================================================="
    log_info "Iniciando downloads de modelos..."
    log_info "================================================="
    download_models
    
    # Passo 8: Aguardar SageAttention (que provavelmente ainda estará rodando)
    wait_for_sageattention
    
    # Passo 9: Iniciar o ComfyUI
    log_info "================================================="
    log_info "Instalação concluída com sucesso!"
    log_info "================================================="
    launch_comfyui
}

# ============================================================================
# EXECUÇÃO DO SCRIPT
# ============================================================================

# Captura interrupções (Ctrl+C)
trap 'log_error "Script interrompido pelo usuário."; kill $(jobs -p) 2>/dev/null; exit 1' INT TERM

# Executa a função principal
main "$@"