#!/bin/bash

# ComfyUI Automated Setup Script
# Supports LightningAI, RunPod, VastAI, and generic Linux environments
# Author: Generated for automated ComfyUI deployment
# Modified based on user request.

set -e
set -o pipefail

# ============================================================================
# GLOBAL CONFIGURATIONS
# ============================================================================

# Civitai API Token
readonly CIVITAI_TOKEN="b1f7d1e6d35dc74f149e767ff7542d96"

# Custom Nodes to Install (from Git, handled by comfy-cli)
readonly 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/TinyTerra/ComfyUI_tinyterraNodes"
    "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/1038lab/ComfyUI-RMBG"
    "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/ltdrdata/ComfyUI-Impact-Pack"
    "https://github.com/pythongosssss/ComfyUI-Custom-Scripts"
    "https://github.com/rgthree/rgthree-comfy"
    "https://github.com/adbrasi/rand0micoUploader"
    "https://github.com/ssitu/ComfyUI_UltimateSDUpscale"
    "https://github.com/adbrasi/Importador"
    "https://github.com/adbrasi/comfydodi"
    "https://github.com/Cezarsaint/blacklisto"
    "https://github.com/omar92/ComfyUI-QualityOfLifeSuit_Omar92"
    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
    "https://github.com/adbrasi/pixivmosaic"
)

# Custom Nodes from Registry
readonly REGISTRY_NODES=(
    "civicomfy"
    "ComfyUI-RMBG"
)

# Files to Download - Format: "URL|target_subdir|filename"
readonly DOWNLOAD_FILES=(
    "https://civitai.com/api/download/models/1761560?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|"
    "https://civitai.com/api/download/models/1522463?type=Model&format=SafeTensor&size=full&fp=fp16|checkpoints|"
    "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/1499397?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1501009?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1304799?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1715751?type=Model&format=SafeTensor|loras"
    "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/adbrasi/testedownload/resolve/main/99coins_anime_girl_face_m_seg.pt|ultralytics/bbox|99coins_anime_girl_face_m_seg.pt"
    "https://huggingface.co/adbrasi/testedownload/resolve/main/ntd11_anime_nsfw_segm_v4_all.pt|ultralytics/bbox|ntd11_anime_nsfw_segm_v4_all.pt"
    "https://huggingface.co/adbrasi/testedownload/resolve/main/penis.pt|ultralytics/bbox|penis.pt"
)

# Model directories to create
readonly MODEL_DIRS=(
    "checkpoints"
    "loras"
    "upscale_models"
    "ultralytics/bbox"
    "Nudenet" # ### NOVO ### Adicionado para garantir que a pasta exista
)

# Workflow configuration
readonly WORKFLOW_URL="https://huggingface.co/adbrasi/testedownload/resolve/main/packreator_setup.json"

# Global variables
BASE_DIR=""
COMFY_INSTALL_DIR=""
MODELS_DIR=""
WORKFLOW_FILE=""
SUDO_PREFIX=""
ARIA2C_AVAILABLE=false
ENV_NAME=""

# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

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
}

is_root() {
    [ "$(id -u)" -eq 0 ]
}

# ============================================================================
# ENVIRONMENT DETECTION
# ============================================================================

detect_environment() {
    log_info "Detecting environment..."
    
    if [ -d "/teamspace/studios/this_studio" ]; then
        BASE_DIR="/teamspace/studios/this_studio"
        ENV_NAME="lightning"
    elif [ -d "/workspace" ] && [ ! -d "/teamspace/studios/this_studio" ]; then
        BASE_DIR="/workspace"
        ENV_NAME="runpod"
    elif is_root && [ ! -d "/teamspace/studios/this_studio" ] && [ ! -d "/workspace" ]; then
        BASE_DIR="/root"
        ENV_NAME="vastai"
    else
        BASE_DIR="$HOME"
        ENV_NAME="generic"
    fi
    
    COMFY_INSTALL_DIR="${BASE_DIR}/my_comfy_setup"
    MODELS_DIR="${COMFY_INSTALL_DIR}/models"
    WORKFLOW_FILE="${COMFY_INSTALL_DIR}/workflow.json"
    
    log_info "Environment detected: $ENV_NAME"
    log_info "Base directory: $BASE_DIR"
    log_info "ComfyUI install directory: $COMFY_INSTALL_DIR"
    log_info "Models directory: $MODELS_DIR"
}

