#!/bin/bash
#-------------------------------------------------------------------------------
# Copyright (C) 2012-2019 British Crown (Met Office) & Contributors.
#
# This file is part of Rose, a framework for meteorological suites.
#
# Rose is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Rose 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 Rose. If not, see <http://www.gnu.org/licenses/>.
#-------------------------------------------------------------------------------
# NAME
#     rose-bash-completion
#
# SYNOPSIS
#     . $ROSE_HOME/etc/rose-bash-completion
#
# DESCRIPTION
#     Set Bash auto-completion for rose/rosie commands.
#
#     Users should source this file in their ~/.bashrc, using something
#     like this:
#     if [[ $- =~ i && -f /path/to/rose-bash-completion ]]; then
#         . /path/to/rose-bash-completion
#     fi
#     where /path/to/rose-bash-completion is replaced by the path to
#     this file.
#
#     Administrators may want to place this file in the
#     /etc/bash_completion.d/ (or equivalent) directory.
#-------------------------------------------------------------------------------

_rose() {
    local current possible_subcommand rose_command ROSE_DIR ok_subcommands
    local subcommand current_option_line current_option_arg_index
    local current_option current_option_arg current_index suboption_lines
    local suboption_sep_lines suboptions sub_args custom_function_name metavar
    local formal_current_option
    COMPREPLY=()
    current="${COMP_WORDS[COMP_CWORD]}"
    possible_subcommand="${COMP_WORDS[1]:-}"
    rose_command="${COMP_WORDS[0]}"
    rose --version 1>/dev/null 2>&1 || return 1
    if grep -q "bin/ros[^/]*$" <<<"$rose_command"; then
        ROSE_DIR=$(dirname "$rose_command")
        eval ROSE_DIR="$ROSE_DIR"
        ROSE_DIR=$(dirname $(cd $ROSE_DIR && pwd -P))
        rose_command=$(basename "$rose_command")
    else
        ROSE_DIR=$(__rose_get_rose_dir)
    fi
    ok_subcommands=$(cd $ROSE_DIR/bin && ls $rose_command-* | \
                     sed "s/^$rose_command-//g")

    # Determine the subcommand (e.g. app-run for rose app-run)
    subcommand=
    if [[ -n $possible_subcommand ]]; then
        for ok_subcommand in $ok_subcommands; do
            if [[ $ok_subcommand == $possible_subcommand ]]; then
                subcommand=$possible_subcommand
                break
            fi
        done
    fi
    if [[ -n $subcommand && -L $ROSE_DIR/bin/rose-$subcommand ]]; then
        subcommand=$(readlink $subcommand)
    fi          

    # Determine option or option argument completion.
    current_option_line=$(printf "%s\n" "${COMP_WORDS[@]}" | grep -n "^-" | \
                          tail -1)
    if [[ -n $current_option_line ]]; then
        current_option_arg_index=$( \
            grep -o "^[0-9][0-9]*" <<<"$current_option_line")
        current_option=$(grep -o '\-.*$' <<<"$current_option_line")
        current_option_arg="${COMP_WORDS[$current_option_arg_index]:-}"
        if [[ $current_option_arg == "=" ]]; then
            current_option_arg_index=$(( current_option_arg_index + 1 ))
            current_option_arg="${COMP_WORDS[$current_option_arg_index]:-}"
        fi
    else
        current_option_arg_index=
        current_option_arg=
        current_option=
    fi
    current_index=${#COMP_WORDS[@]}
    current_index=$(( current_index - 1 ))

    # Supply any completion information if we can.
    if [[ -n $subcommand ]]; then
        suboption_lines=$(__${rose_command}_help $subcommand | \
                          sed -n '/^    -/p')
        suboption_sep_lines=$( \
            sed 's/^ *//g; s/, \(-\)/\n\1/g' <<<"$suboption_lines")
        suboptions=$(sed 's/=/ /g' <<<"$suboption_sep_lines")
        if grep -q '^'$current_option'$' <<<"$suboptions"; then
            # The current option does not take an argument - forget about it.
            current_option=
        fi
        if [[ -z $current_option ]] || \
                [[ $current_index -gt $current_option_arg_index && \
                -n $current_option_arg ]]; then
            # No relevant current option - supply options and arguments.
            custom_function_name="_${rose_command}_"${subcommand//-/_}_ARGS
            sub_args=
            if type -t $custom_function_name 1>/dev/null 2>&1; then
                sub_args=$($custom_function_name "$current")
            fi
            suboptions=$( \
                sed 's/=.*$/=/g; s/  *.*$//g; s/ *$//g;' \
                <<<"$suboption_sep_lines")
            COMPREPLY=( $(compgen -W "$suboptions $sub_args" -- "$current") )
            return 0
        fi
        if grep -q "^$current_option\>" <<<"$suboptions" && \
                [[ $current != $current_option ]]; then
            # The current string must be an argument to the option.
            current_option_args=
            if [[ $current != "=" ]]; then
                current_option_args=$current
            fi
            metavar=$(sed -n "s/^$current_option \([^ ]*\)/\1/p" \
                      <<<"$suboptions")
            if [[ $metavar == "DIR" || $metavar == "PATH" || \
                  $metavar = "FILE" ]]; then
                return 0
            fi
            formal_current_option=$( \
                sed -n "s/^ *\(-[^ =]*\).*, $current_option\b.*/\1/p" \
                <<<"$suboption_lines")
            if [[ -n $formal_current_option ]]; then
                current_option=$formal_current_option
            fi
            custom_function_name="_${rose_command}_"${subcommand//-/_}
            custom_function_name+=${current_option//-/_}
            if type -t $custom_function_name 1>/dev/null 2>&1; then
                $custom_function_name "$current_option_args"
                return $?
            fi
            COMPREPLY=
            return 1
        fi
        suboptions=$(sed 's/=.*$/=/g; s/  *.*$//g; s/ *$//g;' \
                     <<<"$suboption_sep_lines")
        COMPREPLY=( $(compgen -W "$suboptions" -- "$current") )
        return 0
    fi
    # There is an incomplete subcommand or none at all - supply a list of them.
    COMPREPLY=( $(compgen -W "$ok_subcommands" -- "$current") )
    return 0
}

_rose_app_run__app_mode() {
    COMPREPLY=( $(compgen -W "fcm_make rose_ana rose_arch rose_prune" -- $1) )
}

_rose_app_run__command_key() {
    local config_dir command_keys
    config_dir=$(__rose_get_config_dir)
    if [[ -d $config_dir ]]; then
        command_keys=$(rose config --keys --file=$config_dir/rose-app.conf command 2>/dev/null)
        COMPREPLY=( $(compgen -W "$command_keys" -- $1) )
        return 0
    fi
    return 1
}

_rose_app_run__opt_conf_key() {
    local config_dir opt_conf_keys
    config_dir=$(__rose_get_config_dir)
    if [[ -d $config_dir/opt ]]; then
        opt_conf_keys=$(cd $config_dir/opt/ && ls rose-app-*.conf | \
                        sed "s/rose-app-\(.*\).conf/\1/g")
        COMPREPLY=( $(compgen -W "$opt_conf_keys" -- $1) )
        return 0
    fi
    return 1
}

_rose_app_upgrade_ARGS() {
    local config_dir meta_path meta_path_option all_versions_option
    config_dir=$(__rose_get_config_dir)
    meta_path=$(__rose_get_meta_path)
    all_versions_option=
    meta_path_option="--meta-path=$meta_path"
    if [[ -z $meta_path ]]; then
        meta_path_option=
    fi
    if printf "%s " "${COMP_WORDS[@]}" | \
       grep -q " --all-versions\| -a "; then
        all_versions_option="--all-versions"
    fi
    rose app-upgrade $all_versions_option $meta_path_option \
        --config=$config_dir 2>/dev/null | sed 's/^. \(.*\)$/\1/g'
}

_rose_host_select__rank_method() {
    COMPREPLY=( $(compgen -W "load fs mem random" $1) )
}

_rose_macro_ARGS() {
    local config_dir meta_path meta_path_option
    config_dir=$(__rose_get_config_dir)
    meta_path=$(__rose_get_meta_path)
    meta_path_option="--meta-path=$meta_path"
    if [[ -z $meta_path ]]; then
        meta_path_option=
    fi  
    rose macro -q $meta_path_option --config=$config_dir 2>/dev/null | \
        sed "s/^\[.*\] //g"
}

_rose_metadata_graph_ARGS() {
    local config_dir meta_path
    config_dir=$(__rose_get_config_dir)
    meta_path=$(__rose_get_meta_path)
    if [[ -f $config_dir/rose-app.conf ]]; then
        PATH=$PATH:$meta_path rose config --meta --keys \
            --file=$config_dir/rose-app.conf 2>/dev/null | \
            sed '/=/d'
    else
        rose config --keys --file=$config_dir/rose-meta.conf 2>/dev/null | \
            sed '/=/d'
    fi
}

_rose_metadata_graph__property() {
    COMPREPLY=( $(compgen -W "trigger" -- $1) )
    return 0
}

_rose_stem__group() {
    local config_dir groups
    config_dir=$(__rose_stem_get_config_dir)
    if [[ ! -e $config_dir/meta/rose-meta.conf ]]; then
        return 1
    fi
    groups=$(rose config --file=$config_dir/meta/rose-meta.conf \
             jinja2:suite.rc=RUN_NAMES widget[rose-config-edit] | \
             grep -o "\-\-choices.[^ ]*" | sed "s/--choices.//g;" | \
             tr '\n' ',' | tr ',' ' ')
    tasks=$(rose config --file=$config_dir/meta/rose-meta.conf \
            jinja2:suite.rc=RUN_NAMES widget[rose-config-edit] | \
            tr ' ' '\n' | sed -n '/rose.config_editor.*/d; /^\([^-].*\)/p')
    COMPREPLY=( $(compgen -W "$groups $tasks" -- $1) )
    return 0
}

# Should we separate this from _rose_stem__group?
_rose_stem__task() {
    _rose_stem__group "$@"
}

_rose_suite_gcontrol__name() {
    local names
    names=$(cylc print -xy 2>/dev/null)
    COMPREPLY=( $(compgen -W "$names" -- $1) )
}

_rose_suite_hook__mail_cc() {
    local users
    prev_users=$1
    user_tail=${prev_users##*,}
    user_head=${prev_users%,*}
    users=$(__rose_get_users)
    if [[ -n $user_tail ]] && grep -q "\<$user_tail\>" <<<"$users"; then
        # A complete user name, now need a comma.
        users=$(sed "s/^/$prev_users,/; s/ / $prev_users,/g" <<<"$users")
        COMPREPLY=( $(compgen -W "$users" -- $prev_users) )
        return
    fi
    if [[ -n $user_head && $user_head != $prev_users ]]; then
        users=$(sed "s/^/$user_head,/; s/ / $user_head,/g" <<<"$users")
    fi
    COMPREPLY=( $(compgen -W "$users" -- $prev_users) )
}

_rose_suite_log__name() {
    local names
    names=$(cylc print -xy 2>/dev/null)
    COMPREPLY=( $(compgen -W "$names" -- $1) )
}

_rose_suite_log__user() {
    local users
    users=$(__rose_get_users)
    COMPREPLY=( $(compgen -W "$users" -- $1) )
}

_rose_suite_run__host() {
    local hosts
    hosts=$(rose config rose-suite-run hosts 2>/dev/null)
    hosts=$(__rose_get_expanded_hosts "$hosts")
    COMPREPLY=( $(compgen -W "$hosts" -- $1) )
}

_rose_suite_run__opt_conf_key() {
    local config_dir opt_conf_keys
    config_dir=$(__rose_get_config_dir)
    if [[ -d $config_dir/opt ]]; then
        opt_conf_keys=$(cd $config_dir/opt/ && ls rose-suite-*.conf | \
                        sed "s/rose-suite-\(.*\).conf/\1/g")
        COMPREPLY=( $(compgen -W "$opt_conf_keys" -- $1) )
        return 0
    fi
    return 1
}

_rose_suite_run__run() {
    COMPREPLY=( $(compgen -W "reload restart run" -- $1) )
}

_rose_suite_shutdown__name() {
    local names
    names=$(cylc print -xy 2>/dev/null)
    COMPREPLY=( $(compgen -W "$names" -- $1) )
}

_rose_suite_stop__name() {
    _rose_suite_shutdown__name "$@"
    return $?
}

_rose_task_run__app_mode() {
    _rose_app_run__app_mode "$@"
    return $?
}

_rose_task_run__command_key() {
    _rose_app_run__command_key "$@"
    return $?
}

_rose_task_run__opt_conf_key() {
    _rose_app_run__opt_conf_key "$@"
    return $?
}

_rose_test_battery_ARGS() {
    cd $ROSE_DIR/t && ls -d ros*
}

_rosie_create__prefix() {
    __rosie_get_prefixes "$@"
    return $?
}

_rosie_go__prefix() {
    __rosie_get_prefixes "$@"
    return $?
}

_rosie_lookup__prefix() {
    __rosie_get_prefixes "$@"
    return $?
}

_rosie_ls__prefix() {
    __rosie_get_prefixes "$@"
    return $?
}

__rose_get_config_dir() {
    local config_dir
    config_dir=$( \
        printf "%s " "${COMP_WORDS[@]}" | \
        sed -n "s/.*--config = \([^ ]*\).*/\1/p; s/.*-C *\([^ ]*\).*/\1/p")
    if [[ -z $config_dir ]]; then
        config_dir=$PWD
    fi
    eval config_dir="$config_dir"
    cd $config_dir 1>/dev/null 2>&1 && pwd -P
}

__rose_stem_get_config_dir() {
    local config_dir
    config_dir=$(__rose_get_config_dir)
    if [[ ! -e $config_dir/rose-stem/suite.rc ]]; then
        while [[ ! -d $config_dir/rose-stem ]] && \
              svn info $config_dir >/dev/null 2>&1; do
            new_config_dir=$(cd $config_dir/.. && pwd -P)
            if [[ $new_config_dir == $config_dir ]]; then
                return 1
            fi
            config_dir=$new_config_dir
        done
        if [[ ! -e $config_dir/rose-stem/suite.rc ]]; then
            return 1
        fi
    fi
    config_dir=$config_dir/rose-stem
    cd $config_dir 1>/dev/null 2>&1 && pwd -P
}

__rose_get_expanded_hosts() {
    for host in "$@"; do
        rose config rose-host-select group{$host} || echo $host
    done
}

__rose_get_meta_path() {
    local meta_path
    meta_path=$( \
        printf "%s " "${COMP_WORDS[@]}" | \
        sed -n "s/.*--meta-path = \([^ ]*\).*/\1/p; s/.*-M *\([^ ]*\).*/\1/p")
    if [[ -n $meta_path ]]; then
        eval meta_path="$meta_path"
        cd $meta_path 1>/dev/null 2>&1 && pwd -P
    fi
}

__rose_get_rose_dir() {
    rose --version | sed "s/.*(\(.*\))/\1/"
}

__rose_get_users() {
    getent passwd | cut -d: -f1 | LANG=C sort -u
}
    
__rose_help() {
    local subcommand
    subcommand=$1
    if type -t _rose_help__${subcommand//-/_}; then
        _rose_help__${subcommand//-/_} | sed -n '/^OPTIONS$/,/^[^ ]$/p;'
        return
    fi
    rose help $subcommand 2>/dev/null
}

__rosie_get_prefixes() {
    local prefixes
    prefixes=$(rose config --keys rosie-id | \
               sed -n "s/^prefix-location\.\(.*\)/\1/p")
    COMPREPLY=( $(compgen -W "$prefixes" -- $1) )
}

__rosie_help() {
    local subcommand
    subcommand=$1
    if type -t _rosie_help__${subcommand//-/_}; then
        _rosie_help__${subcommand//-/_}
        return
    fi
    rosie help $subcommand
}

_rose_help__stem() {
    rose help stem
    rose help suite-run
}

_rose_help__task_run() {
    rose help task-run
    rose help app-run
    rose help task-env
}

complete -o bashdefault -o default -o nospace -F _rose rose rosie
