#!/bin/bash
#
# bash completion file for core atomic commands
# copied from the docker bash completion
#
#
# This script provides completion of:
#  - commands and their options
#  - container ids and names
#  - image repos and tags
#  - filepaths
#
# To enable the completions either:
#  - place this file in /etc/bash_completion.d
#  or
#  - copy this file to e.g. ~/.atomic-completion.sh and add the line
#    below to your .bashrc after bash completion features are loaded
#    . ~/.atomic-completion.sh
#
# Note:
# Currently, the completions will not work if the atomic daemon is not
# bound to the default communication port/socket
# If the atomic daemon is using a unix socket for communication your user
# must have access to the socket for the completions to function correctly
#
# Note for developers:
# Please arrange options sorted alphabetically by long name with the short 
# options immediately following their corresponding long form.
# This order should be applied to lists, alternatives and code blocks.


__atomic_previous_extglob_setting=$(shopt -p extglob)
shopt -s extglob

__docker_q() {
	docker=$(awk -F ": " '/^default_docker/ {print $2}' /etc/atomic.conf 2>/dev/null || echo docker)
	docker 2>/dev/null "$@"
}

_atomic_atomic() {
	local boolean_options="
		--help -h
	"
	case "$prev" in
		$main_options_with_args_glob )
			return
			;;
	esac

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$boolean_options $main_options_with_args" -- "$cur" ) )
			;;
		*)
			COMPREPLY=( $( compgen -W "${commands[*]} help" -- "$cur" ) )
			;;
	esac
}

__atomic_containers_all() {
	local IFS=$'\n'
	local containers=( $(__docker_q ps -aq --no-trunc) )
	if [ "$1" ]; then
		containers=( $(__docker_q inspect --format "{{if $1}}{{.Id}}{{end}}" "${containers[@]}") )
	fi
	local names=( $(__docker_q inspect --format '{{.Name}}' "${containers[@]}") )
	names=( "${names[@]#/}" ) # trim off the leading "/" from the container names
	unset IFS
	COMPREPLY=( $(compgen -W "${names[*]} ${containers[*]}" -- "$cur") )
}

__atomic_containers_running() {
	__atomic_containers_all '.State.Running'
}

__atomic_containers_stopped() {
	__atomic_containers_all 'not .State.Running'
}

__atomic_containers_pauseable() {
	__atomic_containers_all 'and .State.Running (not .State.Paused)'
}

__atomic_containers_unpauseable() {
	__atomic_containers_all '.State.Paused'
}

__atomic_image_repos() {
	local repos="$(__docker_q images | awk 'NR>1 && $1 != "<none>" { print $1 }')"
	COMPREPLY=( $(compgen -W "$repos" -- "$cur") )
}

__atomic_image_repos_and_tags() {
	local reposAndTags="$(__docker_q images | awk 'NR>1 && $1 != "<none>" { print $1; print $1":"$2 }')"
	COMPREPLY=( $(compgen -W "$reposAndTags" -- "$cur") )
	__ltrim_colon_completions "$cur"
}

__atomic_image_repos_and_tags_and_ids() {
	local images
	if test $_system == "0"; then
		images="$(__docker_q images -a --no-trunc | awk 'NR>1 { print $3; if ($1 != "<none>") { print $1; print $1":"$2 } }')"
	else
		CHECKOUT_DIR=`test -e /etc/atomic.conf && grep "^checkout_path:" /etc/atomic.conf | cut -d' ' -f2`
		images="$(ls -1 ${CHECKOUT_DIR:/var/lib/containers/atomic/} | grep -v \.[01]$)"
	fi
        COMPREPLY=( $(compgen -W "$images" -- "$cur") )
	__ltrim_colon_completions "$cur"
}

__atomic_containers_and_images() {
	__atomic_containers_all
	local containers=( "${COMPREPLY[@]}" )
	__atomic_image_repos_and_tags_and_ids
	COMPREPLY+=( "${containers[@]}" )
}

