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

_docta_completion() {
    # 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)
    }

    debug() {
        # Set COMP_DEBUG=1 if you want debug messages to be logged 
        [[ -n $COMP_DEBUG ]] && echo "$*" >> /tmp/completion_debug.log
    }

    debug "## completion called for '""${COMP_WORDS[@]}""' ##"

    # Set local to avoid namespace pollution
    local cur_word prev_word prev_arg prev_args all_args available multi mapped_arg solos 
    local need_file need_dir need_str need_template need_file_or_dir_or_str
    local files dirs ext need_file_arg need_dir_arg

    local -A arg_map=(
        ["-h"]="--help"        
        ["--help"]="-h"        
        ["-v"]="--version"        
        ["--version"]="-v"        
        ["-q"]="--quiet"        
        ["--quiet"]="-q"        
        ["-d"]="--debug"        
        ["--debug"]="-d"
    )

    local -A ext_map=(
        ["--pdf"]="pdf"        
        ["--html"]="html"
    )

    ext=""

    # All the arguments generally available
    all_args=('-h' '--help' '-v' '--version' '--install' '--bash-completion' '--init' '--check' '--make-config' '-q' '--quiet' '-d' '--debug' '--pdf' '--html')
    # Arguments that can be specified multiple times
    multi=()
    # Arguments that should only be by themselves with no other arguments
    solos=('-h' '--help' '-v' '--version' '--install' '--bash-completion' '--init' '--check' '--make-config')
    # Arguments that need files
    need_file=('--pdf' '--html')
    # Arguments that need directories
    need_dir=()
    # Arguments that need a string argument that can't be completed
    need_str=()
    # Arguments that need a valid template
    need_template=()

    # Current word we're trying to complete
    cur_word="${COMP_WORDS[COMP_CWORD]}"
    # Previous word we completed, needed for identifying options that need files/ strings/ etc
    prev_word="${COMP_WORDS[COMP_CWORD - 1]}"
    # Arguments excluding the one we're currently working on and the name of the program
    if (( COMP_CWORD > 1 )); then
        prev_args=("${COMP_WORDS[@]:1:COMP_CWORD-1}")
    else
        prev_args=()
    fi

    # Do we still need a file argument (or can we take one at least)
    need_file_arg=0    # -1 = unlimited
    # Can we take a directory arg
    need_dir_arg=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
            debug "Found solo option ${arg} so aborting completion"            
            return
        fi
    done

    prev_arg=""

    need_value=("${need_file}" "${need_dir}" "${need_str}")

    for arg in "${prev_args[@]}"; do
        # If we're looking for a positional argument and this looks like it might be a file/ dir
        if { (( need_dir_arg > 0 )) || (( need_file_arg > 0 )) } && [[ ${arg} != -* ]]; then
            # If the previous arg wasn't something that needed a file/dir/str
            if ! string_in_array "${prev_arg}" "${need_value[@]}" || [[ "${prev_arg}" != -* ]]; then
                debug "Don't need a dir or file as we have ${arg}"
                # Cancel out looking for files or dirs
                (( need_dir_arg = need_dir_arg > 0 ? need_dir_arg - 1 : 0 ))
                (( need_file_arg = need_file_arg > 0 ? need_file_arg - 1 : 0 ))
                debug "need_dir_arg = ${need_dir_arg}, need_file_arg = ${need_file_arg}"
            else 
                debug "${prev_arg} is a file param (${COMP_CWORD} words)"
            fi
        fi
        prev_arg="${arg}"
    done

    # Remove existing args that are not able to be used multiple times or that are solo if we already have args
    for arg in "${all_args[@]}"; do
        mapped_arg="${arg_map[$arg]}"
        if string_in_array "${arg}" "${multi[@]}" || string_in_array "${mapped_arg}" "${multi[@]}"; then
            debug "${arg} (or ${mapped_arg}) is in multi so can use again"
            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
        debug "We need a string for ${prev_word} so not offering completion"
        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_file[@]}" && [[ "${cur_word}" != -* ]]; then
        # Our previous arugment needs a file and our current argument doesn't look like a parameter
        ext="${ext_map[${prev_word}]}"
        if [[ -n $ext ]]; then
            debug "Argument ${prev_word} needs a file so offering ${ext} files"
            mapfile -t files < <(compgen -f -- "${cur_word}" | grep -E "\.${ext}$")
        else
            debug "Argument ${prev_word} needs a file but don't know the extention'"
            mapfile -t files < <(compgen -f -- "${cur_word}")
        fi

        mapfile -t dirs < <(compgen -o dirnames -- "$cur_word" | sed 's/[^ ]* */&\//g')

        COMPREPLY=( "${dirs[@]}" "${files[@]}" )
        # 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        
    elif string_in_array "${prev_word}" "${need_template[@]}"; then
        # Provide templates as a word list but don't treat them as files
        templates=($(docta -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_arg}" != 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}" ))
        # 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 (( "${need_dir_arg}" != 0 )); then
            dirs=($(compgen -o dirnames -- "${cur_word}" | sed 's/[^ ]* */&\//g'))
        fi

        if (( "${need_file_arg}" != 0 )); then
            if [[ -z "${ext}" ]]; then 
                debug "Need file arg but don't have a filter"
                files=($(compgen -o dirnames -- "${cur_word}" | sed 's/[^ ]* */&\//g') $(compgen -f -- "${cur_word}"))
            else
                debug "Need file and have a pattern to use ${ext}"
                files=($(compgen -o dirnames -- "${cur_word}" | sed 's/[^ ]* */&\//g') $(compgen -f -- "${cur_word}" | grep -E "\.(${ext})$"))
            fi
        fi

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

complete -o nosort -F _docta_completion docta
