# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#
_have nodetool && _have cqlsh &&
{

    show_keyspaces()
    {
        local ks=$(get_keyspaces)
        COMPREPLY=( $(compgen -W "$ks" -- "$1") )
    }

    get_keyspaces()
    {
        [ -z "$keyspaces" ] && keyspaces=$(echo "DESCRIBE KEYSPACES" | cqlsh | egrep -v '^$')
        echo $keyspaces
    }

    show_datacenters()
    {
        cur=$1
        set|grep -q ^dcs || dcs=$(echo "select data_center from system.peers;"|cqlsh |tail -n +4|sort|uniq|awk '{if(length($1)>1) print $1}'|xargs)
        COMPREPLY=( $(compgen -W "$dcs" -- "$cur") )
    }

    show_cfs()
    {
        local cur prev cfs
        prev=$1
        cur=$2
        cfs=$(get_cfs $1 $2)
        COMPREPLY=( $(compgen -W "$cfs" -- "$cur") )
    }

    get_cfs()
    {
        local prev
        prev=$1
        [ -z "${cf[$prev]}" ] && cf[$prev]=$(echo "DESCRIBE COLUMNFAMILIES" | cqlsh -k ${prev} | egrep -v '^$')
        echo ${cf[$prev]}
    }

    show_last_cfs()
    {
        local cur cfs re
        cur=$1
        re=$(echo ${COMP_WORDS[@]:3:$(($COMP_CWORD - 3))} | sed -e 's/ /\\|/g')
        cfs=$(get_cfs ${COMP_WORDS[2]} | sed -e "s/$re//g")
        COMPREPLY=( $(compgen -W "$cfs" -- "${cur}") )
    }

    _nodetool()
    {
        local cur prev ks
        COMPREPLY=()
        _get_comp_words_by_ref cur prev

        local shopt='
            bootstrap
            clientstats
            compactionhistory
            compactionstats
            decommission
            describecluster
            disableauditlog
            disablebackup
            disablebinary
            disablefullquerylog
            disablegossip
            disablehandoff
            disablehintsfordc
            disableoldprotocolversions
            drain
            enableauditlog
            enablebackup
            enablebinary
            enablefullquerylog
            enablegossip
            enablehandoff
            enablehintsfordc
            enableoldprotocolversions
            failuredetector
            gcstats
            getbatchlogreplaythrottle
            getcompactionthroughput
            getconcurrentcompactors
            getconcurrentviewbuilders
            getinterdcstreamthroughput
            getlogginglevels
            getmaxhintwindow
            getseeds
            getstreamthroughput
            gettimeout
            gettraceprobability
            gossipinfo
            handoffwindow
            help
            invalidatecountercache
            invalidatekeycache
            invalidaterowcache
            join
            listsnapshots
            pausehandoff
            proxyhistograms
            rangekeysample
            refreshsizeestimates
            reloadlocalschema
            reloadseeds
            reloadssl
            reloadtriggers
            repair_admin
            replaybatchlog
            resetfullquerylog
            resetlocalschema
            resumehandoff
            ring
            setbatchlogreplaythrottle
            setconcurrentcompactors
            setconcurrentviewbuilders
            sethintedhandoffthrottlekb
            setinterdcstreamthroughput
            setlogginglevel
            setmaxhintwindow
            settimeout
            status
            statusbackup
            statusbinary
            statusgossip
            statushandoff
            stopdaemon
            tablestats
            tpstats
            version
            '

        local lngopt='
            assassinate
            cleanup
            clearsnapshot
            compact
            describering
            disableautocompaction
            enableautocompaction
            flush
            garbagecollect
            getcompactionthreshold
            getendpoints
            getreplicas
            getsstables
            import
            info
            move
            netstats
            profileload
            rebuild
            rebuild_index
            refresh
            relocatesstables
            removenode
            repair
            scrub
            setcachecapacity
            setcachekeystosave
            setcompactionthreshold
            setcompactionthroughput
            setstreamthroughput
            settraceprobability
            snapshot
            statusautocompaction
            stop
            tablehistograms
            toppartitions
            truncatehints
            upgradesstables
            verify
            viewbuildstatus
	    '

        local optwks='
            cleanup
            clearsnapshot
            compact
            describering
            flush
            garbagecollect
            getcompactionthreshold
            getendpoints
            getreplicas
            getsstables
            import
            profileload
            rebuild_index
            refresh
            relocatesstables
            repair
            scrub
            setcompactionthreshold
            snapshot
            statusautocompaction
            tablehistograms
            toppartitions
            verify
            viewbuildstatus
            '

        local optwcfs='
            cleanup
            compact
            disableautocompaction
            enableautocompaction
            flush
            garbagecollect
            getreplicas
            import
            profileload
            relocatesstables
            repair
            scrub
            statusautocompaction
            toppartitions
            upgradesstables
            verify
            '

        if [[ $COMP_CWORD -eq 1 ]] ; then
            COMPREPLY=( $(compgen -W "${lngopt} ${shopt}" -- "${cur}") )
        elif [[ $(echo "${lngopt}"|egrep -c "\b${prev}\b") -gt 0 ]] ; then
            if echo $optwks|grep -q "\b$prev\b" ; then
                show_keyspaces "${cur}"
            else
                case "${prev}" in
                    removenode)
                        # we don't want to lose time using nodetool status a 2nd time
                        # in case of force or status
                        if [[ "${cur}" =~ ^(f|s) ]] ; then
                            COMPREPLY=( $(compgen -W "status force" -- "${cur}") )
                        else
                            [ -z "$IDS" ] && IDS=$(nodetool status|grep %|awk '{print $7}'|xargs)
                            COMPREPLY=( $(compgen -W "status force $IDS" -- "${cur}") )
                        fi
                        return 0
                        ;;
                    stop)
                        COMPREPLY=( $(compgen -W "COMPACTION VALIDATION CLEANUP SCRUB VERIFY INDEX_BUILD" -- "${cur}") )
                        return 0
                        ;;
                    info)
                        COMPREPLY=( $(compgen -W "-T --tokens" -- "${cur}") )
                        return 0
                        ;;
                    rebuild|disablehintsfordc|enablehintsfordc)
                        show_datacenters "${cur}"
                        return 0
                        ;;
                    upgradesstables)
                        ks=$(get_keyspaces)
                        COMPREPLY=( $(compgen -W "-a --include-all-sstables $ks" -- "${cur}") )
                        return 0
                        ;;
                esac
            fi
        elif [[ $COMP_CWORD -eq 3 ]] ; then
            case "${COMP_WORDS[1]}" in
                cleanup|compact|flush|garbagecollect|getcompactionthreshold|getendpoints|getreplicas|getsstables|import|profileload|rebuild_index|refresh|relocatesstables|repair|scrub|setcompactionthreshold|statusautocompaction|tablehistograms|toppartitions|verify)
                    show_cfs ${prev} ${cur}
                    return 0
                    ;;
                upgradesstables)
                    if [[ ! ${prev} == -* ]]; then
                        show_cfs ${prev} ${cur}
                    fi
                    return 0
                    ;;
                snapshot)
                    COMPREPLY=( $(compgen -W "-cf" -- "${cur}") )
                    return 0
                    ;;
            esac
        elif [[ "${optwcfs}" == *${COMP_WORDS[1]}* ]] ; then
            show_last_cfs ${cur}
        elif [[ $COMP_CWORD -eq 4 && ${COMP_WORDS[1]} == "snapshot" ]] ; then
            show_cfs ${COMP_WORDS[2]} ${cur}
        elif [[ $COMP_CWORD -eq 5 && ${COMP_WORDS[1]} == "snapshot" ]] ; then
            COMPREPLY=( $(compgen -W "-t" -- "${cur}") )
        fi
    }
    complete -F _nodetool nodetool
}
