#!/bin/bash

# ComfyUI Automated Setup Script
# Supports LightningAI, RunPod, VastAI, and generic Linux environments
# Author: Generated for automated ComfyUI deployment

set -e
set -o pipefail

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

# Civitai API Token
readonly CIVITAI_TOKEN="0b698de03658fb8df2eeef90bad3ab9a"

# Custom Nodes to Install
readonly CUSTOM_NODES=(
    "https://github.com/Cezarsaint/custom_prompt_manager"
    "https://github.com/adbrasi/pakreatorio"
)

# Files to Download - Format: "URL|target_subdir|filename"
# If filename is empty, let aria2c/wget determine from Content-Disposition
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/1268294?type=Model&format=SafeTensor|loras|"
    "https://civitai.com/api/download/models/1715330?type=Model&format=SafeTensor|loras|"
    "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/snapshotBRABEX.json|INSTALL_DIR|snapshotBRABEX.json"
)

# Model directories to create
readonly MODEL_DIRS=(
    "checkpoints"
    "loras"
    "upscale_models"
    "ultralytics/bbox"
)

# Global variables
BASE_DIR=""
COMFY_INSTALL_DIR=""
MODELS_DIR=""
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"
    
    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..."
    
    # Update package list
    if command_exists apt-get; then
        ${SUDO_PREFIX}apt-get update -qq || log_warning "Failed to update package list"
        
        # Install basic dependencies
        ${SUDO_PREFIX}apt-get install -y -qq python3 python3-pip git || {
            log_error "Failed to install basic dependencies (python3, pip, git)"
            exit 1
        }
        
        # Try to install aria2c
        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"
        # Check if aria2c is already available
        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
    }
    
    # Check if comfy is in PATH
    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 || {
        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..."
    
    # Create model directories
    for dir in "${MODEL_DIRS[@]}"; do
        mkdir -p "$MODELS_DIR/$dir"
        log_info "Created directory: $MODELS_DIR/$dir"
    done
    
    # Ensure install directory exists for snapshot
    mkdir -p "$COMFY_INSTALL_DIR"
}

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

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

# ============================================================================
# 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"
    
    # Add Civitai token if needed
    url=$(add_civitai_token "$url")
    
    local download_cmd=""
    local success=false
    
    # Try aria2c first
    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
    
    # Try wget if aria2c failed or unavailable
    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
    
    # Try curl as last resort
    if ! $success && command_exists curl; then
        log_warning "Using curl fallback for $url - filename handling may be less robust"
        
        if [ -n "$filename" ]; then
            download_cmd="curl -L -C - -o \"$target_dir/$filename\" \"$url\""
        else
            # For curl with content-disposition, it's more complex
            # Try -J -O with --output-dir (if available) or fallback to simpler approach
            if curl --help 2>/dev/null | grep -q "output-dir"; then
                download_cmd="curl -L -J -O --output-dir \"$target_dir\" \"$url\""
            else
                # Fallback: extract filename from URL or use generic name
                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
        if [ "$target_subdir" = "INSTALL_DIR" ]; then
            target_dir="$COMFY_INSTALL_DIR"
        else
            target_dir="$MODELS_DIR/$target_subdir"
        fi
        
        # Note: Directory creation is handled by download tools automatically
        # No need to mkdir -p here as it can cause issues
        
        local display_name="${filename:-$(basename "$url" | cut -d'?' -f1)}"
        log_info "Starting download: $display_name -> $target_dir"
        
        # Create a temporary log file for this download
        local log_file
        log_file=$(mktemp)
        download_logs+=("$log_file")
        
        # Start download in background
        (
            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
    
    # Wait for all downloads to complete
    local failed_count=0
    for i in "${!pids[@]}"; do
        local pid=${pids[$i]}
        local log_file=${download_logs[$i]}
        
        if wait "$pid"; then
            log_info "Download completed successfully"
        else
            failed_count=$((failed_count + 1))
            log_warning "Download failed - check logs"
        fi
        
        # Show download log
        if [ -f "$log_file" ]; then
            cat "$log_file" >&2
            rm -f "$log_file"
        fi
    done
    
    if [ $failed_count -gt 0 ]; then
        log_warning "$failed_count downloads failed, but continuing..."
    else
        log_info "All downloads completed successfully"
    fi
}

# ============================================================================
# SNAPSHOT RESTORATION
# ============================================================================

restore_snapshot() {
    log_info "Restoring snapshot..."
    
    local snapshot_file="$COMFY_INSTALL_DIR/snapshotBRABEX.json"
    
    if [ -f "$snapshot_file" ]; then
        if comfy --workspace "$COMFY_INSTALL_DIR" node restore-snapshot "$snapshot_file"; then
            log_info "Snapshot restored successfully"
        else
            log_warning "Failed to restore snapshot, ComfyUI will start without it"
        fi
    else
        log_warning "Snapshot file not found: $snapshot_file"
    fi
}

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


launch_comfyui() {
    log_info "Launching ComfyUI on port 8818..."
    log_info "ComfyUI will be accessible at: http://0.0.0.0:8818"
    
    # Launch ComfyUI

    exec comfy --workspace "$COMFY_INSTALL_DIR" launch -- --fast --listen 0.0.0.0 --port 8818
}

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

main() {
    log_info "Starting ComfyUI automated setup..."
    log_info "========================================"
    
    # Step 1: Environment Detection
    detect_environment
    
    # Step 2: Permission Setup
    setup_permissions
    
    # Step 3: System Dependencies
    install_system_dependencies
    
    # Step 4: Install comfy-cli
    install_comfy_cli
    
    # Step 5: Install ComfyUI Core
    install_comfyui_core
    
    # Step 6: Create Directory Structure
    create_model_directories
    
    # Step 7: Install Custom Nodes
    install_custom_nodes
    
    # Step 8: Download Files in Parallel
    download_all_files
    
    # Step 9: Restore Snapshot
    restore_snapshot
    
    # Step 10: Launch ComfyUI
    log_info "========================================"
    log_info "Setup completed successfully!"
    log_info "========================================"
    launch_comfyui
}

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

# Trap to handle script interruption
trap 'log_error "Script interrupted"; exit 1' INT TERM

# Run main function
main "$@"