_migration() {
    local cur prev opts
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"
    opts="upgrade downgrade"

    if [[ ${COMP_CWORD} -eq 1 ]]; then
        COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
    elif [[ ${COMP_CWORD} -eq 2 ]]; then
        if [[ "${prev}" == "upgrade" ]]; then
            COMPREPLY=( $(compgen -W "head" -- ${cur}) )
        elif [[ "${prev}" == "downgrade" ]]; then
            COMPREPLY=( $(compgen -W "base" -- ${cur}) )
        fi
    fi
}
complete -F _migration migration
