#!/bin/bash

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

    # 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 "${prev}" in
        status)
            # No additional completions for status
            COMPREPLY=()
            ;;
        servers)
            # Use the new autocomplete_list_servers command
            local servers=($(dropshell autocomplete_list_servers))
            COMPREPLY=( $(compgen -W "${servers[*]}" -- ${cur}) )
            return 0
            ;;
        templates)
            # No additional completions for templates
            COMPREPLY=()
            ;;
        autocomplete_list_services)
            # Use the new autocomplete_list_servers command for server names
            local servers=($(dropshell autocomplete_list_servers))
            COMPREPLY=( $(compgen -W "${servers[*]}" -- ${cur}) )
            return 0
            ;;
        autocomplete_list_commands)
            # No additional completions needed
            COMPREPLY=()
            return 0
            ;;
        run|install|backup)
            # First argument after run/install/backup is server name
            local servers=($(dropshell autocomplete_list_servers))
            COMPREPLY=( $(compgen -W "${servers[*]}" -- ${cur}) )
            return 0
            ;;
        *)
            # Handle completion for service names and commands after run/install/backup
            if [[ ${COMP_CWORD} -ge 2 ]] && [[ "${COMP_WORDS[1]}" == "run" || "${COMP_WORDS[1]}" == "install" || "${COMP_WORDS[1]}" == "backup" ]]; then
                if [[ ${COMP_CWORD} -eq 3 ]]; then
                    # Second 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
                elif [[ ${COMP_CWORD} -eq 4 && "${COMP_WORDS[1]}" == "run" ]]; then
                    # Third argument is command name (only for run command)
                    # For now, we'll just complete with common commands
                    local common_commands="status start stop update backup"
                    COMPREPLY=( $(compgen -W "${common_commands}" -- ${cur}) )
                    return 0
                fi
            fi
            ;;
    esac
}

# Register the completion function
complete -F _dropshell_completions dropshell