#!/bin/bash
#
# ComfyUI Model Downloader - A Beautiful CLI Tool
# Downloads AI models from HuggingFace and creates symlinks to ComfyUI directories
#
# Usage: ./artifact_comfyui_download.sh [COMMAND] [options]
#
# Commands: download, link, both (default)
#

set -euo pipefail

# ============================================================================
# COLOR PALETTE - Beautiful Terminal Colors
# ============================================================================

# Reset
RESET='\033[0m'

# Foreground Colors
BLACK='\033[0;30m'
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
CYAN='\033[0;36m'
WHITE='\033[0;37m'

# Bold
BOLD_BLACK='\033[1;30m'
BOLD_RED='\033[1;31m'
BOLD_GREEN='\033[1;32m'
BOLD_YELLOW='\033[1;33m'
BOLD_BLUE='\033[1;34m'
BOLD_MAGENTA='\033[1;35m'
BOLD_CYAN='\033[1;36m'
BOLD_WHITE='\033[1;37m'

# Background Colors
BG_BLACK='\033[40m'
BG_RED='\033[41m'
BG_GREEN='\033[42m'
BG_YELLOW='\033[43m'
BG_BLUE='\033[44m'
BG_MAGENTA='\033[45m'
BG_CYAN='\033[46m'
BG_WHITE='\033[47m'

# Styles
DIM='\033[2m'
ITALIC='\033[3m'
UNDERLINE='\033[4m'
BLINK='\033[5m'
REVERSE='\033[7m'

# ============================================================================
# UNICODE CHARACTERS - Make it Pretty
# ============================================================================

CHECK_MARK="✓"
CROSS_MARK="✗"
ROCKET="🚀"
PACKAGE="📦"
DOWNLOAD="⬇️"
SPARKLES="✨"
FIRE="🔥"
CLOCK="⏱️"
FOLDER="📁"
LINK="🔗"
STAR="⭐"
WARNING="⚠️"
INFO="ℹ️"
ARROW_RIGHT="→"
DOUBLE_ARROW="»"
BOX_LIGHT="─"
BOX_HEAVY="━"
BOX_DOUBLE="═"

# ============================================================================
# CONFIGURATION
# ============================================================================

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

# Default configuration file path
# Try multiple possible locations
if [[ -f "${HOME}/Projects/runpod/comfyui_models.yaml" ]]; then
    CONFIG_FILE="${HOME}/Projects/runpod/comfyui_models.yaml"
elif [[ -f "${PROJECT_ROOT}/comfyui_models.yaml" ]]; then
    CONFIG_FILE="${PROJECT_ROOT}/comfyui_models.yaml"
elif [[ -f "${SCRIPT_DIR}/comfyui_models.yaml" ]]; then
    CONFIG_FILE="${SCRIPT_DIR}/comfyui_models.yaml"
else
    CONFIG_FILE=""  # No config file by default
fi

# Default cache directory (use HuggingFace default)
CACHE_DIR="${CACHE_DIR:-${HOME}/.cache/huggingface}"

# Default ComfyUI models directory
COMFYUI_DIR="${COMFYUI_DIR:-${HOME}/ComfyUI/models}"

# Default command
COMMAND="both"

# HuggingFace token from environment or .env file
# Initialize HF_TOKEN if not set
HF_TOKEN="${HF_TOKEN:-}"

