#!/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"; }

drop_caches() {
    sync
    echo 3 > /proc/sys/vm/drop_caches || true
}

flush_cache_if_large() {
    local file_path="$1"
    if [ -f "$file_path" ]; then
        local sz
        sz=$(stat -c%s "$file_path" 2>/dev/null || echo 0)
        if [ "$sz" -ge 209715200 ]; then
            drop_caches
        fi
    else
        drop_caches
    fi
}

# Configuração
COMFY_BASE="/workspace/comfy"
COMFY_DIR="$COMFY_BASE/ComfyUI"
MODELS_DIR="$COMFY_DIR/models"
COMFY_HOST="0.0.0.0"
COMFY_PORT="8818"
VENV_DIR="$COMFY_BASE/.venv"
CIVITAI_TOKEN="4fcb2834969399006a736ee402b061e5"

export DEBIAN_FRONTEND=noninteractive
export HF_HOME="/workspace/.hf"
export HUGGINGFACE_HUB_CACHE="$HF_HOME"
export TRANSFORMERS_CACHE="$HF_HOME"
export TMPDIR="/workspace/.tmp"
export GIT_LFS_SKIP_SMUDGE=1
export PIP_NO_CACHE_DIR=off

PIP_FLAGS="--no-cache-dir --prefer-binary --only-binary=:all:"
ARIA="aria2c -c -x 2 -s 2 --disk-cache=0 --file-allocation=none --console-log-level=warn"

# Garantir diretórios necessários para instalar fora de /root e concentrar caches
mkdir -p "$COMFY_BASE" "$HF_HOME" "$TMPDIR"

