#!/usr/bin/env bash
# setup_comfyui_simple.sh
# Script simples para instalação do ComfyUI

set -euo pipefail

# Cores
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="0.0.0.0"
COMFY_PORT="8818"
VENV_DIR="/root/comfy/.venv"
CIVITAI_TOKEN="4fcb2834969399006a736ee402b061e5"

# Performance
export MAX_JOBS=32
export HF_HUB_ENABLE_HF_TRANSFER=1

# Lista de downloads - Formato: "URL|TIPO|NOME_OPCIONAL"
DOWNLOAD_FILES=(
    # Mega primeiro (sem nome forçado - usar nome original)
    #"mega://https://mega.nz/file/gIRTFQSQ#no6Ay3JLE9LVRi7ib9O-Jc0CW7XmG046kCgpCzDg1tY|loras|"
    
    # Checkpoints
    "https://civitai.com/api/download/models/1828803?type=Model&format=SafeTensor&size=full&fp=fp16|checkpoints|ultimateHentaiAnimeRXTRexAnime_rxV1.safetensors"
    "https://civitai.com/api/download/models/2167369?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|wai15.safetensors"
    
    # LoRAs Civitai
    "https://civitai.com/api/download/models/1268294?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1715330?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1715330?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1499397?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1779002?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1114313?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1780244?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1613410?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1833157?type=Model&format=SafeTensor|embeddings|"


    "https://civitai.com/api/download/models/2121199?type=Model&format=Other|embeddings|"
    
    # Upscalers
    "https://huggingface.co/Kim2091/AnimeSharp/resolve/main/4x-AnimeSharp.pth|upscale_models|4x-AnimeSharp.pth"
    "https://huggingface.co/Kim2091/AnimeSharpV3/resolve/main/2x-AnimeSharpV3.pth|upscale_models|2x-AnimeSharpV3.pth"
    "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2.pth|upscale_models|4x-UltraSharpV2.pth"
    "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2_Lite.pth|upscale_models|4x-UltraSharpV2_Lite.pth"
    
    # Ultralytics
    "https://huggingface.co/adbrasi/testedownload/resolve/main/99coins_anime_girl_face_m_seg.pt|ultralytics/bbox|99coins_anime_girl_face_m_seg.pt"
)

# Custom nodes
CUSTOM_NODES=(
    "https://github.com/adbrasi/huggpackreator"
    "https://github.com/adbrasi/packreator_processor"
    "https://github.com/adbrasi/Packreator_manager"
    "https://github.com/adbrasi/cezarsave34"
    "https://github.com/adbrasi/pageonetor"
    "https://github.com/adbrasi/pakreatorio"
    "https://github.com/adbrasi/WaterMark_bumbumzin"
    "https://github.com/adbrasi/marcadaguita"
    "https://github.com/adbrasi/randomico"
    "https://github.com/adbrasi/groqrouter"
    "https://github.com/adbrasi/find_charakito"
    "https://github.com/adbrasi/randomsizito"
    "https://github.com/adbrasi/importex"
    "https://github.com/adbrasi/storitadifusita"
    "https://github.com/adbrasi/attentionPPM"
    "https://github.com/ClownsharkBatwing/RES4LYF"
    "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
    "https://github.com/sipherxyz/comfyui-art-venture"
    "https://github.com/pamparamm/sd-perturbed-attention"
    "https://github.com/KoreTeknology/ComfyUI-Universal-Styler"
    "https://github.com/WASasquatch/was-node-suite-comfyui"
    "https://github.com/chflame163/ComfyUI_LayerStyle"
    "https://github.com/newtextdoc1111/ComfyUI-Autocomplete-Plus"
    "https://github.com/pythongosssss/ComfyUI-WD14-Tagger"
    "https://github.com/ltdrdata/ComfyUI-Impact-Pack"
    "https://github.com/pythongosssss/ComfyUI-Custom-Scripts"
    "https://github.com/rgthree/rgthree-comfy"
    "https://github.com/ssitu/ComfyUI_UltimateSDUpscale"
    "https://github.com/adbrasi/Importador"
    "https://github.com/adbrasi/GetFirstTag"
    "https://github.com/adbrasi/comfydodi"
    "https://github.com/omar92/ComfyUI-QualityOfLifeSuit_Omar92"
    "https://github.com/Cezarsaint/blacklisto"
    "https://github.com/TinyTerra/ComfyUI_tinyterraNodes"
    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
    "https://github.com/Cezarsaint/rand0micoUploaderLoven"
    "https://github.com/adbrasi/pixivmosaic"
    "https://github.com/adbrasi/futfilter"
    "https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"
    "https://github.com/shiimizu/ComfyUI_smZNodes"
    "https://github.com/CoreyCorza/ComfyUI-CRZnodes"
    "https://github.com/ClownsharkBatwing/RES4LYF"
    "https://github.com/MoonGoblinDev/Civicomfy"
    "https://github.com/ltdrdata/ComfyUI-Impact-Pack"
    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
    "https://github.com/rgthree/rgthree-comfy"
)

# Download HuggingFace
download_hf() {
    local url="$1"
    local target_dir="$2"
    local filename="$3"
    
    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
        log_success "Já existe: $filename"
        return 0
    fi
    
    log_info "Baixando HF: $filename"
    
    if command -v hf >/dev/null 2>&1; then
        cd "$target_dir"
        hf download "$(dirname "$url")" "$(basename "$url")" --local-dir . 2>/dev/null
        cd - >/dev/null
    else
        wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$url"
    fi
}

