bash-completion

# list of commands/arguments
__nnictl_cmds="create resume view update stop trial experiment platform webui config log package tensorboard top ss_gen --version"
__nnictl_create_cmds="--config --port --debug --foreground"
__nnictl_resume_cmds="--port --debug --foreground"
__nnictl_view_cmds="--port"
__nnictl_update_cmds="searchspace concurrency duration trialnum"
__nnictl_update_searchspace_cmds="--filename"
__nnictl_update_concurrency_cmds="--value"
__nnictl_update_duration_cmds="--value"
__nnictl_update_trialnum_cmds="--value"
__nnictl_stop_cmds="--port --all"
__nnictl_trial_cmds="ls kill"
__nnictl_trial_ls_cmds="--head --tail"
__nnictl_trial_kill_cmds="--trial_id"
__nnictl_experiment_cmds="show status list delete export import save load"
__nnictl_experiment_list_cmds="--all"
__nnictl_experiment_delete_cmds="--all"
__nnictl_experiment_export_cmds="--filename --type --intermediate"
__nnictl_experiment_import_cmds="--filename"
__nnictl_experiment_save_cmds="--path --saveCodeDir"
__nnictl_experiment_load_cmds="--path --codeDir --logDir"
__nnictl_platform_cmds="clean"
__nnictl_platform_clean_cmds="--config"
__nnictl_webui_cmds="url"
__nnictl_config_cmds="show"
__nnictl_log_cmds="stdout stderr trial"
__nnictl_log_stdout_cmds="--tail --head --path"
__nnictl_log_stderr_cmds="--tail --head --path"
__nnictl_log_trial_cmds="--trial_id"
__nnictl_package_cmds="install show list uninstall"
__nnictl_package_install_cmds="--name"
__nnictl_package_list_cmds="--all"
__nnictl_tensorboard_cmds="start stop"
__nnictl_tensorboard_start_cmds="--trial_id --port"
__nnictl_top_cmds="--time"
__nnictl_ss_gen_cmds="--trial_command --trial_dir --file"

# list of commands that accept an experiment ID as second argument
__nnictl_2nd_expid_cmds=" resume view stop top "
# list of commands that accept an experiment ID as third argument
__nnictl_3rd_expid_cmds=" update trial experiment webui config log tensorboard "


# remove already set arguments from candidates
__nnictl_remain_args()
{
    local ret=${!1}  # ret = $__nnictl_xxx_cmds
    # prevent that "--trial_id" changes to "--_id" in this situation: "nnictl trial kill --trial_id"
    if [[ ${ret} != "--trial_id" ]]; then
        for arg in "${COMP_WORDS[@]::${#COMP_WORDS[@]}-1}"; do
            local ret=${ret/$arg/}  # remove it from $ret
        done
    fi
    echo $ret
}

# complete files with specific extension
__nnictl_complete_extension()
{
    COMPREPLY=($(compgen -f -X "!*.$1" -- ${COMP_WORDS[-1]}))
    if [[ -z "${COMPREPLY[*]}" ]]; then
        # if there is no matching file here, search in sub-directories
        COMPREPLY=($(compgen -d -S "/" -- ${COMP_WORDS[-1]}))
        compopt -o nospace
    fi
}

_nnictl()
{
    local cur=${COMP_WORDS[-1]}
    local last=${COMP_WORDS[-2]}
    if [[ ${#COMP_WORDS[@]} -eq 2 ]]; then
        # completing first argument from __nnictl_cmds
        COMPREPLY=($(compgen -W "$__nnictl_cmds" -- "${COMP_WORDS[1]}"))

    elif [[ ${#COMP_WORDS[@]} -eq 3 ]]; then
        # completing second argument from __nnictl_${FirstArg}_cmds
        local args=__nnictl_${COMP_WORDS[1]}_cmds
        COMPREPLY=($(compgen -W "${!args}" -- "${COMP_WORDS[2]}"))

        # add experiment IDs to candidates if desired
        if [[ $__nnictl_2nd_expid_cmds =~ " ${COMP_WORDS[1]} " ]]; then
            local experiments=$(ls ~/nni-experiments 2>/dev/null)
            COMPREPLY+=($(compgen -W "$experiments" -- $cur))
        fi

    elif [[ $last != -* || $last == --debug || $last == --foreground || $last == --intermediate || $last == --all ]]; then
        # last argument does not starts with "-", so this one is likely to be "--xxx"
        local args=__nnictl_${COMP_WORDS[1]}_${COMP_WORDS[2]}_cmds
        if [[ $args =~ "-" || $__nnictl_2nd_expid_cmds =~ " ${COMP_WORDS[1]} " ]]; then
            # the second argument starts with "-" or is an experiment id, use __nnictl_${FirstArg}_cmds
            local args=__nnictl_${COMP_WORDS[1]}_cmds
        fi
        # remove already set arguments from candidates
        local remain_args=$(__nnictl_remain_args ${args})
        COMPREPLY=($(compgen -W "$remain_args" -- $cur))

        # if this is 3rd arguments, try adding experiment IDs to candidates
        if [[ ${#COMP_WORDS[@]} -eq 4 ]]; then
            if [[ $__nnictl_3rd_expid_cmds =~ " ${COMP_WORDS[1]} " && ${COMP_WORDS[2]} != "list" ]]; then
                local experiments=$(ls ~/nni-experiments 2>/dev/null)
                COMPREPLY+=($(compgen -W "$experiments" -- $cur))
            fi
        fi

    elif [[ ${COMP_WORDS[2]} == "export" ]]; then
        # "export" command is somewhat unique
        if [[ " --type -t " =~ " $last " ]]; then
            COMPREPLY=($(compgen -W "json csv" -- $cur))

        elif [[ " --filename -f " =~ " $last " ]]; then
            # try to detect whether complete CSV file or JSON file
            [[ "$COMP_LINE" =~ "csv" ]] && local export_csv=1
            [[ "$COMP_LINE" =~ "json" ]] && local export_json=1
            if [[ -n $export_csv && -z $export_json ]]; then
                local ext="csv"         # CSV only
            elif [[ -n $export_json && -z $export_csv ]]; then
                local ext="json"        # JSON only
            else
                local ext="@(csv|json)" # both
            fi
            __nnictl_complete_extension "$ext"
        fi

    elif [[ " --trial_id -T " =~ " $last " ]]; then
        # complete trial ID
        if [[ -e ${HOME}/nni-experiments/${COMP_WORDS[2]} ]]; then
            local trials=$(ls -d ~/nni-experiments/${COMP_WORDS[2]}/trials/* 2>/dev/null | grep -o '[^/]*$')
        elif [[ -e "${HOME}/nni-experiments/${COMP_WORDS[3]}" ]]; then
            local trials=$(ls -d ~/nni-experiments/${COMP_WORDS[3]}/trials/* 2>/dev/null | grep -o '[^/]*$')
        else
            local trials=$(ls -d ~/nni-experiments/*/trials/* 2>/dev/null | grep -o '[^/]*$')
        fi
        COMPREPLY=($(compgen -W "$trials" -- $cur))

    elif [[ " --config -c " =~ " $last " ]]; then
        # complete YAML file
        __nnictl_complete_extension "@(yml|yaml)"

    elif [[ " --filename -f " =~ " $last " ]]; then
        # complete JSON file
        __nnictl_complete_extension "json"

    fi
}

if [[ ${BASH_VERSINFO[0]} -le 4 && ${BASH_VERSINFO[1]} -le 4 ]]; then
    complete -F _nnictl nnictl
else
    complete -o nosort -F _nnictl nnictl
fi