#!/bin/bash
# Autocomplete script for dbzeug bash completion. Stand alone and generated at runtime

# Function to check if a string is in an array
string_in_array() {
    local target="$1"
    shift
    local array=("$@")

    for element in "${array[@]}"; do
        if [[ "$element" == "$target" ]]; then
            return 0  # Return success (found)
        fi
    done

    return 1  # Return failure (not found)
}

_dbzeug_completion() {
    local cur_word prev_word prev_args all_args available multi mapped_arg solos need_dir need_str need_template need_file

    local -A arg_map=(
        ["-h"]="--help"        
        ["--help"]="-h"        
        ["-l"]="--list"        
        ["--list"]="-l"        
        ["-o"]="--output"        
        ["--output"]="-o"        
        ["-t"]="--template"        
        ["--template"]="-t"        
        ["-c"]="--console"        
        ["--console"]="-c"        
        ["-i"]="--include"        
        ["--include"]="-i"        
        ["-x"]="--exclude"        
        ["--exclude"]="-x"        
        ["-j"]="--json"        
        ["--json"]="-j"        
        ["-d"]="--debug"        
        ["--debug"]="-d"
    )

    # All the arguments generally available
    all_args=('-h' '--help' '-l' '--list' '-o' '--output' '-t' '--template' '-c' '--console' '--drops' '-i' '--include' '-x' '--exclude' '--integrate' '-j' '--json' '-d' '--debug' '--dbz' '--version' '--validate' '--fix' '--template-dir' '--user-field' '--verbose' '--install' '--bash-completion')
    # Arguments that can be specified multiple times
    multi=('-t' '--template')
    # Arguments that should only be by themselves with no other arguments
    solos=('-h' '--help' '-l' '--list' '--version' '--install' '--bash-completion')
    # Arguments that need directories
    need_dir=('-o' '--output' '--template-dir')
    # Arguments that need a string argument that can't be completed
    need_str=('-i' '--include' '-x' '--exclude' '--dbz' '--user-field')
    # Arguments that need a valid template
    need_template=('-t' '--template')
        
    # Current word we're trying to complete
    cur_word="${COMP_WORDS[COMP_CWORD]}"
    # Previous word we completed, needed for identifing options that need files/ strings/ etc
    prev_word="${COMP_WORDS[COMP_CWORD - 1]}"
    # Arguments excluding the one we're currently working on
    prev_args=("${COMP_WORDS[@]::${#COMP_WORDS[@]}-1}")
    
    # Do we still need a file argument (or can we take one at least)
    need_file=1
    # Arguments available for completion, to be calculated
    available=()
    
    # Sort arguments to make things prettier
    all_args=($(printf "%s\n" "${all_args[@]}" | sort))

    # Block further completion if the argument should be only by itself
    for arg in "${prev_args[@]}"; do
        mapped_arg="${arg_map[$arg]}"
        if string_in_array "${arg}" "${solos[@]}" || string_in_array "${mapped_arg}" "${solos[@]}"; then
            # BAIL! we've go args that should be used alone
            return
        fi
    done

    # Check if we still need a file
    for arg in "${prev_args[@]}"; do
        if [[ "${arg}" == *.sql ]]; then
            need_file=0
            break
        fi
    done

    # Remove existing args that are not able to be used multiple times
    for arg in "${all_args[@]}"; do
        mapped_arg="${arg_map[$arg]}"
        if string_in_array "${arg}" "${multi[@]}" || string_in_array "${mapped_arg}" "${multi[@]}"; then
            available=("${available[@]}" "${arg}")
        elif ! string_in_array "${arg}" "${prev_args[@]}" && ! string_in_array "${mapped_arg}" "${prev_args[@]}"; then
            available=("${available[@]}" "${arg}")
        fi
    done

    if string_in_array "${prev_word}" "${need_str[@]}"; then
        return        # Do nothing if we want a string
    elif string_in_array "${prev_word}" "${need_dir[@]}"; then
        # If we have a directory ending in a slash for the current completion add it so we can get subdirectories too
        if [[ -d "${cur_word}" && "${cur_word}" =~ /$ ]]; then
            COMPREPLY=("${cur_word}" $(compgen -o dirnames -- "${cur_word}" | sed 's/[^ ]* */&\//g'))
        else
            # Just get directories as per normal but add a slash so they're clearly dirs
            COMPREPLY=($(compgen -o dirnames -- "${cur_word}" | sed 's/[^ ]* */&\//g'))

            # Don't just allow the completion of a directory in case there are sub directories, more tabs will show or
            # complete sub directories
            compopt -o nospace
        fi
    elif string_in_array "${prev_word}" "${need_template[@]}"; then
        # Provide templates as a word list but don't treat them as files
        templates=($(dbzeug -l | grep "^ " | tr -d " "))
        COMPREPLY=($(compgen -W "$(IFS=$'\n'; echo "${templates[*]}")" -- "${cur_word}"))
    elif [[ -d "${cur_word}"  && ! "${cur_word}" =~ /$ ]]; then
        # If we have a directory and it doesn't end with a slash, update it to have a slash and use that
        cur_word="${cur_word}/"
        COMP_WORDS[COMP_CWORD]="${cur_word}"
        COMPREPLY=($(compgen -f -- "${cur_word}" | grep -E '\.sql$'))
        compopt -o nospace
    elif [[ "${cur_word}" == -* ]]; then
        # if we're completing an option, don't check for files
        COMPREPLY=($(compgen -W "$(IFS=$'\n'; echo "${available[*]}")" -- "${cur_word}"))
    elif [[ "${need_file}" != 0 && -n "${cur_word}" ]]; then
        # If we still need a file, build the options
        COMPREPLY=($(compgen -o dirnames -- "${cur_word}" | sed 's/[^ ]* */&\//g') $(compgen -f -- "${cur_word}" | grep -E '\.sql$'))
        # If we don't have a file (eg, only a directory) that's not a valid completion so don't add a space
        if [[ ! -f "${cur_word}" ]]; then
            compopt -o nospace
        fi
    else
        # If no specific option is provided, complete with both arguments and files
        if [[ "${need_file}" != 0 ]]; then
            local files=($(compgen -o dirnames -- "${cur_word}" | sed 's/[^ ]* */&\//g') $(compgen -f -- "${cur_word}" | grep -E '\.sql$'))
        fi

        COMPREPLY=($(compgen -W "$(IFS=$'\n'; echo "${available[*]}")" -- "${cur_word}") "${files[@]}")
    fi
}

complete -o nosort -F _dbzeug_completion dbzeug
