# Copyright (C) Canonical, Ltd.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

set -o pipefail

_multipass_complete()
{
    _multipass_instances()
    {
        local state=$1 instances

        instances=$(multipass list --format=csv --no-ipv4 | \tail -n +2)
        [ -n "$state" ] && instances=$(awk -F',' "\$2 == \"$state\"" <<< "$instances")
        instances=$(echo "$instances" | \cut -d',' -f 1  | \tr '\r\n' ' ')

        _add_nonrepeating_args "$instances"
    }

    _multipass_snapshots()
    {
        local instance=$1 snapshots

        snapshots=$(multipass list --snapshots --format=csv 2>/dev/null | \tail -n +2)
        [ -n "$instance" ] && snapshots=$(awk -F',' "\$1 == \"$instance\"" <<< "$snapshots")
        snapshots=$(echo "$snapshots" | \cut -d',' -f 1,2 | \tr ',' '.' | \tr '\r\n' ' ')

        _add_nonrepeating_args "$snapshots"
    }

    _multipass_instances_and_snapshots()
    {
        _multipass_snapshots
        _multipass_instances
    }

    _multipass_restorable_snapshots()
    {
        local instances
        instances=$( multipass info --no-runtime-information --format=csv \
                     | \tail -n +2 \
                     | \awk -F',|\r?\n' '$2 == "Stopped" && $16 > 0' \
                     | \cut -d',' -f 1 \
                     | \tr '\r\n' ' ')

        for instance in ${instances}; do
            _multipass_snapshots "${instance}"
        done
    }

    # Set $opts to the list of available networks.
    _multipass_networks()
    {
        opts=$( multipass networks --format=csv 2>/dev/null | \tail -n +2 | \cut -d',' -f 1 )
    }

    _multipass_instances_with_colon()
    {
        local opts_bak="${opts}"
        opts=""

        _multipass_instances
        local instances="${opts}"
        opts="$opts_bak"

        for instance in ${instances}; do
            opts="${opts} ${instance}:"
        done
    }

    _multipass_settings_keys()
    {
        local keys

        keys=$( multipass get --keys | \tr '\r\n' ' ' )

        for key in ${keys}; do
            if [[ "${COMP_WORDS[*]}" == *"${key}"* ]]; then
                return
            fi
        done

        _add_nonrepeating_args "$keys"
    }

    # When we are completing the `alias` command, this function makes sure that each argument was specified at most
    # once. It fills $opts with the variables which were not yet specified.
    # Precondition: the command must be "alias".
    # Parameters: none.
    _unspecified_alias_vars()
    {
        opts=""

        _add_nonrepeating_args "--no-map-working-directory --help --verbose"

        local definition_found=0

        for ((i=2; i<COMP_CWORD; ++i)); do
            if [[ "${COMP_WORDS[$i]}" =~ ':' ]]; then definition_found=1; fi
        done

        if [[ ${definition_found} == 0 ]] && [[ "${cur}" != ':' ]]; then
            _multipass_instances_with_colon
        fi
    }

    # Set $multipass_aliases to the list of available aliases.
    _multipass_aliases()
    {
        multipass_aliases=$( multipass aliases --format=csv 2>/dev/null | \tail -n +2 | \cut -d',' -f 1 )
    }

    # Removes the comma or equals sign at the end of $cur.
    _remove_trailing_char()
    {
        if [[ $cur =~ [,=]$ ]]; then
            cur="${cur#*[,=]}"
            return 0
        fi

        return 1
    }

    # When options for --network are specified in the form var=val, this function sets $opts to the values needed
    # depending on the value of var (which is specified as the parameter $1).
    _complete_network_var()
    {
        case "$1" in
            "id")
                _multipass_networks
            ;;
            "mac")
                opts=""
            ;;
            "mode")
                opts="auto manual"
            ;;
        esac
    }

    # Set $last_match to the last argument that matches the provided regex, or the empty string if there is no
    # match. Return its position in COMP_WORDS, or 0 when there is no match.
    _find_last_match()
    {
        local last_pos

        last_match=""
        for ((last_pos=COMP_CWORD; last_pos>0; --last_pos)); do
            if [[ "${COMP_WORDS[$last_pos]}" =~ $1 ]]; then
                last_match=${COMP_WORDS[$last_pos]}
                break
            fi
        done

        return ${last_pos}
    }

    # Given the last argument starting with "--" is "--network", scan from there until the end of the command-line
    # for the names of the variables which were already specified. And fill-in $opts with the variables which were
    # not yet specified.
    # Precondition: the last dash-something argument must be "--network".
    # Parameters: $1 the position of the last "--network".
    _unspecified_network_vars()
    {
        opts=""

        local id_found=0
        local mac_found=0
        local mode_found=0

        for ((i=$1+1; i<COMP_CWORD; ++i)); do
            if [[ "${COMP_WORDS[$i]}" == 'id' ]]; then id_found=1; fi
            if [[ "${COMP_WORDS[$i]}" == 'mac' ]]; then mac_found=1; fi
            if [[ "${COMP_WORDS[$i]}" == 'mode' ]]; then mode_found=1; fi
        done

        if [[ ${id_found} == 0 ]]; then opts="${opts} id="; fi
        if [[ ${mac_found} == 0 ]]; then opts="${opts} mac="; fi
        if [[ ${mode_found} == 0 ]]; then opts="${opts} mode="; fi
    }

    # Adds the passed list to the opts array while filtering out duplicates
    _add_nonrepeating_args()
    {
        local array=$1
        local found

        for item in $array; do
            found=0
            for ((i=2; i<CWORD; i++)); do
                if [[ "${WORDS[i]}" == "${item}" ]]; then
                    found=1
                    break
                fi
            done
            if [ ${found} == 0 ]; then
                opts="${opts} ${item}"
            fi
        done
    }

    # Add comma and equals sign to the list of word separators if they are not already there, useful to separate
    # options of the form "opt1=val1,opt2=val2".
    if [[ ! "${COMP_WORDBREAKS}" =~ "," ]]; then
        COMP_WORDBREAKS="${COMP_WORDBREAKS},"
    fi
    if [[ ! "${COMP_WORDBREAKS}" =~ "=" ]]; then
        COMP_WORDBREAKS="${COMP_WORDBREAKS}="
    fi

    local cur cmd opts prev prev2 prev_opts multipass_aliases
    _get_comp_words_by_ref -n := -w WORDS -i CWORD cur prev
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"
    if [[ ${COMP_CWORD} -gt 1 ]]; then
        prev2="${COMP_WORDS[COMP_CWORD-2]}"
    fi
    cmd="${COMP_WORDS[1]}"
    prev_opts=false
    multipass_cmds="authenticate transfer delete exec find help info launch list mount networks \
                    purge recover restart restore shell snapshot start stop suspend umount version get set \
                    clone alias aliases unalias"

    if [[ "${multipass_cmds}" =~ " ${cmd} " || "${multipass_cmds}" =~ ^${cmd} || "${multipass_cmds}" =~ \ ${cmd}$ ]];
    then
        _add_nonrepeating_args "--help"
        opts="${opts} --verbose"
    fi

    case "${cmd}" in
        "exec")
            opts="${opts} --working-directory --no-map-working-directory"
        ;;
        "info")
            _add_nonrepeating_args "--format --snapshots"
        ;;
        "list"|"ls")
            _add_nonrepeating_args "--format --snapshots"
        ;;
        "networks"|"aliases")
            _add_nonrepeating_args "--format"
        ;;
        "delete")
            _add_nonrepeating_args "--all --purge"
        ;;
        "launch")
            _add_nonrepeating_args "--cpus --disk --memory --name --cloud-init --bridged --mount"
            opts="${opts} --network --mount"
        ;;
        "mount")
            opts="${opts} --gid-map --uid-map"
        ;;
        "recover"|"start"|"suspend"|"restart")
            _add_nonrepeating_args "--all"
        ;;
        "stop")
            _add_nonrepeating_args "--all --cancel --time"
        ;;
        "find")
            _add_nonrepeating_args "--show-unsupported --force-update --format"
        ;;
        "unalias")
            _multipass_aliases
            _add_nonrepeating_args "$multipass_aliases"
        ;;
        "transfer"|"copy-files")
            _add_nonrepeating_args "--parents --recursive"
        ;;
        "snapshot")
            _add_nonrepeating_args "--name --comment"
        ;;
        "restore")
            _add_nonrepeating_args "--destructive"
        ;;
        "clone")
            _add_nonrepeating_args "--name"
        ;;
        "get")
            _add_nonrepeating_args "--raw --keys"
        ;;
    esac

    if [[ ${prev} == -* ]]; then
        case "${prev}" in
            "--format"|"-f")
                opts="table json csv yaml"
                prev_opts=true
            ;;
            "--cloud-init")
                _filedir
                return
            ;;
            "--mount")
                _filedir
                return
            ;;
            "--network")
                _multipass_networks
                if [ "${OLD_BASH}" = true ]; then
                    if [[ "${opts}" != "" ]]; then
                        opts="${opts} bridged"
                    fi
                else
                    if [[ "${opts}" != "" ]]; then
                        opts="${opts} bridged id= mode= mac="
                    fi
                    # TODO: Do use spaces after the completion of an interface name or "bridged".
                    compopt -o nospace
                fi
        esac
    else
        local last_match
        _find_last_match ^-.*
        local last_match_pos=$?

        if [ "${OLD_BASH}" = false ] && [[ ${last_match} == "--network" ]]; then
            # We need to complete the option pairs in two steps. One, when the last input char was an equals sign. Two,
            # when the user starts to input something, because prev changes when she starts to type a value.
            if [[ "${cur}" == '=' ]] && [[ "${prev2}" == ',' || "${prev2}" == '--network' ]] ; then
                compopt -o nospace
                _remove_trailing_char
                _complete_network_var "${prev}"

            elif [[ "${prev}" == '=' ]]; then
                compopt -o nospace
                _complete_network_var "${prev2}"

            # Complete now in the case on which there is a comma separating var=val pairs.
            elif [[ "${cur}" == ',' && "${prev2}" == '=' ]] ||
                 [[ "${prev}" == ',' && "${COMP_WORDS[COMP_CWORD-3]}" == '=' ]]; then
                _remove_trailing_char
                compopt -o nospace
                _unspecified_network_vars ${last_match_pos}
            fi
        fi
    fi

    if [[ "$prev_opts" = false ]]; then
        case "${cmd}" in
            "exec"|"suspend"|"restart")
                _multipass_instances "Running"
            ;;
            "stop")
                _get_comp_words_by_ref -n := -w WORDS -i CWORD cur prev
                for ((i=2; i<CWORD; i++)); do
                    if [[ "${WORDS[i]}" == "--force" ]]; then
                        _multipass_instances "Starting"
                        _multipass_instances "Restarting"
                        _multipass_instances "Suspending"
                        _multipass_instances "Suspended"
                        break
                    fi
                done

                _multipass_instances "Running"
            ;;
            "connect"|"sh"|"shell")
                _multipass_instances "Running"
                _multipass_instances "Stopped"
                _multipass_instances "Suspended"
            ;;
            "start")
                _multipass_instances "Stopped"
                _multipass_instances "Suspended"
            ;;
            "umount"|"unmount")
                _multipass_instances
            ;;
            "delete"|"info")
                _multipass_instances_and_snapshots
            ;;
            "recover")
                _multipass_instances "Deleted"
            ;;
            "snapshot")
                _multipass_instances "Stopped"
            ;;
            "clone")
                _multipass_instances "Stopped"
            ;;
            "restore")
                _multipass_restorable_snapshots
            ;;
            "get")
                _multipass_settings_keys
            ;;
            "set")
                _multipass_settings_keys

                COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
                if [[ ${#COMPREPLY[@]} -eq 1 && "${COMPREPLY[0]}" != "--help" && "${COMPREPLY[0]}" != "--verbose" ]]; then
                    opts="${COMPREPLY[0]}="
                    compopt -o nospace
                fi
            ;;
            "mount")
                local source_set=0
                local prev
                _get_comp_words_by_ref -n := -w WORDS -i CWORD cur prev
                # Scan through the current command line to detect if the source
                # positional arg has been set.
                for ((i=2; i<CWORD; i++)); do
                    if [[ "${WORDS[i]}" != -* ]] && \
                       ([[ "${WORDS[i-1]}" != -* ]] || [[ "${WORDS[i-1]}" == *=* ]]); then
                            source_set=1
                            break
                    fi
                done

                if [ ${source_set} == 0 ] ; then
                    if [[ ${prev} != -* ]] || ([[ ${prev} == -* ]] && [[ ${prev} == *=* ]]); then
                        _filedir -d
                        return
                    fi
                elif [ ${source_set} == 1 ] && [[ ${prev} != -* ]]; then
                    _multipass_instances "Running"
                    _multipass_instances "Stopped"
                    _multipass_instances "Suspended"
                fi
            ;;
            "transfer"|"copy-files")
                _multipass_instances "Running"

                COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
                if [[ "${#COMPREPLY[@]}" == "0" ]]; then
                    _filedir
                    return
                fi
            ;;
            "alias")
                _unspecified_alias_vars
                # TODO: Do use spaces after the completion of an option starting with "--".
                compopt -o nospace
            ;;
            "help")
                opts=$multipass_cmds
            ;;
        esac
    fi

    if [[ ${COMP_CWORD} -eq 1 ]]; then
        _multipass_aliases
        opts="${opts} ${multipass_cmds} ${multipass_aliases}"
    fi

    if [[ -n "${opts}" ]]; then
        COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
    fi

    return 0
}

# Check the Bash version. Bash < 4 does not have `compopt`, which we use for completing `launch --network`. We
# disable some completions on that command based on this test.
if [[ $(echo ${BASH_VERSION} | sed s/\\..*//g) -gt 3 ]]; then
    OLD_BASH=false
else
    OLD_BASH=true
fi

complete -F _multipass_complete multipass
