#compdef cargo

autoload -U regexp-replace

_cargo() {
    local curcontext="$curcontext" ret=1
    local -a command_scope_spec common jobs parallel features manifest msgfmt triple target registry
    local -a state line state_descr # These are set by _arguments
    typeset -A opt_args

    common=(
        '(-q --quiet)*'{-v,--verbose}'[use verbose output]'
        '(-q --quiet -v --verbose)'{-q,--quiet}'[no output printed to stdout]'
        '-Z+[pass unstable (nightly-only) flags to cargo]: :_cargo_unstable_flags'
        '--offline[run without accessing the network]'
        '--frozen[require that Cargo.lock and cache are up-to-date]'
        '--locked[require that Cargo.lock is up-to-date]'
        '--color=[specify colorization option]:coloring:(auto always never)'
        '(- 1 *)'{-h,--help}'[show help message]'
    )

    # leading items in parentheses are an exclusion list for the arguments following that arg
    # See: http://zsh.sourceforge.net/Doc/Release/Completion-System.html#Completion-Functions
    #   - => exclude all other options
    #   1 => exclude positional arg 1
    #   * => exclude all other args
    #   +blah => exclude +blah
    _arguments -s -S -C $common \
        '(- 1 *)--list[list installed commands]' \
        '(- 1 *)--explain=[provide a detailed explanation of an error message]:error code' \
        '(- 1 *)'{-V,--version}'[show version information]' \
        '(+beta +nightly)+stable[use the stable toolchain]' \
        '(+stable +nightly)+beta[use the beta toolchain]' \
        '(+stable +beta)+nightly[use the nightly toolchain]' \
        '1: :_cargo_cmds' \
        '*:: :->args'

    # These flags are mutually exclusive specifiers for the scope of a command; as
    # they are used in multiple places without change, they are expanded into the
    # appropriate command's `_arguments` where appropriate.
    command_scope_spec=(
        '(--bin --example --test --lib)--bench=[specify benchmark name]: :_cargo_benchmark_names'
        '(--bench --bin --test --lib)--example=[specify example name]:example name:_cargo_example_names'
        '(--bench --example --test --lib)--bin=[specify binary name]:binary name'
        '(--bench --bin --example --test)--lib=[specify library name]:library name'
        '(--bench --bin --example --lib)--test=[specify test name]:test name'
    )

    jobs=(
        '(-j --jobs)'{-j+,--jobs=}'[specify number of parallel jobs]:jobs [# of CPUs]'
    )

    parallel=(
        "${jobs[@]}"
        '--keep-going[do not abort build on first build error]'
    )

    features=(
        '(--all-features)'{-F+,--features=}'[specify features to activate]:feature'
        '(--features -F)--all-features[activate all available features]'
        "--no-default-features[don't build the default features]"
    )

    msgfmt='--message-format=[specify error format]:error format [human]:(human json short)'
    triple='--target=[specify target triple]:target triple:_cargo_target_triple'
    target='--target-dir=[specify directory for all generated artifacts]:directory:_directories'
    manifest='--manifest-path=[specify path to manifest]:path:_directories'
    registry='--registry=[specify registry to use]:registry'

    case $state in
        args)
            curcontext="${curcontext%:*}-${words[1]}:"
            case ${words[1]} in
                add)
                    _arguments -s -A "^--" $common $manifest $registry \
                        {-F+,--features=}'[specify features to activate]:feature' \
                        "--default-features[enable the default features]" \
                        "--no-default-features[don't enable the default features]" \
                        "--optional[mark the dependency as optional]" \
                        "--no-optional[mark the dependency as required]" \
                        "--dev[add as a dev dependency]" \
                        "--build[add as a build dependency]" \
                        "--target=[add as a dependency to the given target platform]" \
                        "--rename=[rename the dependency]" \
                        "--dry-run[don't actually write the manifest]" \
                        '--branch=[branch to use when adding from git]:branch' \
                        '--git=[specify URL from which to add the crate]:url:_urls' \
                        '--path=[local filesystem path to crate to add]: :_directories' \
                        '--rev=[specific commit to use when adding from git]:commit' \
                        '--tag=[tag to use when adding from git]:tag' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]' \
                        '1: :_guard "^-*" "crate name"' \
                        '*:args:_default'
                        ;;
                bench)
                    _arguments -s -A "^--" $common $jobs $features $msgfmt $triple $target $manifest \
                        "${command_scope_spec[@]}" \
                        '--all-targets[benchmark all targets]' \
                        "--no-run[compile but don't run]" \
                        '(-p --package)'{-p+,--package=}'[specify package to run benchmarks for]:package:_cargo_package_names' \
                        '--exclude=[exclude packages from the benchmark]:spec' \
                        '--no-fail-fast[run all benchmarks regardless of failure]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]' \
                        '1: :_guard "^-*" "bench name"' \
                        '*:args:_default'
                        ;;

                build | b)
                    _arguments -s -S $common $parallel $features $msgfmt $triple $target $manifest \
                        '--all-targets[equivalent to specifying --lib --bins --tests --benches --examples]' \
                        "${command_scope_spec[@]}" \
                        '(-p --package)'{-p+,--package=}'[specify package to build]:package:_cargo_package_names' \
                        '--release[build in release mode]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]'
                        ;;

                check | c)
                    _arguments -s -S $common $parallel $features $msgfmt $triple $target $manifest \
                        '--all-targets[equivalent to specifying --lib --bins --tests --benches --examples]' \
                        "${command_scope_spec[@]}" \
                        '(-p --package)'{-p+,--package=}'[specify package to check]:package:_cargo_package_names' \
                        '--release[check in release mode]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]'
                        ;;

                clean)
                    _arguments -s -S $common $triple $target $manifest \
                        '(-p --package)'{-p+,--package=}'[specify package to clean]:package:_cargo_package_names' \
                        '--release[clean release artifacts]' \
                        '--doc[clean just the documentation directory]'
                        ;;

                doc | d)
                    _arguments -s -S $common $parallel $features $msgfmt $triple $target $manifest \
                        '--no-deps[do not build docs for dependencies]' \
                        '--document-private-items[include non-public items in the documentation]' \
                        '--open[open docs in browser after the build]' \
                        '(-p --package)'{-p+,--package=}'[specify package to document]:package:_cargo_package_names' \
                        '--release[build artifacts in release mode, with optimizations]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]'
                        ;;

                fetch)
                    _arguments -s -S $common $triple $manifest
                        ;;

                fix)
                    _arguments -s -S $common $parallel $features $msgfmt $triple $target $manifest \
                        "${command_scope_spec[@]}" \
                        '--broken-code[fix code even if it already has compiler errors]' \
                        '--edition[fix in preparation for the next edition]' \
                        '--edition-idioms[fix warnings to migrate to the idioms of an edition]' \
                        '--allow-no-vcs[fix code even if a VCS was not detected]' \
                        '--allow-dirty[fix code even if the working directory is dirty]' \
                        '--allow-staged[fix code even if the working directory has staged changes]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]'
                ;;

                generate-lockfile)
                    _arguments -s -S $common $manifest
                        ;;

                help)
                    _cargo_cmds
                        ;;
                info)
                    _arguments -s -A "^--" $common $registry \
                        '--index=[specify registry index]:index' \
                        '*: :_guard "^-*" "crate"'
                        ;;

                init)
                    _arguments -s -S $common $registry \
                        '--lib[use library template]' \
                        '--edition=[specify edition to set for the crate generated]:edition:(2015 2018 2021)' \
                        '--vcs=[initialize a new repo with a given VCS]:vcs:(git hg pijul fossil none)' \
                        '--name=[set the resulting package name]:name' \
                        '1:path:_directories'
                        ;;

                install)
                    _arguments -s -S $common $parallel $features $triple $registry \
                        '(-f --force)'{-f,--force}'[force overwriting of existing crates or binaries]' \
                        '--bin=[only install the specified binary]:binary' \
                        '--branch=[branch to use when installing from git]:branch' \
                        '--debug[Build in debug mode (with the "dev" profile) instead of release mode]' \
                        '--example=[install the specified example instead of binaries]:example:_cargo_example_names' \
                        '--git=[specify URL from which to install the crate]:url:_urls' \
                        '--path=[local filesystem path to crate to install]: :_directories' \
                        '--rev=[specific commit to use when installing from git]:commit' \
                        '--root=[directory to install packages into]: :_directories' \
                        '--tag=[tag to use when installing from git]:tag' \
                        '--version=[version to install from crates.io]:version' \
                        '--list[list all installed packages and their versions]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]' \
                        '*: :_guard "^-*" "crate"'
                        ;;

                locate-project)
                    _arguments -s -S $common $manifest \
                        '--message-format=[specify output representation]:output representation [json]:(json plain)' \
                        '--workspace[locate Cargo.toml of the workspace root]'
                        ;;

                login)
                    _arguments -s -S $common $registry \
                        '*: :_guard "^-*" "token"'
                        ;;

                metadata)
                    _arguments -s -S $common $features $manifest \
                        "--no-deps[output information only about the root package and don't fetch dependencies]" \
                        '--format-version=[specify format version]:version [1]:(1)'
                        ;;

                new)
                    _arguments -s -S $common $registry \
                        '--lib[use library template]' \
                        '--vcs:initialize a new repo with a given VCS:(git hg none)' \
                        '--name=[set the resulting package name]'
                        ;;

                owner)
                    _arguments -s -S $common $registry \
                        '(-a --add)'{-a,--add}'[specify name of a user or team to invite as an owner]:name' \
                        '--index=[specify registry index]:index' \
                        '(-l --list)'{-l,--list}'[list owners of a crate]' \
                        '(-r --remove)'{-r,--remove}'[specify name of a user or team to remove as an owner]:name' \
                        '--token=[specify API token to use when authenticating]:token' \
                        '*: :_guard "^-*" "crate"'
                        ;;

                package)
                    _arguments -s -S $common $parallel $features $triple $target $manifest $registry \
                        '--index=[specify registry index]:index' \
                        '(-l --list)'{-l,--list}'[print files included in a package without making one]' \
                        '--no-metadata[ignore warnings about a lack of human-usable metadata]' \
                        '--allow-dirty[allow dirty working directories to be packaged]' \
                        "--no-verify[don't build to verify contents]"
                        ;;

                pkgid)
                    _arguments -s -S $common $manifest \
                        '(-p --package)'{-p+,--package=}'[specify package to get ID specifier for]:package:_cargo_package_names' \
                        '*: :_guard "^-*" "spec"'
                        ;;

                publish)
                    _arguments -s -S $common $parallel $features $triple $target $manifest $registry \
                        '--index=[specify registry index]:index' \
                        '--allow-dirty[allow dirty working directories to be packaged]' \
                        "--no-verify[don't verify the contents by building them]" \
                        '--dry-run[perform all checks without uploading]'
                        ;;

                remove | rm)
                    _arguments -s -A "^--" $common $manifest \
                        "--dev[remove as a dev dependency]" \
                        "--build[remove as a build dependency]" \
                        "--target=[remove as a dependency from the given target platform]" \
                        "--dry-run[don't actually write the manifest]" \
                        '(-p --package)'{-p+,--package=}'[package to remove from]:package:_cargo_package_names' \
                        '1: :_guard "^-*" "crate name"' \
                        '*:args:_default'
                        ;;

                run | r)
                    _arguments -s -S $common $parallel $features $msgfmt $triple $target $manifest \
                        '--example=[name of the bin target]:name:_cargo_example_names' \
                        '--bin=[name of the bin target]:name' \
                        '(-p --package)'{-p+,--package=}'[specify package with the target to run]:package:_cargo_package_names' \
                        '--release[build in release mode]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]' \
                        '*: :_default'
                        ;;

                rustc)
                    _arguments -s -S $common $parallel $features $msgfmt $triple $target $manifest \
                        '(-p --package)'{-p+,--package=}'[specify package to build]:package:_cargo_package_names' \
                        '--profile=[specify profile to build the selected target for]:profile' \
                        '--release[build artifacts in release mode, with optimizations]' \
                        "${command_scope_spec[@]}" \
                        '--ignore-rust-version[Ignore rust-version specification in packages]' \
                        '*: : _dispatch rustc rustc -default-'
                        ;;

                rustdoc)
                    _arguments -s -S $common $parallel $features $msgfmt $triple $target $manifest \
                        '--document-private-items[include non-public items in the documentation]' \
                        '--open[open the docs in a browser after the operation]' \
                        '(-p --package)'{-p+,--package=}'[specify package to document]:package:_cargo_package_names' \
                        '--release[build artifacts in release mode, with optimizations]' \
                        "${command_scope_spec[@]}" \
                        '--ignore-rust-version[Ignore rust-version specification in packages]' \
                        '*: : _dispatch rustdoc rustdoc -default-'
                        ;;

                search)
                    _arguments -s -S $common $registry \
                        '--index=[specify registry index]:index' \
                        '--limit=[limit the number of results]:results [10]' \
                        '*: :_guard "^-*" "query"'
                        ;;

                test | t)
                    _arguments -s -S $common $jobs $features $msgfmt $triple $target $manifest \
                        '--test=[test name]: :_cargo_test_names' \
                        '--no-fail-fast[run all tests regardless of failure]' \
                        '--no-run[compile but do not run]' \
                        '(-p --package)'{-p+,--package=}'[package to run tests for]:package:_cargo_package_names' \
                        '--all[test all packages in the workspace]' \
                        '--release[build artifacts in release mode, with optimizations]' \
                        '1: :_cargo_test_names' \
                        '(--doc --bin --example --test --bench)--lib[only test library]' \
                        '(--lib --bin --example --test --bench)--doc[only test documentation]' \
                        '(--lib --doc --example --test --bench)--bin=[binary name]' \
                        '(--lib --doc --bin --test --bench)--example=[example name]:_cargo_example_names' \
                        '(--lib --doc --bin --example --bench)--test=[test name]' \
                        '(--lib --doc --bin --example --test)--bench=[benchmark name]' \
                        '--ignore-rust-version[Ignore rust-version specification in packages]' \
                        '*: :_default'
                        ;;

                tree)
                    _arguments -s -S $common $features $triple $manifest \
                        '(-p --package)'{-p+,--package=}'[package to use as the root]:package:_cargo_package_names' \
                        '(-i --invert)'{-i+,--invert=}'[invert the tree for the given package]:package:_cargo_package_names' \
                        '--prefix=[line prefix]:prefix:(depth indent none)' \
                        '--no-dedupe[repeat shared dependencies]' \
                        '(-d --duplicates)'{-d,--duplicates}'[packages with multiple versions]' \
                        '--charset=[utf8 or ascii]:charset:(utf8 ascii)' \
                        '(-f --format)'{-f,--format=}'[format string]:format' \
                        '(-e --edges)'{-e,--edges=}'[edge kinds]:kind:(features normal build dev all no-dev no-build no-normal)' \
                        ;;

                uninstall)
                    _arguments -s -S $common \
                        '(-p --package)'{-p+,--package=}'[specify package to uninstall]:package:_cargo_package_names' \
                        '--bin=[only uninstall the specified binary]:name' \
                        '--root=[directory to uninstall packages from]: :_files -/' \
                        '*:crate:_cargo_installed_crates -F line'
                        ;;

                update)
                    _arguments -s -S $common $manifest \
                        '--aggressive=[force dependency update]' \
                        '--recursive=[force dependency update]' \
                        "--dry-run[don't actually write the lockfile]" \
                        '(-p --package)'{-p+,--package=}'[specify package to update]:package:_cargo_package_names' \
                        '--precise=[update single dependency to precise release]:release' \
                        '*:package:_cargo_package_names'
                        ;;

                version)
                    _arguments -s -S $common
                        ;;

                yank)
                    _arguments -s -S $common $registry \
                        '--version=[specify yank version]:version' \
                        '--undo[undo a yank, putting a version back into the index]' \
                        '--index=[specify registry index to yank from]:registry index' \
                        '--token=[specify API token to use when authenticating]:token' \
                        '*: :_guard "^-*" "crate"'
                        ;;
                *)
                    # allow plugins to define their own functions
                    if ! _call_function ret _cargo-${words[1]}; then
                        # fallback on default completion for unknown commands
                        _default && ret=0
                    fi
                    (( ! ret ))
                ;;
            esac
            ;;
    esac
}