# ============================================================================
# PERMISSION MANAGEMENT
# ============================================================================

setup_permissions() {
    log_info "Setting up permissions..."
    
    if is_root; then
        SUDO_PREFIX=""
        log_info "Running as root, no sudo needed"
    elif command_exists sudo; then
        SUDO_PREFIX="sudo "
        log_info "Non-root user with sudo available"
    else
        SUDO_PREFIX=""
        log_warning "Non-root user without sudo, some system operations may fail"
    fi
}

# ============================================================================
# SYSTEM DEPENDENCIES
# ============================================================================

install_system_dependencies() {
    log_info "Installing system dependencies..."
    
    if command_exists apt-get; then
        ${SUDO_PREFIX}apt-get update -qq || log_warning "Failed to update package list"
        ${SUDO_PREFIX}apt-get install -y -qq python3 python3-pip git || { log_error "Failed to install basic dependencies (python3, pip, git)"; exit 1; }
        if ${SUDO_PREFIX}apt-get install -y -qq aria2 >/dev/null 2>&1; then ARIA2C_AVAILABLE=true; log_info "aria2c installed successfully"; else ARIA2C_AVAILABLE=false; log_info "aria2c installation failed, will use wget/curl as fallback"; fi
    else
        log_warning "apt-get not available, skipping system package installation"
        if command_exists aria2c; then ARIA2C_AVAILABLE=true; log_info "aria2c already available"; fi
    fi
}

# ============================================================================
# COMFY-CLI INSTALLATION
# ============================================================================

install_comfy_cli() {
    log_info "Installing comfy-cli..."
    
    python3 -m pip install --upgrade comfy-cli || { log_error "Failed to install comfy-cli"; exit 1; }
    
    if ! command_exists comfy; then
        if [ -d "$HOME/.local/bin" ]; then export PATH="$HOME/.local/bin:$PATH"; log_info "Added $HOME/.local/bin to PATH"; fi
        if ! command_exists comfy; then log_error "comfy command not found in PATH after installation"; exit 1; fi
    fi
    
    log_info "comfy-cli version: $(comfy --version 2>/dev/null || echo 'unknown')"
}

# ============================================================================
# COMFYUI CORE INSTALLATION
# ============================================================================

install_comfyui_core() {
    log_info "Installing ComfyUI core..."
    comfy --workspace "$COMFY_INSTALL_DIR" --skip-prompt install --nvidia --version 0.3.42 || { log_error "Failed to install ComfyUI core"; exit 1; }
    log_info "ComfyUI core installed successfully"
}

# ============================================================================
# DIRECTORY STRUCTURE
# ============================================================================

create_model_directories() {
    log_info "Creating model directories..."
    for dir in "${MODEL_DIRS[@]}"; do mkdir -p "$MODELS_DIR/$dir"; log_info "Created directory: $MODELS_DIR/$dir"; done
    mkdir -p "$COMFY_INSTALL_DIR"
}

# ============================================================================
# CUSTOM NODES INSTALLATION
# ============================================================================

install_custom_nodes() {
    log_info "Installing custom nodes..."
    
    # Install nodes from registry
    for node in "${REGISTRY_NODES[@]}"; do
        log_info "Installing custom node from registry: $node"
        comfy --workspace "$COMFY_INSTALL_DIR" node registry-install "$node" || { log_warning "Failed to install custom node from registry: $node"; }
    done

    # Install nodes from Git repositories
    for node_url in "${CUSTOM_NODES[@]}"; do
        log_info "Installing custom node from Git: $node_url"
        comfy --workspace "$COMFY_INSTALL_DIR" node install "$node_url" || { log_warning "Failed to install custom node: $node_url"; }
    done
}