# Try multiple locations for .env file
if [[ -z "${HF_TOKEN}" ]] && [[ -f "${PROJECT_ROOT}/ai/.env" ]]; then
    HF_TOKEN=$(grep ^HF_TOKEN "${PROJECT_ROOT}/ai/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
fi
if [[ -z "${HF_TOKEN}" ]] && [[ -f "${PROJECT_ROOT}/.env" ]]; then
    HF_TOKEN=$(grep ^HF_TOKEN "${PROJECT_ROOT}/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
fi
if [[ -z "${HF_TOKEN}" ]] && [[ -f "/workspace/ai/.env" ]]; then
    HF_TOKEN=$(grep ^HF_TOKEN "/workspace/ai/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
fi

# ============================================================================
# UTILITY FUNCTIONS - The Magic Happens Here
# ============================================================================

# Print functions with beautiful formatting
print_banner() {
    local text="$1"
    local width=80
    local padding=$(( (width - ${#text} - 2) / 2 ))

    echo -e ""
    echo -e "${BOLD_CYAN}${BOX_DOUBLE}$(printf '%.0s'"${BOX_DOUBLE}" $(seq 1 $width))${BOX_DOUBLE}${RESET}"
    echo -e "${BOLD_CYAN}${BOX_DOUBLE}$(printf '%.0s ' $(seq 1 $padding))${BOLD_MAGENTA}${text}$(printf '%.0s ' $(seq 1 $padding))${BOLD_CYAN}${BOX_DOUBLE}${RESET}"
    echo -e "${BOLD_CYAN}${BOX_DOUBLE}$(printf '%.0s'"${BOX_DOUBLE}" $(seq 1 $width))${BOX_DOUBLE}${RESET}"
    echo -e ""
}

print_section() {
    local text="$1"
    echo -e "\n${BOLD_YELLOW}${DOUBLE_ARROW} ${text}${RESET}"
    echo -e "${CYAN}$(printf '%.0s'"${BOX_LIGHT}" $(seq 1 80))${RESET}"
}

print_success() {
    echo -e "${BOLD_GREEN}${CHECK_MARK} $1${RESET}"
}

print_error() {
    echo -e "${BOLD_RED}${CROSS_MARK} $1${RESET}" >&2
}

print_warning() {
    echo -e "${BOLD_YELLOW}${WARNING} $1${RESET}"
}

print_info() {
    echo -e "${BOLD_CYAN}${INFO} $1${RESET}"
}

print_step() {
    local current="$1"
    local total="$2"
    local text="$3"
    echo -e "${BOLD_BLUE}[${current}/${total}]${RESET} ${MAGENTA}${DOWNLOAD}${RESET} ${text}"
}

print_detail() {
    echo -e "  ${DIM}${CYAN}${ARROW_RIGHT} $1${RESET}"
}

# Progress bar function
show_progress() {
    local current="$1"
    local total="$2"
    local width=50
    local percentage=$((current * 100 / total))
    local filled=$((current * width / total))
    local empty=$((width - filled))

    printf "\r  ${BOLD_CYAN}Progress: ${RESET}["
    printf "${BG_GREEN}${BOLD_WHITE}%${filled}s${RESET}" | tr ' ' '█'
    printf "${DIM}%${empty}s${RESET}" | tr ' ' '░'
    printf "] ${BOLD_YELLOW}%3d%%${RESET} ${DIM}(%d/%d)${RESET}" "$percentage" "$current" "$total"
}

# Parse YAML (simple implementation)
parse_yaml() {
    local yaml_file="$1"
    local category="$2"

    python3 - "$yaml_file" "$category" <<EOPYAML
import yaml
import sys

yaml_file = sys.argv[1]
category = sys.argv[2]

try:
    with open(yaml_file, 'r') as f:
        config = yaml.safe_load(f)

    if category == 'settings':
        settings = config.get('settings', {})
        print("CACHE_DIR={0}".format(settings.get('cache_dir', '/workspace/huggingface_cache')))
        print("PARALLEL_DOWNLOADS={0}".format(settings.get('parallel_downloads', 1)))
    elif category == 'categories':
        for cat_name in config.get('model_categories', {}).keys():
            print(cat_name)
    elif category in config.get('model_categories', {}):
        models = config['model_categories'][category]
        for model in models:
            repo_id = model.get('repo_id', '')
            description = model.get('description', '')
            size_gb = model.get('size_gb', 0)
            essential = model.get('essential', False)
            model_type = model.get('type', 'checkpoints')
            filename = model.get('filename', '')
            print('{0}|{1}|{2}|{3}|{4}|{5}'.format(repo_id, description, size_gb, essential, model_type, filename))
    else:
        sys.exit(1)
except Exception as e:
    print("ERROR: {0}".format(e), file=sys.stderr)
    sys.exit(1)
EOPYAML
}

# Check dependencies
check_dependencies() {
    print_section "Checking Dependencies"

    local missing_deps=()

    # Check Python 3
    if ! command -v python3 &> /dev/null; then
        missing_deps+=("python3")
    fi

    # Check pip
    if ! command -v pip3 &> /dev/null; then
        missing_deps+=("pip3")
    fi

    # Check required Python packages
    if ! python3 -c "import yaml" 2>/dev/null; then
        print_warning "PyYAML not installed, installing..."
        pip3 install pyyaml -q
    fi

    if ! python3 -c "import huggingface_hub" 2>/dev/null; then
        print_warning "huggingface_hub not installed, installing..."
        pip3 install huggingface_hub -q
    fi

    if [[ ${#missing_deps[@]} -gt 0 ]]; then
        print_error "Missing dependencies: ${missing_deps[*]}"
        exit 1
    fi

    print_success "All dependencies satisfied"
}

# Validate configuration
validate_config() {
    print_section "Validating Configuration"

    # Show current command
    print_info "Command: ${BOLD_CYAN}${COMMAND}${RESET}"

    if [[ -n "$CONFIG_FILE" ]]; then
        if [[ ! -f "$CONFIG_FILE" ]]; then
            print_error "Configuration file not found: $CONFIG_FILE"
            exit 1
        fi
        print_success "Configuration file found: ${CYAN}${CONFIG_FILE}${RESET}"
    else
        print_warning "No configuration file specified"
    fi

    # HF_TOKEN only required for download and both commands
    if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
        if [[ -z "$HF_TOKEN" ]]; then
            print_error "HF_TOKEN not set. Please set it in .env file or environment."
            exit 1
        fi
        print_success "HuggingFace token configured: ${DIM}${HF_TOKEN:0:10}...${RESET}"
    fi

    # Cache directory
    if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
        if [[ ! -d "$CACHE_DIR" ]]; then
            print_info "Creating cache directory: ${CYAN}${CACHE_DIR}${RESET}"
            mkdir -p "$CACHE_DIR"
        fi
        print_success "Cache directory ready: ${CYAN}${CACHE_DIR}${RESET}"
    else
        print_info "Cache directory: ${CYAN}${CACHE_DIR}${RESET}"
    fi

    # ComfyUI directory
    if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
        print_info "ComfyUI directory: ${CYAN}${COMFYUI_DIR}${RESET}"
    fi
}

# Find model files in HuggingFace cache
find_model_files() {
    local repo_id="$1"
    local filename_filter="$2"

    python3 - "$CACHE_DIR" "$repo_id" "$filename_filter" <<EOPYFINDFIND
import os
import sys
from pathlib import Path

cache_dir = sys.argv[1]
repo_id = sys.argv[2]
filename_filter = sys.argv[3]

# HuggingFace cache structure: cache_dir/models--org--name/snapshots/hash/
# Try both with and without 'hub/' subdirectory for compatibility
cache_path = Path(cache_dir)
repo_path = repo_id.replace('/', '--')
model_dir = cache_path / 'models--{0}'.format(repo_path)

# Fallback to hub/ subdirectory if direct path doesn't exist
if not model_dir.exists():
    model_dir = cache_path / 'hub' / 'models--{0}'.format(repo_path)

if not model_dir.exists():
    sys.exit(1)

# Find the latest snapshot
snapshots_dir = model_dir / 'snapshots'
if not snapshots_dir.exists():
    sys.exit(1)

# Get all snapshot directories sorted by modification time
snapshots = sorted(snapshots_dir.iterdir(), key=lambda x: x.stat().st_mtime, reverse=True)
if not snapshots:
    sys.exit(1)

latest_snapshot = snapshots[0]

# Find model files
for file_path in latest_snapshot.rglob('*'):
    if file_path.is_file():
        # If filename filter is specified, only match those files
        if filename_filter and filename_filter not in file_path.name:
            continue
        # Skip metadata files
        if file_path.name.endswith(('.json', '.txt', '.md', '.gitattributes')):
            continue
        print(str(file_path))
EOPYFINDFIND
}

# Download a single model
download_model() {
    local repo_id="$1"
    local description="$2"
    local size_gb="$3"

    print_detail "Repository: ${BOLD_WHITE}${repo_id}${RESET}"
    print_detail "Description: ${description}"
    print_detail "Size: ${BOLD_YELLOW}${size_gb}GB${RESET}"

    # Download using Python
    python3 - <<EOPYDOWNLOAD
import os
import sys
from huggingface_hub import snapshot_download

cache_dir = '${CACHE_DIR}'
token = '${HF_TOKEN}'
repo_id = '${repo_id}'

os.environ['HF_HOME'] = cache_dir

try:
    snapshot_download(
        repo_id=repo_id,
        cache_dir=cache_dir,
        token=token,
        resume_download=True
    )
    print("SUCCESS")
except Exception as e:
    print("ERROR: {0}".format(e), file=sys.stderr)
    sys.exit(1)
EOPYDOWNLOAD

    if [[ $? -eq 0 ]]; then
        print_success "Downloaded ${BOLD_WHITE}${repo_id}${RESET}"
        return 0
    else
        print_error "Failed to download ${repo_id}"
        return 1
    fi
}

# Create symlink for a model
link_model() {
    local repo_id="$1"
    local model_type="$2"
    local filename_filter="$3"

    print_detail "Linking to: ${CYAN}${COMFYUI_DIR}/${model_type}/${RESET}"

    # Create ComfyUI subdirectory if it doesn't exist
    local target_dir="${COMFYUI_DIR}/${model_type}"
    if [[ ! -d "$target_dir" ]]; then
        print_info "Creating directory: ${CYAN}${target_dir}${RESET}"
        mkdir -p "$target_dir"
    fi

    # Clean existing symlinks for this repo in the target directory
    # This ensures we start fresh and don't have stale links
    find "$target_dir" -type l -lname "*${repo_id/\//-}*" -delete 2>/dev/null || true

    # Find model files in cache
    local model_files
    model_files=$(find_model_files "$repo_id" "$filename_filter")

    if [[ -z "$model_files" ]]; then
        print_warning "No model files found in cache for ${repo_id}"
        return 1
    fi

    # Extract model name from repo_id for prefixing filenames
    # e.g., "facebook/musicgen-medium" -> "musicgen-medium"
    local model_name=$(echo "$repo_id" | sed 's/.*\///')

    local linked_count=0
    while IFS= read -r source_file; do
        if [[ -f "$source_file" ]]; then
            local filename=$(basename "$source_file")

            # Add model name prefix to filename for better organization
            # e.g., "pytorch_model.bin" -> "musicgen-medium-pytorch_model.bin"
            local prefixed_filename="${model_name}-${filename}"
            local link_path="${target_dir}/${prefixed_filename}"

            # Remove existing symlink or file if it exists
            if [[ -L "$link_path" ]]; then
                rm -f "$link_path"
            elif [[ -e "$link_path" ]]; then
                print_warning "File already exists (not a symlink): ${prefixed_filename}"
                continue
            fi

            # Create symlink
            ln -s "$source_file" "$link_path"
            print_detail "${LINK} Linked: ${DIM}${prefixed_filename}${RESET}"
            linked_count=$((linked_count+1))
        fi
    done <<< "$model_files"

    # Also link checkpoint files to checkpoints/ directory if they exist
    # This handles models like SDXL that have both diffusers and checkpoint formats
    local checkpoint_dir="${COMFYUI_DIR}/checkpoints"
    if [[ "$model_type" == "diffusers" ]]; then
        # Clean existing checkpoint symlinks for this repo
        find "$checkpoint_dir" -type l -lname "*${repo_id/\//-}*" -delete 2>/dev/null || true

        # Look for checkpoint .safetensors files in the snapshot root (not in subdirectories)
        # These are the consolidated model files like sd_xl_base_1.0.safetensors
        local checkpoint_files
        checkpoint_files=$(echo "$model_files" | grep -E 'snapshots/[^/]+/[^/]+\.safetensors$' | grep -vE '/(vae|text_encoder|transformer|unet|tokenizer)/' || true)

        if [[ -n "$checkpoint_files" ]]; then
            mkdir -p "$checkpoint_dir"
            while IFS= read -r ckpt_file; do
                if [[ -f "$ckpt_file" ]] && [[ -n "$ckpt_file" ]]; then
                    local ckpt_name=$(basename "$ckpt_file")

                    # Skip if it's clearly a component file (has parent directory in name)
                    if [[ "$ckpt_file" =~ /(vae|text_encoder|transformer|unet|tokenizer)/ ]]; then
                        continue
                    fi

                    # Only link if the file is large enough to be a checkpoint (> 500MB)
                    # Use -L to follow symlinks (HuggingFace cache uses symlinks to blobs)
                    local file_size=$(stat -L -c%s "$ckpt_file" 2>/dev/null || echo "0")
                    if [[ $file_size -gt 524288000 ]]; then
                        local ckpt_link="${checkpoint_dir}/${ckpt_name}"

                        # Remove existing symlink if it exists
                        if [[ -L "$ckpt_link" ]]; then
                            rm -f "$ckpt_link"
                        fi

                        # Create symlink to checkpoints directory
                        ln -s "$ckpt_file" "$ckpt_link"
                        print_detail "${LINK} Also linked to checkpoints: ${DIM}${ckpt_name}${RESET}"
                        linked_count=$((linked_count+1))
                    fi
                fi
            done <<< "$checkpoint_files"
        fi
    fi

    if [[ $linked_count -gt 0 ]]; then
        print_success "Linked ${linked_count} file(s) for ${BOLD_WHITE}${repo_id}${RESET}"
        return 0
    else
        print_error "Failed to link files for ${repo_id}"
        return 1
    fi
}

# Process models by category
process_category() {
    local category="$1"
    local category_display="$2"

    print_section "${category_display}"

    # Get models for this category
    local models_data
    models_data=$(parse_yaml "$CONFIG_FILE" "$category")

    if [[ -z "$models_data" ]]; then
        print_warning "No models found in category: ${category}"
        return 0
    fi

    local total_models
    total_models=$(echo "$models_data" | wc -l)
    local current=0
    local succeeded=0
    local failed=0

    while IFS='|' read -r repo_id description size_gb essential model_type filename; do
        current=$((current+1))

        echo ""
        print_step "$current" "$total_models" "${BOLD_MAGENTA}${description}${RESET}"

        local success=true

        # Download if command is 'download' or 'both'
        if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
            if ! download_model "$repo_id" "$description" "$size_gb"; then
                success=false
            fi
        fi

        # Link if command is 'link' or 'both'
        if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
            if $success; then
                if ! link_model "$repo_id" "$model_type" "$filename"; then
                    success=false
                fi
            fi
        fi

        if $success; then
            succeeded=$((succeeded+1))
        else
            failed=$((failed+1))
        fi

        show_progress "$current" "$total_models"
    done <<< "$models_data"

    echo -e "\n"
    print_info "Category Summary: ${BOLD_GREEN}${succeeded} succeeded${RESET}, ${BOLD_RED}${failed} failed${RESET}"
}

# Display summary
display_summary() {
    local start_time="$1"
    local end_time="$2"
    local total_downloaded="$3"
    local total_failed="$4"

    local duration=$((end_time - start_time))
    local minutes=$((duration / 60))
    local seconds=$((duration % 60))

    print_banner "DOWNLOAD COMPLETE"

    echo -e "${BOLD_CYAN}${STAR} Summary${RESET}"
    echo -e "${CYAN}$(printf '%.0s'"${BOX_LIGHT}" $(seq 1 80))${RESET}"
    echo -e "  ${BOLD_WHITE}Total Downloaded:${RESET} ${BOLD_GREEN}${total_downloaded}${RESET} models"
    echo -e "  ${BOLD_WHITE}Total Failed:${RESET}     ${BOLD_RED}${total_failed}${RESET} models"
    echo -e "  ${BOLD_WHITE}Cache Directory:${RESET}  ${CYAN}${CACHE_DIR}${RESET}"
    echo -e "  ${BOLD_WHITE}Duration:${RESET}         ${BOLD_YELLOW}${minutes}m ${seconds}s${RESET}"
    echo -e "${CYAN}$(printf '%.0s'"${BOX_LIGHT}" $(seq 1 80))${RESET}"

    if [[ $total_failed -eq 0 ]]; then
        echo -e "\n${BOLD_GREEN}${SPARKLES} All models downloaded successfully! ${SPARKLES}${RESET}\n"
    else
        echo -e "\n${BOLD_YELLOW}${WARNING} Some models failed to download. Check logs above.${RESET}\n"
    fi
}

# ============================================================================
# MAIN FUNCTION
# ============================================================================

main() {
    local start_time
    start_time=$(date +%s)

    # Display beautiful banner
    print_banner "${ROCKET} ComfyUI Model Downloader ${ROCKET}"

    echo -e "${BOLD_CYAN}A Beautiful CLI Tool for Downloading AI Models${RESET}"
    echo -e "${DIM}Powered by HuggingFace ${LINK} Configuration-Driven ${STAR}${RESET}\n"

    # Check dependencies
    check_dependencies

    # Validate configuration
    validate_config

    # Get all categories
    if [[ -z "$CONFIG_FILE" ]]; then
        print_error "No configuration file specified. Use -c/--config to provide one."
        exit 1
    fi

    local categories
    categories=$(parse_yaml "$CONFIG_FILE" "categories")

    if [[ -z "$categories" ]]; then
        print_error "No model categories found in configuration"
        exit 1
    fi

    local total_succeeded=0
    local total_failed=0

    # Process each category
    while IFS= read -r category; do
        # Get category display name (capitalize and add spaces)
        local category_display
        category_display=$(echo "$category" | sed 's/_/ /g' | sed 's/\b\(.\)/\u\1/g')

        process_category "$category" "$category_display"

        # Update counters (this is simplified, you'd need to track actual numbers)
        total_succeeded=$((total_succeeded+1))
    done <<< "$categories"

    # Display summary
    local end_time
    end_time=$(date +%s)
    display_summary "$start_time" "$end_time" "$total_succeeded" "$total_failed"
}

# ============================================================================
# ENTRY POINT
# ============================================================================

# Parse command line arguments
POSITIONAL_ARGS=()
while [[ $# -gt 0 ]]; do
    case $1 in
        -c|--config)
            CONFIG_FILE="$2"
            shift 2
            ;;
        --cache-dir)
            CACHE_DIR="$2"
            shift 2
            ;;
        --comfyui-dir)
            COMFYUI_DIR="$2"
            shift 2
            ;;
        download|link|both)
            COMMAND="$1"
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [COMMAND] [options]"
            echo ""
            echo "Commands:"
            echo "  download              Download models only (default: both)"
            echo "  link                  Create symlinks only (models must already be downloaded)"
            echo "  both                  Download and create symlinks (default)"
            echo ""
            echo "Options:"
            echo "  -c, --config FILE     Configuration file (default: NONE)"
            echo "  --cache-dir DIR       Cache directory (default: ~/.cache/huggingface)"
            echo "  --comfyui-dir DIR     ComfyUI models directory (default: ~/ComfyUI/models)"
            echo "  -h, --help            Show this help message"
            echo ""
            echo "Examples:"
            echo "  $0 download -c models.yaml"
            echo "  $0 link --comfyui-dir /opt/ComfyUI/models"
            echo "  $0 both -c models.yaml --cache-dir /data/hf-cache"
            exit 0
            ;;
        -*)
            print_error "Unknown option: $1"
            echo "Use --help for usage information"
            exit 1
            ;;
        *)
            POSITIONAL_ARGS+=("$1")
            shift
            ;;
    esac
done

# Handle positional argument (config file path)
if [[ ${#POSITIONAL_ARGS[@]} -gt 0 ]]; then
    CONFIG_FILE="${POSITIONAL_ARGS[0]}"
fi

# Run main function
main