__atomic_pos_first_nonflag() {
	local argument_flags=$1

	local counter=$cpos
	while [ $counter -le $cword ]; do
		if [ -n "$argument_flags" ] && eval "case '${words[$counter]}' in $argument_flags) true ;; *) false ;; esac"; then
			(( counter++ ))
		else
			case "${words[$counter]}" in
				-*)
					;;
				*)
					break
					;;
			esac
		fi
		(( counter++ ))
	done

	echo $counter
}

# Transforms a multiline list of strings into a single line string
# with the words separated by "|".
# This is used to prepare arguments to __atomic_pos_first_nonflag().
__atomic_to_alternatives() {
	local parts=( $1 )
	local IFS='|'
	echo "${parts[*]}"
}

# Transforms a multiline list of options into an extglob pattern
# suitable for use in case statements.
__atomic_to_extglob() {
	local extglob=$( __atomic_to_alternatives "$1" )
	echo "@($extglob)"
}

_atomic_help() {
	local counter=$(__atomic_pos_first_nonflag)
	if [ $cword -eq $counter ]; then
		COMPREPLY=( $( compgen -W "${commands[*]}" -- "$cur" ) )
	fi
}

_atomic_info() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "-h" -- "$cur" ) )
			;;
		*)
		    __atomic_image_repos_and_tags
		    ;;
	esac
}

# Subcommand processing.
# Locates the first occurrence of any of the subcommands contained in the
# first argument. In case of a match, calls the corresponding completion
# function and returns 0.
# If no match is found, 1 is returned. The calling function can then
# continue processing its completion.
#
# TODO if the preceding command has options that accept arguments and an
# argument is equal ot one of the subcommands, this is falsely detected as
# a match.
__atomic_subcommands() {
	local subcommands="$1"

	local counter=$(($command_pos + 1))
	while [ $counter -lt $cword ]; do
		case "${words[$counter]}" in
			$(__atomic_to_extglob "$subcommands") )
				subcommand_pos=$counter
				local subcommand=${words[$counter]}
				local completions_func=_atomic_${command}_${subcommand}
				declare -F $completions_func >/dev/null && $completions_func
				return 0
				;;
		esac
		(( counter++ ))
	done
	return 1
}

_atomic_scan() {
    local options_with_args="
		--fetch_cves
	"
	local all_options="$options_with_args
		--help -h
		--all
		--images
		--containers
		--rootfs
		--list
		--scanner
		--scan_type
		--verbose
	"
	[ "$command" = "scan" ] && all_options="$all_options"

	local options_with_args_glob=$(__atomic_to_extglob "$options_with_args")

	case "$prev" in
		$options_with_args_glob )
			return 0
			;;
	esac

    case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
		*)

			local counter=$( __atomic_pos_first_nonflag $( __atomic_to_alternatives "$options_with_args" ) )

			if [ $cword -eq $counter ]; then
				__atomic_containers_and_images
				return 0
			fi

			COMPREPLY=( $( compgen -d "$cur" ) )
			return 0
			;;
	esac
	return 0

}

_atomic_top() {
    local options_with_args="
		--optional
		-o
		-d
        -n
	"
    local all_options="$options_with_args"
    [ "$command" = "top" ] && all_options="$all_options"
    local options_with_args_glob=$(__atomic_to_extglob "$options_with_args")

    case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
		*)
		        __atomic_containers_running
			;;
	esac
	return 0
}

_atomic_verify() {
	case "$cur" in
		*)
		    __atomic_image_repos_and_tags
		    ;;
	esac
}

_atomic_version() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "--recurse -r" -- "$cur" ) )
			;;
		*)
		    __atomic_image_repos_and_tags
		    ;;
	esac
}

_atomic_pull() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "--storage" -- "$cur" ) )
			;;
	esac
}

_atomic_update() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "--force -f --set --container" -- "$cur" ) )
			;;
		*)
		    __atomic_image_repos_and_tags
		    ;;
	esac
}

_atomic_uninstall() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W -- "$cur" ) )
			;;
		*)
		    __atomic_image_repos_and_tags_and_ids
		    ;;
	esac
}