# 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/2320393?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|perfectxl.safetensors"
    "https://huggingface.co/ChenkinNoob/ChenkinNoob-XL-V0.2/resolve/main/ChenkinNoob-XL-V0.2.safetensors|checkpoints|ChenkinNoob-XL-V0.2.safetensors"
    "https://civitai.com/api/download/models/1828803?type=Model&format=SafeTensor&size=full&fp=fp16|checkpoints|Hentai_Anime_RX.safetensors"
    
    # LoRAs Civitai
    "https://civitai.com/api/download/models/1818185?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2295041?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1245271?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2310692?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1164079?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1031892?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1452729?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1969737?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1941518?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1941518?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1431530?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1777446?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1565457?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2326498?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/938035?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1444623?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1295936?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2068919?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1280931?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2238110?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1453075?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1630029?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2296528?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1268294?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1148809?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/1804885?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1809575?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2135873?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2332149?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1517104?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1112990?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1379428?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1755071?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2298959?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1936647?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1583448?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1873259?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1570568?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1563658?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1445868?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1415706?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1373187?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/2066157?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1741819?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1833157?type=Model&format=SafeTensor|embeddings|"
    
    "https://huggingface.co/Anzhc/Anzhcs_YOLOs/resolve/main/Anzhc%20Face%20seg%201024%20v2%20y8n.pt|ultralytics/bbox|Anzhc_Faceseg_1024_v2_y8n.pt"
    "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth|sams|sam_vit_b_01ec64.pth"
    "https://huggingface.co/adbrasi/wanlotest/resolve/main/Eyeful_v2-Individual.pt|ultralytics/bbox|Eyeful_v2-Paired.pt"
    "https://huggingface.co/Bingsu/adetailer/resolve/main/hand_yolov8s.pt|ultralytics/bbox|hand_yolov8s.pt"
    "https://huggingface.co/adbrasi/wanlotest/resolve/main/ntd11_anime_nsfw_segm_v5-variant1.pt|ultralytics/bbox|ntd11_anime_nsfw_segm_v5-variant1.pt"
    "https://civitai.com/api/download/models/2121199?type=Model&format=Other|embeddings|"
    "https://civitai.com/api/download/models/1195487?type=Negative&format=Other|embeddings|"
    "https://civitai.com/api/download/models/1470551?type=Model&format=SafeTensor|embeddings|"
    # Upscalers
    "https://huggingface.co/adbrasi/wanlotest/resolve/main/2x-AnimeSharpV4_RCAN_fp16_op17.onnx|upscale_models|2x-AnimeSharpV4_RCAN_fp16_op17.onnx"
    "https://huggingface.co/Kim2091/AnimeSharp/resolve/main/4x-AnimeSharp.pth|upscale_models|4x-AnimeSharp.pth"
    "https://huggingface.co/adbrasi/wanlotest/resolve/main/2x-AnimeSharpV4_Fast_RCAN_PU.safetensors|upscale_models|2x-AnimeSharpV4_Fast_RCAN_PU.safetensors"
    "https://huggingface.co/Kim2091/AnimeSharpV3/resolve/main/2x-AnimeSharpV3.pth|upscale_models|2x-AnimeSharpV3.pth"
    "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.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"

    "https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/resolve/main/diffusion_pytorch_model_promax.safetensors|controlnet|controlnet_union_promax.safetensors"
    
    # 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/Cezarsaint/Packreator_managerMEita"
    "https://github.com/adbrasi/cezarsave34"
    "https://github.com/adbrasi/pageonetor"
    "https://github.com/QuietNoise/comfyui_queue_manager"
    "https://github.com/adbrasi/pakreatorio"
    "https://github.com/adbrasi/WaterMark_bumbumzin"
    "https://github.com/adbrasi/marcadaguita"
    "https://github.com/adbrasi/randomico"
    "https://github.com/kijai/ComfyUI-KJNodes"
    "https://github.com/ltdrdata/ComfyUI-Inspire-Pack"
    "https://github.com/Goshe-nite/comfyui-gps-supplements"
    "https://github.com/adbrasi/groqrouter"
    "https://github.com/adbrasi/find_charakito"
    "https://github.com/adbrasi/randomsizito"
    "https://github.com/zombieyang/sd-ppp"
    "https://github.com/adbrasi/importex"
    "https://github.com/adbrasi/storitadifusita"
    "https://github.com/adbrasi/attentionPPM"
    "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/Fannovel16/comfyui_controlnet_aux"
    "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/MoonGoblinDev/Civicomfy"
    "https://github.com/regiellis/ComfyUI-EasyIllustrious"
    "https://github.com/adbrasi/captionizador"
    "https://github.com/Jonseed/ComfyUI-Detail-Daemon"
    "https://github.com/WASasquatch/was-node-suite-comfyui"
    "https://github.com/fearnworks/ComfyUI_FearnworksNodes"
    "https://github.com/aria1th/ComfyUI-LogicUtils"
    "https://github.com/Extraltodeus/ComfyUI-AutomaticCFG"
    "https://github.com/pythongosssss/ComfyUI-Custom-Scripts"
    "https://github.com/ltdrdata/ComfyUI-Impact-Pack"
    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
    "https://github.com/huchenlei/ComfyUI_DanTagGen"
    "https://github.com/adbrasi/captionsaver"
    "https://github.com/adbrasi/captionizador"
)

# 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

    if [ -n "$filename" ]; then
        flush_cache_if_large "$target_dir/$filename"
    else
        drop_caches
    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
    drop_caches
}

# 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
            $ARIA --dir="$target_dir" --out="$filename" "$url" || \
            wget -q --show-progress -c -O "$target_dir/$filename" "$url" || \
            log_error "Download falhou: $filename"
        else
            $ARIA --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

    if [ -n "$filename" ]; then
        flush_cache_if_large "$target_dir/$filename"
    else
        drop_caches
    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 (comportamento original)
    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 --no-install-recommends python3-venv aria2 megatools git wget curl 2>/dev/null
apt-get clean
rm -rf /var/lib/apt/lists/*

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 ${PIP_FLAGS} -U pip wheel setuptools -q
python -m pip install ${PIP_FLAGS} -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 --workspace "$COMFY_DIR" install --fast-deps --nvidia 
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 ${PIP_FLAGS} --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/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 "========================================="

rm -rf "$TMPDIR"
mkdir -p "$TMPDIR"

cd "$COMFY_DIR"
exec comfy --workspace "$COMFY_DIR" launch -- --listen "$COMFY_HOST" --preview-method latent2rgb --front-end-version Comfy-Org/ComfyUI_frontend@latest --port "$COMFY_PORT"