#!/bin/bash
#
# Interactive Model Sync from Google Drive
# Usage: ./sync-models.sh [destination_path]
#

set -e

DEST_PATH="${1:-./models}"
REMOTE_NAME="drive"
REMOTE_PATH="models"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

print_header() {
    echo -e "${CYAN}"
    echo "╔═══════════════════════════════════════════════════════════════╗"
    echo "║           ComfyUI Model Sync from Google Drive                ║"
    echo "╚═══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

print_step() {
    echo -e "${BLUE}▶ $1${NC}"
}

print_success() {
    echo -e "${GREEN}✓ $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠ $1${NC}"
}

print_error() {
    echo -e "${RED}✗ $1${NC}"
}

# Check and install rclone
install_rclone() {
    if command -v rclone &> /dev/null; then
        print_success "rclone is already installed ($(rclone version | head -1))"
        return 0
    fi

    print_step "Installing rclone..."

    if [[ "$OSTYPE" == "darwin"* ]]; then
        if command -v brew &> /dev/null; then
            brew install rclone
        else
            curl https://rclone.org/install.sh | sudo bash
        fi
    elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
        curl https://rclone.org/install.sh | sudo bash
    else
        print_error "Unsupported OS. Install rclone manually: https://rclone.org/install/"
        exit 1
    fi

    print_success "rclone installed"
}

# Check if running headless
is_headless() {
    if [[ -z "$DISPLAY" ]] || [[ -n "$SSH_CLIENT" && -z "$DISPLAY" ]]; then
        return 0
    fi
    return 1
}

# Configure rclone for Google Drive
configure_rclone() {
    local config_file
    config_file=$(rclone config file 2>/dev/null | grep -v "^Configuration" | head -1)
    print_step "Config file: ${config_file:-default location}"

    local existing_remotes
    existing_remotes=$(rclone listremotes 2>/dev/null)

    if [[ -z "$existing_remotes" ]]; then
        echo "  No remotes configured"
    else
        echo "  Found remotes: $(echo $existing_remotes | tr '\n' ' ')"
    fi

    # Check if drive remote exists and works
    if echo "$existing_remotes" | grep -q "^${REMOTE_NAME}:$"; then
        print_step "Testing '${REMOTE_NAME}:' connection..."

        if rclone lsd "${REMOTE_NAME}:" &> /dev/null; then
            print_success "Google Drive connected"
            return 0
        else
            print_warning "Remote exists but connection failed"
            read -p "Delete and reconfigure? [y/N]: " confirm
            if [[ ! "$confirm" =~ ^[Yy] ]]; then
                exit 1
            fi
            rclone config delete "${REMOTE_NAME}" 2>/dev/null || true
        fi
    fi

    print_step "Configuring Google Drive access..."
    echo ""

    if is_headless; then
        echo -e "${CYAN}═══════════════════════════════════════════════════════════════${NC}"
        echo -e "${CYAN}  HEADLESS AUTHENTICATION${NC}"
        echo -e "${CYAN}═══════════════════════════════════════════════════════════════${NC}"
        echo ""
        echo -e "${GREEN}Step 1:${NC} On your LOCAL machine (with browser), run:"
        echo -e "  ${CYAN}rclone authorize \"drive\"${NC}"
        echo ""
        echo -e "${GREEN}Step 2:${NC} Complete Google sign-in in browser"
        echo ""
        echo -e "${GREEN}Step 3:${NC} Copy the token rclone prints"
        echo ""
        echo -e "${CYAN}═══════════════════════════════════════════════════════════════${NC}"
        echo ""
        echo -e "${YELLOW}Now running 'rclone config' - answer 'n' to browser question${NC}"
        echo -e "${YELLOW}and paste your token at the config_token> prompt${NC}"
        echo ""
        read -p "Press Enter when ready..."
        rclone config
    else
        echo -e "${YELLOW}Opening browser for authentication...${NC}"
        rclone config create "${REMOTE_NAME}" drive
    fi

    # Verify
    print_step "Verifying connection..."
    if rclone lsd "${REMOTE_NAME}:" &> /dev/null; then
        print_success "Google Drive configured"
    else
        print_error "Configuration failed"
        exit 1
    fi
}

# Main
main() {
    print_header

    install_rclone
    configure_rclone

    print_step "Fetching models from Google Drive..."

    # Get list into array - simple and reliable
    local models_raw
    models_raw=$(rclone lsf "${REMOTE_NAME}:${REMOTE_PATH}" -R 2>/dev/null | grep -E '\.(safetensors|pth|pt|gguf|onnx|bin)$' | sort)

    if [[ -z "$models_raw" ]]; then
        print_error "No models found in ${REMOTE_NAME}:${REMOTE_PATH}"
        exit 1
    fi

    # Convert to array
    local -a MODELS
    while IFS= read -r line; do
        [[ -n "$line" ]] && MODELS+=("$line")
    done <<< "$models_raw"

    print_success "Found ${#MODELS[@]} models"
    echo ""

    # Print numbered list
    echo -e "${CYAN}Available models:${NC}"
    echo "─────────────────────────────────────────────────────────────────"
    local i=1
    for model in "${MODELS[@]}"; do
        echo -e "  ${YELLOW}${i}${NC}) ${model}"
        i=$((i + 1))
    done
    echo ""
    echo "Enter numbers separated by spaces (e.g., 1 3 5)"
    echo "Or 'all' to download everything, 'q' to quit"
    echo ""
    read -p "> " INPUT

    [[ "$INPUT" = "q" ]] && exit 0

    # Build download list
    local -a TO_DOWNLOAD
    if [[ "$INPUT" = "all" ]]; then
        TO_DOWNLOAD=("${MODELS[@]}")
    else
        for num in $INPUT; do
            local idx=$((num - 1))
            if [[ $idx -ge 0 ]] && [[ $idx -lt ${#MODELS[@]} ]]; then
                TO_DOWNLOAD+=("${MODELS[$idx]}")
            fi
        done
    fi

    if [[ ${#TO_DOWNLOAD[@]} -eq 0 ]]; then
        print_warning "Nothing selected"
        exit 0
    fi

    echo ""
    print_step "Downloading ${#TO_DOWNLOAD[@]} model(s) to ${DEST_PATH}..."
    mkdir -p "$DEST_PATH"

    local current=0
    for model in "${TO_DOWNLOAD[@]}"; do
        current=$((current + 1))
        echo ""
        echo -e "${BLUE}[$current/${#TO_DOWNLOAD[@]}]${NC} $model"
        local dir
        dir=$(dirname "$model")
        mkdir -p "${DEST_PATH}/${dir}"
        rclone copy "${REMOTE_NAME}:${REMOTE_PATH}/${model}" "${DEST_PATH}/${dir}" -P
    done

    echo ""
    print_success "All downloads complete!"
}

main "$@"