_atomic_diff() {
	local all_options="$options_with_args
		--help
		--json
        --display
        --names-only
        --no-files -n
        --rpms -r
        --verbose -v
	"

	[ "$command" = "diff" ] && all_options="$all_options"


    case "$cur" in
    -*)
        COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
        ;;
    *)
        local counter=$( __atomic_pos_first_nonflag $( __atomic_to_alternatives "$options_with_args" ) )

        if [ $cword -eq $counter ]; then
            __atomic_image_repos_and_tags_and_ids
        fi
        ;;
	esac
}

_atomic_run() {
	local options_with_args="
		--name -n
	"

	local all_options="$options_with_args
		--help
		--spc
               --display
        --quiet
	"

	[ "$command" = "run" ] && all_options="$all_options"

	local options_with_args_glob=$(__atomic_to_extglob "$options_with_args")

	case "$prev" in
		$options_with_args_glob )
			return
			;;
	esac

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
		*)
			local counter=$( __atomic_pos_first_nonflag $( __atomic_to_alternatives "$options_with_args" ) )

			if [ $cword -eq $counter ]; then
				__atomic_image_repos_and_tags_and_ids
			fi
			;;
	esac
}

_atomic_install() {
	local options_with_args="
		--name -n
	"
	local all_options="$options_with_args
		--help
               --display
               --system
               --set
	"

	local options_with_args_glob=$(__atomic_to_extglob "$options_with_args")

	case "$prev" in
		$options_with_args_glob )
			return
			;;
	esac

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
		*)
		    __atomic_image_repos_and_tags_and_ids
			;;
	esac
}

_atomic_images() {
	local all_options="
		--prune
		--help -h
		--noheading -n
	"

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
	esac
}

_atomic_mount() {
	local options_with_args="
		--options -o
	"
	local all_options="$options_with_args
		--help -h
		--live
		--shared
	"

	local options_with_args_glob=$(__atomic_to_extglob "$options_with_args")

	case "$prev" in
		$options_with_args_glob )
			return 0
			;;
	esac

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
		*)

			local counter=$( __atomic_pos_first_nonflag $( __atomic_to_alternatives "$options_with_args" ) )

			if [ $cword -eq $counter ]; then
				__atomic_containers_and_images
				return 0
			fi

			COMPREPLY=( $( compgen -d "$cur" ) )
			return 0
			;;
	esac
	return 0
}

_atomic_upload() {
	local options_with_args="
		--URL -U
		--username -u
		--password -p
		--repo_id --repository_id
		--activation_key -a
	"
	local all_options="$options_with_args
		--help -h
		--pulp
		--satellite
		--verify_ssl
		--debug
	"

	local options_with_args_glob=$(__atomic_to_extglob "$options_with_args")

	case "$prev" in
		$options_with_args_glob )
			return 0
			;;
	esac

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
		*)

			local counter=$( __atomic_pos_first_nonflag $( __atomic_to_alternatives "$options_with_args" ) )

			if [ $cword -eq $counter ]; then
				__atomic_containers_and_images
				return 0
			fi

			COMPREPLY=( $( compgen -d "$cur" ) )
			return 0
			;;
	esac
	return 0
}


_atomic_unmount() {
	case "$prev" in
		"unmount")
			COMPREPLY=( $( compgen -d "$cur" ) )
			return 0
			;;
	esac
	return 0
}

_atomic_stop() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "--name -n" -- "$cur" ) )
			;;
		*)
			__atomic_containers_running
			;;
	esac
}

_atomic_host_deploy() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "-r --reboot --os --preview -h --help" -- "$cur" ) )
			;;
		*)
			;;
	esac
}

_atomic_host_rollback() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "-r --reboot -h --help" -- "$cur" ) )
			;;
		*)
			;;
	esac
}

_atomic_host_status() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "-p --pretty -h --help" -- "$cur" ) )
			;;
		*)
			;;
	esac
}

_atomic_host_upgrade() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "-r --reboot --os --allow-downgrade --check-diff -h --help" -- "$cur" ) )
			;;
		*)
			;;
	esac
}

_atomic_host_unlock() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "--hotfix -h --help" -- "$cur" ) )
			;;
		*)
			;;
	esac
}

