#!/usr/bin/env bash

declare -gA _samba_tool_help_cache=()
declare -gA _samba_tool_subcommand_cache=()
declare -gA _samba_tool_option_cache=()
declare -gA _samba_tool_option_spec_cache=()
declare -ga _samba_tool_command_path=()

# Cache keys are built from the command path: "user", "user setpassword", etc.
_samba_tool_cache_key()
{
    local IFS=$'\037'

    if [[ $# -eq 0 ]]; then
        printf '%s' __root__
        return
    fi

    printf '%s' "$*"
}

_samba_tool_help()
{
    local key

    key=$(_samba_tool_cache_key "$@")
    _samba_tool_ensure_help "$@"

    printf '%s\n' "${_samba_tool_help_cache[$key]}"
}

_samba_tool_ensure_help()
{
    local key

    key=$(_samba_tool_cache_key "$@")
    if [[ ! -v _samba_tool_help_cache[$key] ]]; then
        _samba_tool_help_cache[$key]=$(samba-tool "$@" --help 2>&1)
    fi
}

# Parse the "Available subcommands:" block printed by samba-tool help.
_samba_tool_ensure_subcommands()
{
    local key line in_block=0

    key=$(_samba_tool_cache_key "$@")
    if [[ ! -v _samba_tool_subcommand_cache[$key] ]]; then
        _samba_tool_subcommand_cache[$key]=
        _samba_tool_ensure_help "$@"
        while IFS= read -r line; do
            if [[ $line == "Available subcommands:" ]]; then
                in_block=1
                continue
            fi

            (( in_block )) || continue

            if [[ -z ${line//[[:space:]]/} || $line == "For more help"* ]]; then
                break
            fi

            if [[ $line =~ ^[[:space:]]+([^[:space:]]+)[[:space:]]+- ]]; then
                _samba_tool_subcommand_cache[$key]+="${BASH_REMATCH[1]}"$'\n'
            fi
        done <<< "${_samba_tool_help_cache[$key]}"
    fi
}

_samba_tool_subcommands()
{
    local key token

    key=$(_samba_tool_cache_key "$@")
    _samba_tool_ensure_subcommands "$@"

    while IFS= read -r token; do
        [[ -n $token ]] && printf '%s\n' "$token"
    done <<< "${_samba_tool_subcommand_cache[$key]}"
}

# Parse only real option lines from the "Options:" section.  Examples in the
# long help text may contain strings like --newpassword=passw0rd; those are not
# useful completion candidates.
_samba_tool_ensure_options()
{
    local key line in_options=0 matched option option_spec option_re
    local -A seen=()

    key=$(_samba_tool_cache_key "$@")
    if [[ ! -v _samba_tool_option_cache[$key] ]]; then
        _samba_tool_option_cache[$key]=
        _samba_tool_option_spec_cache[$key]=
        _samba_tool_ensure_help "$@"
        option_re='(^|[ ,])(--?[[:alnum:]][[:alnum:]_-]*(=[^ ,[:space:]]+)?)'
        while IFS= read -r line; do
            if [[ $line == "Options:" ]]; then
                in_options=1
                continue
            fi

            (( in_options )) || continue
            [[ $line =~ ^[[:space:]]*- ]] || continue

            while [[ $line =~ $option_re ]]; do
                matched=${BASH_REMATCH[0]}
                option_spec=${BASH_REMATCH[2]}
                option=$option_spec
                if [[ $option == --*=* ]]; then
                    option=${option%%=*}=
                fi
                if [[ ! -v seen[$option] ]]; then
                    seen[$option]=1
                    _samba_tool_option_cache[$key]+="$option"$'\n'
                fi
                _samba_tool_option_spec_cache[$key]+="$option_spec"$'\n'
                line=${line#*"$matched"}
            done
        done <<< "${_samba_tool_help_cache[$key]}"
    fi
}

_samba_tool_options()
{
    local key token

    key=$(_samba_tool_cache_key "$@")
    _samba_tool_ensure_options "$@"

    while IFS= read -r token; do
        [[ -n $token ]] && printf '%s\n' "$token"
    done <<< "${_samba_tool_option_cache[$key]}"
}

_samba_tool_complete_words()
{
    local current_word=$1
    shift

    mapfile -t COMPREPLY < <(compgen -W "$*" -- "$current_word")
}

_samba_tool_contains_word()
{
    local needle=$1 word
    shift

    for word in "$@"; do
        [[ $word == "$needle" ]] && return 0
    done

    return 1
}

_samba_tool_current_word()
{
    local line_prefix

    if [[ -n ${COMP_LINE-} && -n ${COMP_POINT-} ]]; then
        line_prefix=${COMP_LINE:0:COMP_POINT}
        printf '%s' "${line_prefix##*[[:space:]]}"
        return
    fi

    printf '%s' "${COMP_WORDS[COMP_CWORD]}"
}

# If the current token is an option value with an enum form, such as
# --color=always|never|auto, complete the value after "=".
_samba_tool_complete_option_value()
{
    local current_word=$1
    local completion_word option option_prefix option_values value
    local -a options
    shift

    [[ $current_word == --*=* ]] || return 1

    completion_word=${COMP_WORDS[COMP_CWORD]-}
    option_prefix=${current_word%%=*}=
    _samba_tool_ensure_options "$@"
    mapfile -t options <<< "${_samba_tool_option_spec_cache[$(_samba_tool_cache_key "$@")]}"

    for option in "${options[@]}"; do
        [[ -n $option ]] || continue
        [[ $option == "$option_prefix"* && $option == *"|"* ]] || continue
        option_values=${option#*=}

        while IFS= read -r value; do
            if [[ $completion_word == "$current_word" ]]; then
                COMPREPLY+=("$option_prefix$value")
            else
                COMPREPLY+=("$value")
            fi
        done < <(compgen -W "${option_values//|/ }" -- "${current_word#*=}")

        return 0
    done

    return 1
}

# Walk already typed words while they are known subcommands.  The resulting path
# tells us which `samba-tool ... --help` output should be used for candidates.
_samba_tool_current_command_path()
{
    local i word
    local -a path subcommands

    path=()
    for (( i = 1; i < COMP_CWORD; i++ )); do
        word=${COMP_WORDS[i]}

        [[ $word == -- ]] && break
        [[ $word == -* ]] && continue

        _samba_tool_ensure_subcommands "${path[@]}"
        mapfile -t subcommands <<< "${_samba_tool_subcommand_cache[$(_samba_tool_cache_key "${path[@]}")]}"
        _samba_tool_contains_word "$word" "${subcommands[@]}" || break
        path+=("$word")
    done

    _samba_tool_command_path=("${path[@]}")
}

_samba_tool_complete()
{
    local current_word
    local -a command_path subcommands options candidates

    if declare -F _init_completion >/dev/null; then
        _init_completion -n : || return
    else
        COMPREPLY=()
    fi
    current_word=$(_samba_tool_current_word)

    _samba_tool_current_command_path
    command_path=("${_samba_tool_command_path[@]}")

    if _samba_tool_complete_option_value "$current_word" "${command_path[@]}"; then
        return
    fi

    if [[ $current_word == --*=* ]]; then
        compopt -o default 2>/dev/null || true
        return
    fi

    _samba_tool_ensure_subcommands "${command_path[@]}"
    _samba_tool_ensure_options "${command_path[@]}"
    mapfile -t subcommands <<< "${_samba_tool_subcommand_cache[$(_samba_tool_cache_key "${command_path[@]}")]}"
    mapfile -t options <<< "${_samba_tool_option_cache[$(_samba_tool_cache_key "${command_path[@]}")]}"

    if [[ $current_word == -* ]]; then
        _samba_tool_complete_words "$current_word" "${options[@]}"

        if [[ $current_word == --*=* ]] || [[ " ${COMPREPLY[*]} " == *"= "* ]]; then
            compopt -o nospace 2>/dev/null || true
        fi
        [[ ${#COMPREPLY[@]} -eq 0 ]] && compopt -o default 2>/dev/null || true
        return
    fi

    candidates=("${subcommands[@]}" "${options[@]}")
    _samba_tool_complete_words "$current_word" "${candidates[@]}"

    [[ ${#COMPREPLY[@]} -eq 0 ]] && compopt -o default 2>/dev/null || true
}

complete -F _samba_tool_complete samba-tool
