#!/bin/zsh

# Ensure nvm is loaded
[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"

# Get the script name reliably
SCRIPT_NAME=$(basename "$0")

# Function to display help
show_help() {
    echo "Usage: $SCRIPT_NAME [OPTIONS]"
    echo
    echo "Manage Node.js versions installed via nvm."
    echo
    echo "Options:"
    echo "  --help          Show this help message and exit"
    echo "  --dry-run       Run in dry-run mode (no changes will be made)"
    echo "  -n, --number N  Keep the latest version of the N most recent major versions (default: 4)"
    echo "  --update        Check for and prompt to install updates for kept versions"
    echo "  -f, --force     Skip all confirmation prompts"
    echo
    echo "Examples:"
    echo "  $SCRIPT_NAME --dry-run            # Dry run with default settings"
    echo "  $SCRIPT_NAME -n 3                 # Keep only the 3 most recent major versions"
    echo "  $SCRIPT_NAME --update             # Update kept versions and remove old ones"
    echo "  $SCRIPT_NAME --update --dry-run   # Check for updates in dry-run mode"
    echo "  $SCRIPT_NAME --force              # Run without any confirmation prompts"
}

# ... (rest of the script remains the same)
# Default values
dry_run=false
versions_to_keep=4
update_versions=false
force=false

# Parse command line arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        --help) show_help; exit 0 ;;
        --dry-run) dry_run=true ;;
        -n|--number) versions_to_keep="$2"; shift ;;
        --update) update_versions=true ;;
        -f|--force) force=true ;;
        *) echo "Unknown parameter passed: $1"; show_help; exit 1 ;;
    esac
    shift
done

if $dry_run; then
    echo "Running in dry-run mode. No versions will be removed or installed."
fi
echo "Keeping the latest version of the $versions_to_keep most recent major versions."

# Get all installed versions
installed_versions=($(nvm ls --no-colors --no-alias | grep -oE "v[0-9]+\.[0-9]+\.[0-9]+" | sort -V))

# Get unique major versions
major_versions=($(echo ${installed_versions[@]} | tr ' ' '\n' | cut -d'.' -f1 | sort -u -V | tail -n $versions_to_keep))

# Function to get the latest version for a given major version
get_latest_version() {
    local major_version=$1
    echo ${installed_versions[@]} | tr ' ' '\n' | grep "^$major_version" | sort -V | tail -n 1
}

# Get the latest version for each of the most recent major versions
versions_to_keep_array=()
for major_version in ${major_versions[@]}; do
    versions_to_keep_array+=($(get_latest_version $major_version))
done

# Function to get the latest available version for a given major version
get_latest_available_version() {
    local major_version=$1
    nvm ls-remote --no-colors | grep -oE "$major_version\.[0-9]+\.[0-9]+" | sort -V | tail -n 1
}

# Update function
update_versions() {
    echo "Checking for updates..."
    for major_version in ${major_versions[@]}; do
        local current_version=$(get_latest_version $major_version)
        local latest_version=$(get_latest_available_version $major_version)
        if [[ "$current_version" != "$latest_version" ]]; then
            echo "Update available for $major_version: $current_version -> $latest_version"
            if ! $dry_run; then
                if $force; then
                    echo "Installing Node.js version $latest_version"
                    nvm install $latest_version
                    nvm use $latest_version
                else
                    read "response?Do you want to install $latest_version? (y/N) "
                    if [[ "$response" =~ ^[Yy]$ ]]; then
                        echo "Installing Node.js version $latest_version"
                        nvm install $latest_version
                        nvm use $latest_version
                    else
                        echo "Skipping installation of $latest_version"
                    fi
                fi
            else
                echo "Would install Node.js version $latest_version"
            fi
        else
            echo "$major_version is already up to date ($current_version)"
        fi
    done
}

# Removal function
remove_old_versions() {
    for version in ${installed_versions[@]}; do
        if [[ ! " ${versions_to_keep_array[@]} " =~ " ${version} " ]]; then
            if $dry_run; then
                echo "Would remove Node.js version $version"
            else
                echo "Removing Node.js version $version"
                nvm uninstall $version
            fi
        else
            echo "Keeping Node.js version $version"
        fi
    done
}

# Main execution
if $update_versions; then
    update_versions
fi

if ! $force; then
    echo "Proceed with removing old versions?"
    if ! $dry_run; then
        read "response?Do you want to continue? (y/N) "
        if [[ ! "$response" =~ ^[Yy]$ ]]; then
            echo "Operation cancelled."
            exit 0
        fi
    fi
fi

remove_old_versions

if $dry_run; then
    echo "Dry run complete. To perform the actual operations, run the script without the --dry-run option."
else
    echo "Operations complete"
fi