_atomic_host_host() {
	local boolean_options="
		--help -h
	"
	case "$prev" in
		$main_options_with_args_glob )
			return
			;;
	esac

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$boolean_options $main_options_with_args" -- "$cur" ) )
			;;
		*)
			COMPREPLY=( $( compgen -W "${commands[*]} help" -- "$cur" ) )
			;;
	esac
}

_atomic_help() {
	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
		*)

			local counter=$( __atomic_pos_first_nonflag $( __atomic_to_alternatives "$options_with_args" ) )

			if [ $cword -eq $counter ]; then
				__atomic_containers_and_images
				return 0
			fi

			COMPREPLY=( $( compgen -d "$cur" ) )
			return 0
			;;
	esac
	return 0
}

_atomic_host() {
	local commands=(
		deploy
	        rollback
		status
		upgrade
	        unlock	
	)

	local completions_func=_atomic_host_${prev}
	declare -F $completions_func >/dev/null && $completions_func

	return 0
}

_atomic_storage() {
	local subcommands="
		export
		import
		reset
		modify
	"

	__atomic_subcommands "$subcommands" && return

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "--help" -- "$cur" ) )
			;;
		*)
			COMPREPLY=( $( compgen -W "$subcommands" -- "$cur" ) )
			;;
	esac
	return 0
}

_atomic_storage_modify() {
	local options_with_args="
		--driver
		--add-device
	"
	local all_options="$options_with_args
		--help
		-h
	"

	local options_with_args_glob=$(__atomic_to_extglob "$options_with_args")

	case "$prev" in
		--add-device)
			COMPREPLY=( $( compgen -f "/dev/" ) )
			return 0
			;;

		--driver)
			COMPREPLY=( $( compgen -W "devicemapper overlay" -- "$cur" ) )
			return 0
			;;
	esac

	case "$cur" in
		-*)
			COMPREPLY=( $( compgen -W "$all_options" -- "$cur" ) )
			;;
	esac
}

_atomic_storage_storage() {
        local boolean_options="
                --help -h
        "
        case "$prev" in
                $main_options_with_args_glob )
                        return
                        ;;
        esac

        case "$cur" in
                -*)
                        COMPREPLY=( $( compgen -W "$boolean_options $main_options_with_args" -- "$cur" ) )
                        ;;
                *)
                        COMPREPLY=( $( compgen -W "${commands[*]} help" -- "$cur" ) )
                        ;;
        esac
}


_atomic_storage_export() {
	case "$cur" in
                -*)
                        COMPREPLY=( $( compgen -W "--graph --dir -h --help" -- "$cur" ) )
                        ;;
                *)
                        ;;
        esac
}

_atomic_storage_import() {
	case "$cur" in
                -*)
                        COMPREPLY=( $( compgen -W "--graph --dir -h --help" -- "$cur" ) )
                        ;;
                *)
                        ;;
        esac

}


_atomic() {
	local previous_extglob_setting=$(shopt -p extglob)
	shopt -s extglob

	local commands=(
		host
		diff
		info
		install
		images
		storage
		mount
		pull
		run
		scan
		stop
		top
		verify
		version
		uninstall
		unmount
		update
		upload
	)

	local main_options_with_args="
		--help
	"

	local main_options_with_args_glob=$(__atomic_to_extglob "$main_options_with_args")

	COMPREPLY=()
	local cur prev words cword

	_get_comp_words_by_ref -n : cur prev words cword

	local command='atomic' cpos=0
	local counter=1
        _system=0
	while [ $counter -lt $cword ]; do
		if test "${words[$counter]}" == "--system"; then
			_system=1
			break
		fi
		(( counter++ ))
	done

	counter=1
	while [ $counter -lt $cword ]; do
		case "${words[$counter]}" in
			$main_options_with_args_glob )
				(( counter++ ))
				;;
			"--system")
				_system=1;;
			-*)
				;;
			*)
				command="${words[$counter]}"
				cpos=$counter
				(( cpos++ ))
				break
				;;
		esac
		(( counter++ ))
	done

	local completions_func=_atomic_${command}
	declare -F $completions_func >/dev/null && $completions_func

	eval "$previous_extglob_setting"
	return 0
}

complete -F _atomic atomic
