shopt -s extglob

_hg_command_list()
{
    hg --debug help 2>/dev/null | \
	awk 'function command_line(line) {
		 gsub(/,/, "", line)
		 gsub(/:.*/, "", line)
		 split(line, aliases)
		 command = aliases[1]
		 delete aliases[1]
		 print command
		 for (i in aliases)
		     if (index(command, aliases[i]) != 1)
			 print aliases[i]
	     }
	     /^list of commands:/ {commands=1}
	     commands && /^ debug/ {a[i++] = $0; next;}
	     commands && /^ [^ ]/ {command_line($0)}
	     /^global options:/ {exit 0}
	     END {for (i in a) command_line(a[i])}'

}

_hg_option_list()
{
    hg -v help $1 2> /dev/null | \
        awk '/^ *-/ {
		 for (i = 1; i <= NF; i ++) {
		    if (index($i, "-") != 1)
			 break;
		    print $i;
		 }
	     }'
}


_hg_commands()
{
    local all commands result

    all=$(_hg_command_list)
    commands=${all%%$'\n'debug*}
    result=$(compgen -W '$commands' -- "$cur")

    # hide debug commands from users, but complete them if
    # there is no other possible command
    if [ "$result" = "" ]; then
	local debug
	debug=debug${all#*$'\n'debug}
	result=$(compgen -W '$debug' -- "$cur")
    fi

    COMPREPLY=(${COMPREPLY[@]:-} $result)
}

_hg_paths()
{
    local paths="$(hg paths 2> /dev/null | sed -e 's/ = .*$//')"
    COMPREPLY=(${COMPREPLY[@]:-} $( compgen -W '$paths' -- "$cur" ))
}

_hg_repos()
{
    local i
    for i in $( compgen -d -- "$cur" ); do
        test ! -d "$i"/.hg || COMPREPLY=(${COMPREPLY[@]:-} "$i")
    done
}

_hg_status()
{
    local files="$( hg status -n$1 . 2> /dev/null)"
    COMPREPLY=(${COMPREPLY[@]:-} $( compgen -W '$files' -- "$cur" ))
}

_hg_tags()
{
    local tags="$(hg tags 2> /dev/null |
                      sed -e 's/[0-9]*:[a-f0-9]\{40\}$//; s/ *$//')"
    COMPREPLY=( ${COMPREPLY[@]:-} $(compgen -W '$tags' -- "$cur") )
}

# this is "kind of" ugly...
_hg_count_non_option()
{
    local i count=0
    local filters="$1"

    for (( i=1; $i<=$COMP_CWORD; i++ )); do
	if [[ "${COMP_WORDS[i]}" != -* ]]; then
	    if [[ ${COMP_WORDS[i-1]} == @($filters|$global_args) ]]; then
		continue
	    fi
	    count=$(($count + 1))
	fi
    done

    echo $(($count - 1))
}

_hg()
{
    local cur prev cmd opts i
    # global options that receive an argument
    local global_args='--cwd|-R|--repository'

    COMPREPLY=()
    cur="$2"
    prev="$3"

    # searching for the command
    # (first non-option argument that doesn't follow a global option that
    #  receives an argument)
    for (( i=1; $i<=$COMP_CWORD; i++ )); do
	if [[ ${COMP_WORDS[i]} != -* ]]; then
	    if [[ ${COMP_WORDS[i-1]} != @($global_args) ]]; then
		cmd="${COMP_WORDS[i]}"
		break
	    fi
	fi
    done

    if [[ "$cur" == -* ]]; then
	opts=$(_hg_option_list $cmd)

	COMPREPLY=( ${COMPREPLY[@]:-} $(compgen -W '$opts' -- "$cur") )
	return
    fi

    # global options
    case "$prev" in
	-R|--repository)
	    _hg_repos
	    return
	;;
	--cwd)
	    # Stick with default bash completion
	    return
	;;
    esac

    if [ -z "$cmd" ] || [ $COMP_CWORD -eq $i ]; then
	_hg_commands
	return
    fi

    # canonicalize command name
    cmd=$(hg -q help "$cmd" 2> /dev/null | sed -e 's/^hg //; s/ .*//; 1q')

    if [ "$cmd" != status ] && [ "$prev" = -r ] || [ "$prev" = --rev ]; then
	_hg_tags
	return
    fi

    case "$cmd" in
	help)
	    _hg_commands
	;;
	export|manifest|update)
	    _hg_tags
	;;
	pull|push|outgoing|incoming)
	    _hg_paths
	    _hg_repos
	;;
	paths)
	    _hg_paths
	;;
	add)
	    _hg_status "u"
	;;
	commit)
	    _hg_status "mar"
	;;
	remove)
	    _hg_status "d"
	;;
	forget)
	    _hg_status "a"
	;;
	diff)
	    _hg_status "mar"
	;;
	revert)
	    _hg_status "mard"
	;;
	clone)
	    local count=$(_hg_count_non_option)
	    if [ $count = 1 ]; then
		_hg_paths
	    fi
            _hg_repos
	;;
	debugindex|debugindexdot)
	    COMPREPLY=(${COMPREPLY[@]:-} $( compgen -f -X "!*.i" -- "$cur" ))
	;;
	debugdata)
	    COMPREPLY=(${COMPREPLY[@]:-} $( compgen -f -X "!*.d" -- "$cur" ))
	;;
    esac

}

complete -o bashdefault -o default -F _hg hg 2> /dev/null \
    || complete -o default -F _hg hg