# Download Mega
download_mega() {
    local url="$1"
    local target_dir="$2"
    
    # Remover prefixo mega://
    url="${url#mega://}"
    
    log_info "Baixando do Mega..."
    
    cd "$target_dir"
    # Usar megadl simples sem forçar nome
    if megadl "$url" 2>/dev/null; then
        log_success "Mega download OK"
    else
        log_warn "Mega download falhou"
    fi
    cd - >/dev/null
}

# Download normal
download_file() {
    local url="$1"
    local target_dir="$2"
    local filename="$3"
    
    # Adicionar token Civitai se necessário
    if [[ "$url" == *"civitai.com"* ]] && [[ "$url" != *"token="* ]]; then
        url="${url}&token=${CIVITAI_TOKEN}"
    fi
    
    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
        log_success "Já existe: $filename"
        return 0
    fi
    
    log_info "Baixando: ${filename:-arquivo}"
    
    # Tentar com aria2c primeiro
    if command -v aria2c >/dev/null 2>&1; then
        if [ -n "$filename" ]; then
            aria2c -c -x 4 -s 4 --console-log-level=warn --dir="$target_dir" --out="$filename" "$url" || \
            wget -q --show-progress -c -O "$target_dir/$filename" "$url" || \
            log_error "Download falhou: $filename"
        else
            aria2c -c -x 4 -s 4 --console-log-level=warn --dir="$target_dir" "$url" || \
            (cd "$target_dir" && wget -q --show-progress -c "$url") || \
            log_error "Download falhou: $url"
        fi
    elif [ -n "$filename" ]; then
        wget -q --show-progress -c -O "$target_dir/$filename" "$url" || \
        log_error "Download falhou: $filename"
    else
        (cd "$target_dir" && wget -q --show-progress -c "$url") || \
        log_error "Download falhou: $url"
    fi
}

# Processar downloads
process_downloads() {
    for entry in "${DOWNLOAD_FILES[@]}"; do
        IFS='|' read -r url type filename <<< "$entry"
        
        # Criar diretório
        local target_dir="$MODELS_DIR/$type"
        mkdir -p "$target_dir"
        
        # Baixar arquivo
        if [[ "$url" == mega://* ]]; then
            download_mega "$url" "$target_dir"
        elif [[ "$url" == *.huggingface.co/* ]]; then
            local hf_path="${url#*huggingface.co/}"
            download_hf "$hf_path" "$target_dir" "$filename"
        else
            download_file "$url" "$target_dir" "$filename"
        fi
    done
}

# Clone git repo
clone_repo() {
    local url="$1"
    local dest="$2"
    
    if [ -d "$dest/.git" ]; then
        git -C "$dest" pull --ff-only 2>/dev/null || true
    else
        git clone --depth 1 "$url" "$dest" 2>/dev/null || true
    fi
    
    # Instalar requirements se existir
    if [ -f "$dest/requirements.txt" ]; then
        python -m pip install -q -r "$dest/requirements.txt" 2>/dev/null || true
    fi
}

# ========== INSTALAÇÃO ==========

log_info "========================================="
log_info " ComfyUI Setup"
log_info "========================================="

# 1. Instalar dependências
log_info "[1/6] Instalando dependências do sistema..."

apt-get update -qq
apt-get install -y -qq python3-venv aria2 megatools git wget curl 2>/dev/null

log_success "Dependências do sistema instaladas"

log_info "[2/6] 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 -m pip install -U pip wheel setuptools -q
python -m pip install -U "huggingface_hub[cli,hf_transfer]" comfy-cli -q

log_success "Ambiente virtual pronto e comfy-cli instalado"

# 3. Instalar ComfyUI
log_info "[3/6] Instalando ComfyUI..."

if [ -f "$COMFY_DIR/main.py" ]; then
    log_warn "ComfyUI já existe"
else
    comfy --skip-prompt install --fast-deps --nvidia --version "0.3.58" 
fi

log_success "ComfyUI instalado"

# 4. Corrigir PyTorch para RTX 5090 se detectado
log_info "[4/6] Verificando GPU..."

GPU_INFO=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "")

if [[ "$GPU_INFO" == *"5090"* ]] || [[ "$GPU_INFO" == *"5080"* ]]; then
    log_warn "RTX 5090/5080 detectada - instalando PyTorch nightly"
    python -m pip install --force-reinstall --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
else
    log_info "GPU: ${GPU_INFO:-Não detectada}"
fi

log_success "PyTorch configurado"

# 5. Baixar modelos
log_info "[5/6] Baixando modelos..."

process_downloads

log_success "Downloads processados"

# 6. Instalar custom nodes
log_info "[6/6] Instalando custom nodes..."

CN_DIR="$COMFY_DIR/custom_nodes"
mkdir -p "$CN_DIR"

for repo in "${CUSTOM_NODES[@]}"; do
    node_name=$(basename "$repo")
    clone_repo "$repo" "$CN_DIR/$node_name"
done

log_success "Custom nodes instalados"

# Verificar e iniciar
log_info "========================================="
log_success "Instalação concluída!"
log_info "Iniciando ComfyUI..."
log_info "URL: http://localhost:$COMFY_PORT"
log_info "========================================="

cd "$COMFY_DIR"
exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"