#!/bin/bash

_dropshell_completions() {
    local cur prev opts
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"
    root="${COMP_WORDS[1]}"
    
    # List of main commands
    opts="help version servers templates autocomplete_list_servers autocomplete_list_services autocomplete_list_commands install backup"

    # add all commands to opts
    local commands=($(dropshell autocomplete_list_commands))
    opts="${opts} ${commands[*]}"

    # If we're completing the first argument, show all commands
    if [[ ${COMP_CWORD} -eq 1 ]] ; then
        COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
        return 0
    fi

    # Command-specific completions
    case "${root}" in
        templates|autocomplete_list_servers|autocomplete_list_services|autocomplete_list_commands)
            # No additional completions for these commands
            COMPREPLY=()
            return 0
            ;;
        servers)
            # Show available server names for servers command
            local servers=($(dropshell autocomplete_list_servers))
            COMPREPLY=( $(compgen -W "${servers[*]}" -- ${cur}) )
            return 0
            ;;
        install|backup)
            # Handle completion for run/install/backup commands
            if [[ ${COMP_CWORD} -eq 2 ]]; then
                # Second argument is server name
                local servers=($(dropshell autocomplete_list_servers))
                COMPREPLY=( $(compgen -W "${servers[*]}" -- ${cur}) )
                return 0
            elif [[ ${COMP_CWORD} -eq 3 ]]; then
                # Third argument is service name
                local server_name="${COMP_WORDS[2]}"
                local services=($(dropshell autocomplete_list_services "$server_name"))
                COMPREPLY=( $(compgen -W "${services[*]}" -- ${cur}) )
                return 0
            fi
            ;;
        *)
            # if ${prev} is not in ${commands}
            if [[ ! " ${commands[*]} " =~ " ${root} " ]]; then
                COMPREPLY=()
                return 0
            fi
            # Handle completion for template commands
            if [[ ${COMP_CWORD} -eq 2 ]]; then
                # Second argument is server name
                local servers=($(dropshell autocomplete_list_servers))
                COMPREPLY=( $(compgen -W "${servers[*]}" -- ${cur}) )
                return 0
            elif [[ ${COMP_CWORD} -eq 3 ]]; then
                # Third argument is service name
                local server_name="${COMP_WORDS[2]}"
                local services=($(dropshell autocomplete_list_services "$server_name"))
                COMPREPLY=( $(compgen -W "${services[*]}" -- ${cur}) )
                return 0
            fi
            ;;
    esac
}

# Register the completion function
complete -F _dropshell_completions dropshell