#!/usr/bin/env bash
# setup_comfyui_wan22.sh
# Script baseado no CODE1 com modelos/nodes do CODE2 e melhorias de robustez

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="${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:-}"
HF_TOKEN="${HF_TOKEN:-}"
VENV_DIR="${VENV_DIR:-/root/comfy/.venv}"

# Suprimir avisos do pip quando executado como root
export PIP_ROOT_USER_ACTION=ignore

# Performance
export MAX_JOBS="${MAX_JOBS:-32}"
export HF_HUB_ENABLE_HF_TRANSFER=1
export HF_TRANSFER_CONCURRENCY="${HF_TRANSFER_CONCURRENCY:-16}"
export NVCC_APPEND_FLAGS="${NVCC_APPEND_FLAGS:---threads 8}"
export UV_SYSTEM_PYTHON=1
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

PYTHON_BIN=""

# Lista de downloads (formato: URL|TIPO|NOME_OPCIONAL)
DOWNLOAD_FILES=(
    "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/resolve/main/I2V/Wan2_2-I2V-A14B-LOW_fp8_e4m3fn_scaled_KJ.safetensors|diffusion_models|"
    "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/resolve/main/I2V/Wan2_2-I2V-A14B-HIGH_fp8_e4m3fn_scaled_KJ.safetensors|diffusion_models|"
    "https://huggingface.co/jrewingwannabe/Wan2.2-Lightning_I2V-A14B-4steps-lora/resolve/main/Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors|loras|"
    "https://huggingface.co/jrewingwannabe/Wan2.2-Lightning_I2V-A14B-4steps-lora/resolve/main/Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors|loras|"
    "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/LoRAs/AniSora/Wan2_2_I2V_AniSora_3_2_HIGH_rank_64_fp16.safetensors|loras|"
    "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Wan2_1_VAE_bf16.safetensors|vae|"
    "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/open-clip-xlm-roberta-large-vit-huge-14_visual_fp16.safetensors|clip_vision|"
    "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/umt5-xxl-enc-bf16.safetensors|text_encoders|"
    "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/general/Swin2SR_ClassicalSR_X2_64.pth|upscale_models|"
    "https://huggingface.co/adbrasi/download_models/resolve/main/22-nsfw-HIGH-e6.safetensors|loras|"
    "https://huggingface.co/adbrasi/download_models/resolve/main/NSFW-22-H-e8.safetensors|loras|"
    "https://huggingface.co/adbrasi/download_models/resolve/main/NSFW-22-L-e8.safetensors|loras|"
    "https://huggingface.co/adbrasi/download_models/resolve/main/smoothLoopv2_000002450_high_noise.safetensors|loras|"
    "https://huggingface.co/adbrasi/download_models/resolve/main/smoothLoopv2_000002450_low_noise.safetensors|loras|"
    "https://huggingface.co/adbrasi/download_models/resolve/main/ultimate_sex_000006500_high_noise.safetensors|loras|"
    "https://huggingface.co/adbrasi/download_models/resolve/main/ultimate_sex_000006500_low_noise.safetensors|loras|"
)

# Custom nodes (do CODE2)
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/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/adbrasi/upavidito"
    "https://github.com/wallish77/wlsh_nodes"
    "https://github.com/raindrop313/ComfyUI-WanVideoStartEndFrames"
    "https://github.com/MushroomFleet/DJZ-Nodes"
    "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
    "https://github.com/MoonGoblinDev/Civicomfy"
)