# ============================================================================
# SPECIAL NODE: WD14 TAGGER (GIT CLONE METHOD)
# ============================================================================

install_wd14_tagger_node() {
    log_info "Installing WD14 Tagger node via git clone..."
    local node_repo="https://github.com/pythongosssss/ComfyUI-WD14-Tagger.git"
    local node_name="ComfyUI-WD14-Tagger"
    local custom_nodes_dir="${COMFY_INSTALL_DIR}/custom_nodes"
    local target_dir="${custom_nodes_dir}/${node_name}"
    local requirements_file="${target_dir}/requirements.txt"

    # Ensure custom nodes directory exists
    mkdir -p "$custom_nodes_dir"

    if [ -d "$target_dir" ]; then
        log_info "WD14 Tagger node already exists. Skipping clone."
    else
        log_info "Cloning WD14 Tagger repository into $target_dir"
        if ! git clone "$node_repo" "$target_dir"; then
            log_error "Failed to clone WD14 Tagger repository."
            return 1
        fi
    fi

    if [ -f "$requirements_file" ]; then
        log_info "Installing WD14 Tagger dependencies from $requirements_file..."
        if ! python3 -m pip install -r "$requirements_file"; then
            log_warning "Failed to install WD14 Tagger dependencies. This might cause issues."
        else
            log_info "WD14 Tagger dependencies installed successfully."
        fi
    else
        log_info "No requirements.txt found for WD14 Tagger. Skipping dependency installation."
    fi
}

# ============================================================================
# SPECIAL NODE: STORY DIFFUSION
# ============================================================================

install_story_diffusion_node() {
    log_info "Installing special node: StoryDiffusion..."
    local node_repo="https://github.com/smthemex/ComfyUI_StoryDiffusion.git"
    local node_name="ComfyUI_StoryDiffusion"
    local custom_nodes_dir="${COMFY_INSTALL_DIR}/custom_nodes"
    local target_dir="${custom_nodes_dir}/${node_name}"
    local requirements_file="${target_dir}/requirements.txt"
    local commit_hash="e06e6d3f6543f6a849252b81e84d4fd636f09798"

    # Garante que o diretório de nós customizados exista
    mkdir -p "$custom_nodes_dir"

    if [ -d "$target_dir" ]; then
        log_info "StoryDiffusion node already exists. Pulling latest changes..."
        if ! git -C "$target_dir" fetch --all; then
            log_warning "Failed to fetch updates for StoryDiffusion. Continuing..."
        fi
    else
        log_info "Cloning StoryDiffusion repository into $target_dir"
        if ! git clone "$node_repo" "$target_dir"; then
            log_error "Failed to clone StoryDiffusion repository."
            return 1 
        fi
    fi

    # Checkout no commit específico
    log_info "Checking out commit $commit_hash..."
    if ! git -C "$target_dir" checkout "$commit_hash"; then
        log_error "Failed to checkout commit $commit_hash for StoryDiffusion."
        return 1
    fi

    # Instala as dependências se houver
    if [ -f "$requirements_file" ]; then
        log_info "Installing StoryDiffusion dependencies from $requirements_file..."
        if ! python3 -m pip install -r "$requirements_file"; then
            log_warning "Failed to install StoryDiffusion dependencies. This might cause issues."
        else
            log_info "StoryDiffusion dependencies installed successfully."
        fi
    else
        log_info "No requirements.txt found for StoryDiffusion. Skipping dependency installation."
    fi
}


# ============================================================================
# ### NOVO ### SPECIAL NODE: NUDENET
# ============================================================================