_cargo_unstable_flags() {
    local flags
    flags=( help ${${${(M)${(f)"$(_call_program flags cargo -Z help)"}:#*--*}/ #-- #/:}##*-Z } )
    _describe -t flags 'unstable flag' flags
}

_cargo_installed_crates() {
    local expl
    _description crates expl 'crate'
    compadd "$@" "$expl[@]" - ${${${(f)"$(cargo install --list)"}:# *}%% *}
}

_cargo_cmds() {
    local -a commands
    # This uses Parameter Expansion Flags, which are a built-in Zsh feature.
    # See more: http://zsh.sourceforge.net/Doc/Release/Expansion.html#Parameter-Expansion-Flags
    # and       http://zsh.sourceforge.net/Doc/Release/Expansion.html#Parameter-Expansion
    #
    # # How this work?
    #
    # First it splits the result of `cargo --list` at newline, then it removes the first line.
    # Then it removes indentation (4 whitespaces) before each items. (Note the x## pattern [1]).
    # Then it replaces those spaces between item and description with a `:`
    #
    # [1]: https://github.com/zsh-users/zsh-completions/blob/master/zsh-completions-howto.org#patterns
    commands=( ${${${(M)"${(f)$(_call_program commands cargo --list)}":#    *}/ ##/}/ ##/:} )
    _describe -t commands 'command' commands
}

_cargo_target_triple() {
    local -a result

    if (( $+commands[rustup] )); then
        result=( ${(f)"$(rustup target list --installed)"} )
    else
        result=( ${(f)"$(rustc --print target-list)"} )
    fi

    _describe 'target triple' result
}

#FIXME: Disabled until fixed
#gets package names from the manifest file
_cargo_package_names() {
    _message -e packages package
}

# Extracts the values of "name" from the array given in $1 and shows them as
# command line options for completion
_cargo_names_from_array() {
    local manifest=$(cargo locate-project --message-format plain)
    if [[ -z $manifest ]]; then
        return 0
    fi

    local last_line
    local -a names;
    local in_block=false
    local block_name=$1
    names=()
    while read -r line; do
        if [[ $last_line == "[[$block_name]]" ]]; then
            in_block=true
        else
            if [[ $last_line =~ '\s*\[\[.*' ]]; then
                in_block=false
            fi
        fi

        if [[ $in_block == true ]]; then
            if [[ $line =~ '\s*name\s*=' ]]; then
                regexp-replace line '^\s*name\s*=\s*|"' ''
                names+=( "$line" )
            fi
        fi

        last_line=$line
    done < "$manifest"
    _describe "$block_name" names

}

#Gets the test names from the manifest file
_cargo_test_names() {
    _cargo_names_from_array "test"
}

#Gets the bench names from the manifest file
_cargo_benchmark_names() {
    _cargo_names_from_array "bench"
}

_cargo_example_names() {
    if [[ -d examples ]]; then
        local -a files=(${(@f)$(echo examples/*.rs(:t:r))})
        _values 'example' "${files[@]}"
    fi
}

_cargo