download_hf() {
    local hf_path="$1"
    local target_dir="$2"
    local filename="$3"
    local revision="${4:-}"

    if [[ -z "$filename" ]]; then
        filename="$(basename "${hf_path%%\?*}")"
    fi

    mkdir -p "$target_dir"
    if [[ -f "$target_dir/$filename" ]]; then
        local existing_size
        existing_size=$(stat -c%s "$target_dir/$filename" 2>/dev/null || echo 0)
        if [[ "$existing_size" -gt 1048576 ]]; then
            log_success "Já existe: $filename"
            return 0
        fi
    fi

    local owner="${hf_path%%/*}"
    local remainder="${hf_path#*/}"
    if [[ "$owner" == "$hf_path" ]]; then
        log_error "Caminho HuggingFace inválido: $hf_path"
        return 1
    fi

    local repo_segment="${remainder%%/*}"
    if [[ "$repo_segment" == "$remainder" ]]; then
        log_error "Arquivo não especificado no caminho: $hf_path"
        return 1
    fi

    remainder="${remainder#*/}"
    if [[ -z "$remainder" ]]; then
        log_error "Arquivo não especificado no caminho: $hf_path"
        return 1
    fi

    if [[ "$repo_segment" == *@* ]]; then
        if [[ -z "$revision" ]]; then
            revision="${repo_segment#*@}"
        fi
        repo_segment="${repo_segment%@*}"
    fi

    local repo_id="$owner/$repo_segment"
    local file_path="$remainder"
    local attempt success temp_dir
    for attempt in 1 2 3; do
        temp_dir="$(mktemp -d)"
        success=false

        log_info "[HF] Tentativa $attempt/3: ${repo_id}/${file_path}${revision:+ @}${revision:+$revision}"

        if command -v hf >/dev/null 2>&1; then
            local hf_cmd=(hf download "$repo_id" "$file_path")
            if [[ -n "$revision" ]]; then
                hf_cmd+=(--revision "$revision")
            fi
            hf_cmd+=(--local-dir "$temp_dir" --quiet)
            if [[ -n "$HF_TOKEN" ]]; then
                hf_cmd+=(--token "$HF_TOKEN")
            fi
            if HF_HUB_ENABLE_HF_TRANSFER=1 "${hf_cmd[@]}"; then
                success=true
            fi
        fi

        if [[ "$success" == false ]]; then
            local fallback_revision="${revision:-main}"
            local direct_url="https://huggingface.co/${repo_id}/resolve/${fallback_revision}/${file_path}"
            log_warn "[HF] Tentativa $attempt usando fallback HTTP"
            if command -v aria2c >/dev/null 2>&1; then
                local aria_cmd=(aria2c -c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10 --dir "$temp_dir" --out "$filename")
                if [[ -n "$HF_TOKEN" ]]; then
                    aria_cmd+=(--header="Authorization: Bearer $HF_TOKEN")
                fi
                aria_cmd+=("$direct_url")
                if "${aria_cmd[@]}"; then
                    success=true
                fi
            fi
            if [[ "$success" == false ]] && command -v curl >/dev/null 2>&1; then
                local curl_cmd=(curl -L -# -o "$temp_dir/$filename")
                if [[ -n "$HF_TOKEN" ]]; then
                    curl_cmd+=(-H "Authorization: Bearer $HF_TOKEN")
                fi
                curl_cmd+=("$direct_url")
                if "${curl_cmd[@]}"; then
                    success=true
                fi
            fi
            if [[ "$success" == false ]] && command -v wget >/dev/null 2>&1; then
                local wget_cmd=(wget -q --show-progress -c -O "$temp_dir/$filename")
                if [[ -n "$HF_TOKEN" ]]; then
                    wget_cmd+=(--header="Authorization: Bearer $HF_TOKEN")
                fi
                wget_cmd+=("$direct_url")
                if "${wget_cmd[@]}"; then
                    success=true
                fi
            fi
        fi

        if [[ "$success" == true ]]; then
            local downloaded
            downloaded="$(find "$temp_dir" -type f -name "$filename" -print -quit)"
            if [[ -z "$downloaded" ]]; then
                downloaded="$(find "$temp_dir" -type f -print -quit)"
            fi
            if [[ -n "$downloaded" ]]; then
                mv "$downloaded" "$target_dir/$filename"
                rm -rf "$temp_dir"
                if [[ -s "$target_dir/$filename" ]]; then
                    log_success "Download concluído: $filename"
                    return 0
                fi
                log_warn "Arquivo vazio após download: $filename"
                rm -f "$target_dir/$filename"
            fi
        fi

        rm -rf "$temp_dir"
        log_warn "Tentativa $attempt falhou para $filename"
        sleep $((2 * attempt))
    done

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


download_file() {
    local url="$1"
    local target_dir="$2"
    local filename="$3"

    if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]] && [[ -n "$CIVITAI_TOKEN" ]]; then
        if [[ "$url" == *"?"* ]]; then
            url="${url}&token=${CIVITAI_TOKEN}"
        else
            url="${url}?token=${CIVITAI_TOKEN}"
        fi
    fi
    mkdir -p "$target_dir"

    local final_name="$filename"
    if [[ -z "$final_name" ]]; then
        final_name="$(basename "${url%%\?*}")"
    fi

    if [[ -f "$target_dir/$final_name" ]]; then
        local existing_size
        existing_size=$(stat -c%s "$target_dir/$final_name" 2>/dev/null || echo 0)
        if [[ "$existing_size" -gt 524288 ]]; then
            log_success "Já existe: $final_name"
            return 0
        fi
    fi

    local attempt success
    for attempt in 1 2 3; do
        success=false
        log_info "[DL] Tentativa $attempt/3: ${final_name}"
        if command -v aria2c >/dev/null 2>&1; then
            if aria2c -c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10 --dir "$target_dir" --out "$final_name" "$url"; then
                success=true
            fi
        fi
        if [[ "$success" == false ]] && command -v wget >/dev/null 2>&1; then
            if wget -q --show-progress -c -O "$target_dir/$final_name" "$url"; then
                success=true
            fi
        fi
        if [[ "$success" == false ]] && command -v curl >/dev/null 2>&1; then
            if curl -L -# -C - -o "$target_dir/$final_name" "$url"; then
                success=true
            fi
        fi
        if [[ "$success" == true ]]; then
            if [[ -s "$target_dir/$final_name" ]]; then
                log_success "Download concluído: $final_name"
                return 0
            fi
            log_warn "Arquivo vazio após download: $final_name"
            rm -f "$target_dir/$final_name"
        fi
        log_warn "Tentativa $attempt falhou para $final_name"
        sleep $((2 * attempt))
    done

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

download_mega() {
    local url="$1"
    local target_dir="$2"

    url="${url#mega://}"
    mkdir -p "$target_dir"

    if ! command -v megadl >/dev/null 2>&1; then
        log_warn "megadl não encontrado; tente instalar megatools manualmente"
        return 1
    fi

    local attempt
    for attempt in 1 2 3; do
        log_info "[MEGA] Tentativa $attempt/3"
        local temp_dir
        temp_dir="$(mktemp -d)"
        if megadl --limit-speed 0 --path "$temp_dir" "$url"; then
            local moved=0
            while IFS= read -r -d '' file; do
                if [[ ! -f "$file" ]]; then
                    continue
                fi
                local base
                base="$(basename "$file")"
                if [[ -f "$target_dir/$base" ]]; then
                    log_success "Já existe: $base"
                else
                    if mv "$file" "$target_dir/$base"; then
                        log_success "Download Mega.nz concluído: $base"
                    else
                        log_warn "Falha ao mover $base para destino final"
                        continue
                    fi
                fi
                moved=1
            done < <(find "$temp_dir" -type f -print0)

            rm -rf "$temp_dir"
            if (( moved )); then
                return 0
            fi
            log_warn "[MEGA] Nenhum arquivo detectado após download"
            sleep $((2 * attempt))
        else
            log_warn "Tentativa $attempt falhou para Mega.nz"
            rm -rf "$temp_dir"
            sleep $((2 * attempt))
        fi
    done

    log_error "Falha ao baixar de Mega.nz: $url"
    return 1
}



process_downloads() {
    local failed=0
    for entry in "${DOWNLOAD_FILES[@]}"; do
        IFS='|' read -r raw_url type filename <<< "$entry"
        local target_dir="$MODELS_DIR/$type"
        mkdir -p "$target_dir"

        if [[ "$raw_url" == mega://* ]]; then
            if ! download_mega "$raw_url" "$target_dir"; then
                ((failed++))
            fi
        elif [[ "$raw_url" == hf://* ]]; then
            local hf_path="${raw_url#hf://}"
            if ! download_hf "$hf_path" "$target_dir" "$filename"; then
                ((failed++))
            fi
        elif [[ "$raw_url" == https://huggingface.co/* ]]; then
            local sanitized_url="${raw_url%%\?*}"
            local hf_relative="${sanitized_url#https://huggingface.co/}"
            IFS='/' read -ra hf_parts <<< "$hf_relative"
            if (( ${#hf_parts[@]} >= 4 )) && [[ "${hf_parts[2]}" == "resolve" || "${hf_parts[2]}" == "blob" || "${hf_parts[2]}" == "raw" ]]; then
                local hf_owner="${hf_parts[0]}"
                local hf_repo="${hf_parts[1]}"
                local hf_revision="${hf_parts[3]}"
                local start_index=4
                local hf_file_path=""
                for ((i=start_index; i<${#hf_parts[@]}; i++)); do
                    if [[ -n "$hf_file_path" ]]; then
                        hf_file_path+='/'
                    fi
                    hf_file_path+="${hf_parts[i]}"
                done
                if [[ -n "$hf_file_path" ]]; then
                    local hf_path_combined="${hf_owner}/${hf_repo}/${hf_file_path}"
                    if ! download_hf "$hf_path_combined" "$target_dir" "$filename" "$hf_revision"; then
                        ((failed++))
                    fi
                else
                    if ! download_file "$raw_url" "$target_dir" "$filename"; then
                        ((failed++))
                    fi
                fi
            else
                if ! download_file "$raw_url" "$target_dir" "$filename"; then
                    ((failed++))
                fi
            fi
        else
            if ! download_file "$raw_url" "$target_dir" "$filename"; then
                ((failed++))
            fi
        fi
    done

    if (( failed > 0 )); then
        return 1
    fi
    return 0
}

clone_repo() {
    local url="$1"
    local dest="$2"

    local name
    name="$(basename "$dest")"

    if [[ -d "$dest/.git" ]]; then
        if ! git -C "$dest" pull --ff-only --no-rebase >/dev/null 2>&1; then
            log_warn "Falha ao atualizar $name"
        else
            log_success "Atualizado: $name"
        fi
    else
        if ! git clone --depth 1 "$url" "$dest" >/dev/null 2>&1; then
            log_warn "Falha ao clonar $name"
            return
        fi
        log_success "Clonado: $name"
    fi

    if [[ -f "$dest/requirements.txt" ]]; then
        if [[ -n "$PYTHON_BIN" ]]; then
            if ! "$PYTHON_BIN" -m pip install -q -r "$dest/requirements.txt" >/dev/null 2>&1; then
                log_warn "Falha ao instalar requirements de $name"
            else
                log_info "Dependências instaladas para $name"
            fi
        else
            log_warn "PYTHON_BIN não definido; pulei requirements de $name"
        fi
    fi
}

# ========== INSTALAÇÃO ==========
log_info "========================================="
log_info " ComfyUI + Wan 2.2 Setup"
log_info "========================================="

# 1. Dependências
log_info "[1/7] Instalando dependências do sistema..."
if command -v apt-get >/dev/null 2>&1; then
    apt-get update -qq || log_warn "apt-get update falhou; seguindo assim mesmo"
    apt-get install -y -qq python3-venv aria2 megatools git wget curl || log_warn "apt-get install retornou erro"
else
    log_warn "apt-get indisponível; certifique-se de que python3-venv, aria2, megatools, git, wget e curl estão instalados"
fi
log_success "Dependências do sistema tratadas"

# 2. Preparar ambiente virtual
log_info "[2/7] Preparando ambiente virtual..."
if [[ ! -d "$VENV_DIR/bin" ]]; then
    python3 -m venv "$VENV_DIR"
    log_success "Ambiente virtual criado em $VENV_DIR"
else
    log_info "Ambiente virtual já existe em $VENV_DIR"
fi
. "$VENV_DIR/bin/activate"
PYTHON_BIN="$VENV_DIR/bin/python"

"$PYTHON_BIN" -m pip install -U pip wheel setuptools -q || log_warn "Falha na atualização do pip"
"$PYTHON_BIN" -m pip install -U "huggingface_hub[cli,hf_transfer]" comfy-cli -q || log_warn "Falha instalando huggingface_hub/comfy-cli"
log_success "Ambiente virtual pronto e comfy-cli instalado"

# 3. Instalar ComfyUI
log_info "[3/7] Instalando ComfyUI..."
if [[ -f "$COMFY_DIR/main.py" ]]; then
    log_warn "ComfyUI já presente"
else
    comfy --skip-prompt install --fast-deps --nvidia || log_error "Instalação do ComfyUI falhou"
fi

REQ_FILE="$COMFY_DIR/requirements.txt"
if [[ -f "$REQ_FILE" ]]; then
    log_info "Verificando dependências do ComfyUI em: $REQ_FILE"
    if [[ -n "$PYTHON_BIN" ]]; then
        if "$PYTHON_BIN" -m pip install -q -r "$REQ_FILE"; then
            log_success "Requirements do ComfyUI instalados/validados"
        else
            log_warn "Falha ao instalar requirements do ComfyUI; verifique o log acima e reexecute se necessário"
        fi
    else
        log_warn "PYTHON_BIN não definido; pulei requirements do ComfyUI"
    fi
else
    log_warn "requirements.txt do ComfyUI não encontrado em $REQ_FILE"
fi

log_success "ComfyUI pronto"

# 4. Instalar PyTorch (CUDA 12.8)
log_info "[4/7] Instalando PyTorch CUDA 12.8..."
"$PYTHON_BIN" -m pip install --force --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 || log_error "Falha ao instalar PyTorch"
log_success "PyTorch configurado"

# 5. Instalar SageAttention (condicional)
log_info "[5/7] Instalando SageAttention..."
if "$PYTHON_BIN" - <<'PY' >/dev/null 2>&1
import torch, sys
sys.exit(0 if torch.cuda.is_available() else 1)
PY
then
    if "$PYTHON_BIN" -m pip install "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp312-cp312-linux_x86_64.whl" -q; then
        if "$PYTHON_BIN" - <<'PY' >/dev/null 2>&1
import sageattention
PY
        then
            log_success "SageAttention instalado"
        else
            log_warn "SageAttention instalado, mas não pôde ser importado"
        fi
    else
        log_warn "Falha ao instalar SageAttention"
    fi
else
    log_warn "CUDA indisponível; pulando SageAttention"
fi

# 6. Downloads de modelos
log_info "[6/7] Baixando modelos..."
mkdir -p "$MODELS_DIR"
if process_downloads; then
    log_success "Downloads concluídos (ou já presentes)"
else
    log_warn "Alguns downloads falharam; reexecute o script para tentar novamente"
fi

# 7. Custom nodes
log_info "[7/7] Instalando custom nodes..."
CN_DIR="$COMFY_DIR/custom_nodes"
mkdir -p "$CN_DIR"
node_total=${#CUSTOM_NODES[@]}
node_index=0
set +e
for repo in "${CUSTOM_NODES[@]}"; do
    ((node_index++))
    log_info "[$node_index/$node_total] $(basename "$repo")"
    clone_repo "$repo" "$CN_DIR/$(basename "$repo")"
done
set -e
log_success "Custom nodes processados"


SAGE_FLAG=""
if "$PYTHON_BIN" - <<'PY' >/dev/null 2>&1
import sageattention
PY
then
    SAGE_FLAG="--use-sage-attention"
fi

log_info "========================================="
log_info "Iniciando ComfyUI em http://${COMFY_HOST}:${COMFY_PORT}"
log_info "========================================="

cd "$COMFY_DIR"
exec comfy launch -- ${SAGE_FLAG:+$SAGE_FLAG} --listen "$COMFY_HOST" --port "$COMFY_PORT"