install_nudenet_node() {
    log_info "Installing special node: NudeNet..."
    local node_repo="https://github.com/phuvinh010701/ComfyUI-NudeNet.git"
    local node_name="ComfyUI-NudeNet"
    local custom_nodes_dir="${COMFY_INSTALL_DIR}/custom_nodes"
    local target_dir="${custom_nodes_dir}/${node_name}"
    local requirements_file="${target_dir}/requirements.txt"

    # Etapa 1: Clonar o repositório se não existir
    mkdir -p "$custom_nodes_dir"
    if [ -d "$target_dir" ]; then
        log_info "NudeNet node directory already exists. Skipping clone."
    else
        log_info "Cloning NudeNet repository into $target_dir"
        if ! git clone "$node_repo" "$target_dir"; then
            log_error "Failed to clone NudeNet repository."
            return 1 # Abortar a instalação deste node específico
        fi
    fi

    # Etapa 2: Instalar dependências do requirements.txt
    if [ -f "$requirements_file" ]; then
        log_info "Installing NudeNet dependencies from $requirements_file..."
        if ! python3 -m pip install -r "$requirements_file"; then
            log_warning "Failed to install NudeNet dependencies. This might cause issues."
        else
            log_info "NudeNet dependencies installed successfully."
        fi
    else
        log_info "No requirements.txt found for NudeNet. Skipping dependency installation."
    fi

    # Etapa 3: Baixar o modelo .onnx para a pasta correta
    log_info "Setting up NudeNet model..."
    local nudenet_model_dir="${MODELS_DIR}/Nudenet"
    local nudenet_model_url="https://d2xl8ijk56kv4u.cloudfront.net/models/nudenet.onnx"
    local nudenet_model_filename="nudenet.onnx"
    local nudenet_model_path="${nudenet_model_dir}/${nudenet_model_filename}"

    mkdir -p "$nudenet_model_dir" # Garante que o diretório de destino exista
    
    if [ -f "$nudenet_model_path" ]; then
        log_info "NudeNet model already exists at $nudenet_model_path. Skipping download."
    else
        log_info "Downloading NudeNet model to $nudenet_model_path..."
        if ! download_file "$nudenet_model_url" "$nudenet_model_dir" "$nudenet_model_filename"; then
            log_error "Failed to download the NudeNet model. The node may not work correctly."
        else
            log_info "NudeNet model downloaded successfully."
        fi
    fi
}

# ============================================================================
# DOWNLOAD FUNCTIONS
# ============================================================================

add_civitai_token() { local url="$1"; if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]]; then echo "${url}&token=${CIVITAI_TOKEN}"; else echo "$url"; fi; }

download_file() {
    local url="$1"; local target_dir="$2"; local filename="$3"; url=$(add_civitai_token "$url"); local download_cmd=""; local success=false
    if $ARIA2C_AVAILABLE && command_exists aria2c; then local aria_opts="-c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=0"; if [ -n "$filename" ]; then download_cmd="aria2c $aria_opts --dir=\"$target_dir\" --out=\"$filename\" \"$url\""; else download_cmd="aria2c $aria_opts --dir=\"$target_dir\" --content-disposition=true \"$url\""; fi; if eval "$download_cmd" >/dev/null 2>&1; then success=true; fi; fi
    if ! $success && command_exists wget; then local wget_opts="-q --show-progress -c"; if [ -n "$filename" ]; then download_cmd="wget $wget_opts -O \"$target_dir/$filename\" \"$url\""; else download_cmd="wget $wget_opts --content-disposition -P \"$target_dir\" \"$url\""; fi; if eval "$download_cmd"; then success=true; fi; fi
    if ! $success && command_exists curl; then log_warning "Using curl fallback for $url"; if [ -n "$filename" ]; then download_cmd="curl -L -C - -o \"$target_dir/$filename\" \"$url\""; else if curl --help 2>/dev/null | grep -q "output-dir"; then download_cmd="curl -L -J -O --output-dir \"$target_dir\" \"$url\""; else local fallback_name; fallback_name=$(basename "$url" | cut -d'?' -f1); [ -z "$fallback_name" ] && fallback_name="downloaded_file"; download_cmd="curl -L -C - -o \"$target_dir/$fallback_name\" \"$url\""; fi; fi; if eval "$download_cmd"; then success=true; fi; fi
    if ! $success; then log_error "Failed to download: $url"; return 1; fi; return 0
}

# ============================================================================
# PARALLEL DOWNLOADS
# ============================================================================

download_all_files() {
    log_info "Starting parallel downloads..."; local pids=(); local download_logs=()
    for file_spec in "${DOWNLOAD_FILES[@]}"; do IFS='|' read -r url target_subdir filename <<< "$file_spec"; local target_dir="$MODELS_DIR/$target_subdir"; local display_name="${filename:-$(basename "$url" | cut -d'?' -f1)}"; log_info "Starting download: $display_name -> $target_dir"; local log_file; log_file=$(mktemp); download_logs+=("$log_file"); ( if download_file "$url" "$target_dir" "$filename" >"$log_file" 2>&1; then echo "SUCCESS: Downloaded $display_name" >> "$log_file"; else echo "FAILED: Failed to download $display_name" >> "$log_file"; exit 1; fi ) & pids+=($!); done
    local failed_count=0; for i in "${!pids[@]}"; do if wait "${pids[$i]}"; then log_info "Download completed successfully"; else failed_count=$((failed_count + 1)); log_warning "Download failed - check logs"; fi; if [ -f "${download_logs[$i]}" ]; then cat "${download_logs[$i]}" >&2; rm -f "${download_logs[$i]}"; fi; done
    if [ $failed_count -gt 0 ]; then log_warning "$failed_count downloads failed, but continuing..."; else log_info "All downloads completed successfully"; fi
}

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

download_workflow() {
    log_info "Baixando o workflow de: $WORKFLOW_URL"; log_info "Salvando como: $WORKFLOW_FILE"; local success=false
    if command_exists wget; then if wget -q --show-progress -O "$WORKFLOW_FILE" "$WORKFLOW_URL"; then success=true; fi; fi
    if ! $success && command_exists curl; then if curl -L -o "$WORKFLOW_FILE" "$WORKFLOW_URL"; then success=true; fi; fi
    if ! $success && command_exists aria2c; then if aria2c -q -o "$(basename "$WORKFLOW_FILE")" -d "$(dirname "$WORKFLOW_FILE")" "$WORKFLOW_URL"; then success=true; fi; fi
    if ! $success; then log_error "Nenhuma ferramenta de download (wget, curl, aria2c) funcionou ou está disponível."; exit 1; fi
    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."
}

# ============================================================================
# WORKFLOW DEPENDENCIES
# ============================================================================

install_workflow_dependencies() {
    log_info "Instalando dependências do workflow a partir de: $WORKFLOW_FILE"
    if [ ! -f "$WORKFLOW_FILE" ]; then log_error "Arquivo de workflow não encontrado: $WORKFLOW_FILE"; return 1; fi
    comfy --workspace "$COMFY_INSTALL_DIR" node install-deps --workflow="$WORKFLOW_FILE" || { log_warning "Algumas dependências do workflow podem ter falhado ao instalar."; log_info "Isso pode ser normal. O ComfyUI mostrará os nós ausentes ao iniciar."; }
    log_info "Instalação de dependências do workflow concluída."
}

# ============================================================================
# COMFYUI LAUNCH
# ============================================================================

launch_comfyui() {
    log_info "Launching ComfyUI on port 8818..."; log_info "ComfyUI will be accessible at: http://0.0.0.0:8818"
    exec comfy --workspace "$COMFY_INSTALL_DIR" launch -- --fast --preview-method latent2rgb --listen 0.0.0.0 --port 8818
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

main() {
    log_info "Starting ComfyUI automated setup..."
    log_info "========================================"
    
    detect_environment
    setup_permissions
    install_system_dependencies
    install_comfy_cli
    install_comfyui_core
    create_model_directories
    install_custom_nodes
    
    # Install special nodes and their dependencies
    install_wd14_tagger_node
    install_story_diffusion_node
    install_nudenet_node # ### NOVO ### Chamada para a nova função
    
    download_all_files
    download_workflow
    #install_workflow_dependencies
    
    log_info "========================================"
    log_info "Setup completed successfully!"
    log_info "========================================"
    launch_comfyui
}

# ============================================================================
# SCRIPT EXECUTION
# ============================================================================

trap 'log_error "Script interrupted"; exit 1' INT TERM

main "$@"