#compdef conda
#description:conda package manager
#
# ZSH Completion for conda (http://conda.pydata.org/)
#
# Author: Valentin Haenel <valentin@haenel.co>  (https://github.com/esc/)
# Licence: WTFPL (http://sam.zoy.org/wtfpl/)
# Version: 0.3-dev
#
# This completion depends on Python for a json parser, sorry. Unfortunately
# there is no such thing in zsh (yet).
#
# To use this completion drop it somewhere in you '$fpath', e.g.:
#
#     $ git clone $CLONEURL
#     $ fpath+=$PWD/conda-zsh-completion
#     $ compinit conda
#
# To activate the completion cache for packages, add the following to your
# '.zshrc':
#
#     zstyle ':completion::complete:*' use-cache 1
#
# To forcefully clean the completion cache, look in '~/.zcompcache' and remove
# file starting with 'conda_'.
#
# CHANGELOG
# ---------
#
# * v0.3
#
#   *
#
# * v0.2
#
#   * completion cache for packages
#   * complete all channels defined in .condarc
#
# * v0.1
#
#   * inital release
#
# TODO
# ----
#
# * Subcommand grouping is still alpha.
# * Completion for 'conda config' is somewhat incomplete
#   * 'conda config --get' only completes a single option.
#   * 'conda config --add' completes only the key, not the value.
#   * Similar issues for other options.
# * Example of activating cache only for conda completion
# * Make cache policy configurable
# * Completion for version numbers is entirely missing:
#   e.g. 'conda install numpy==<TAB>' should complete all available versions of
#   numpy.
# * Completion for 'conda remove' isn't aware of any environment given on the
#   command line and so completion of existing packages may be incorrect if
#   that option is given.
# * Completion of channels is rudimentary, conda isn't queried for defined
#   channels.
# * Configuration for external commands, e.g. 'build' is entirely missing.
# * Seperate boolean and list options for 'config'
# * Properly handle package specs on the command line

local state line context
local -A opt_args

__conda_envs(){
    local -a envs
    envs=($( conda info -e | sed "1,2d" | cut -f1 -d' '))
    _describe -t envs 'conda environments' envs
}

__conda_packages_installed(){
    zstyle ":completion:${curcontext}:" cache-policy _conda_caching_policy
    local -a installed_packages
    if _cache_invalid conda_installed_packages || ! _retrieve_cache conda_installed_packages ; then
        installed_packages=($(conda list | sed 1,2d | cut -f1 -d' '))
        _store_cache conda_installed_packages installed_packages
    fi
    _describe -t installed_packages 'installed packages' installed_packages
}

__conda_packages_search(){
    zstyle ":completion:${curcontext}:" cache-policy _conda_caching_policy
    local -a available_packages
    if _cache_invalid conda_available_packages || ! _retrieve_cache conda_available_packages ; then
        available_packages=($(conda search --use-index-cache --json |
        python -c "
import json, sys
parsed = json.load(sys.stdin)
for k in parsed.keys():
    print(k)
        "))
        _store_cache conda_available_packages available_packages
    fi
    _describe -t available_packages 'available packages' available_packages
}

__conda_config_values(){
    local -a config_values search_term
    search_term="$1"
    config_values=($(conda config --json --get "$search_term"|
        python -c "
import json, sys
try:
    values = json.load(sys.stdin)['get']['$search_term']
    for v in values:
        print(v)
except KeyError:
    pass
        "))
    print -l $config_values
}

__conda_config_values_describe(){
    local -a config_values search_term
    search_term="$1"
    config_values=($( __conda_config_values $search_term ))
    if [ "${#config_values}" == 0 ] ; then
        _message "no values found for '$search_term'!"
    else
        _describe -t config_values 'configuration values' config_values
    fi
}

__conda_boolean_config_value(){
    local -a config_values
    config_values=(True False)
    _describe -t config_values 'boolean configuration values' config_values

}

__conda_channels(){
    local -a channels
    channels=$( __conda_config_values "channels" )
    channels+=(system)
    _describe -t channels 'conda channels' channels
}

boolean_config_options=(
    'add_binstar_token'
    'always_yes'
    'allow_softlinks'
    'changeps1'
    'use_pip'
    'offline'
    'binstar_upload'
    'binstar_personal'
    'show_channel_urls'
    'allow_other_channels'
    'ssl_verify'
    )

list_config_options=(
    'channels'
    'disallow'
    'create_default_packages'
    'track_features'
    'envs_dirs'
    )

__boolean_conda_config(){
    _describe -t boolean_config_options 'boolean keys' boolean_config_options
}

__list_conda_config(){
    _describe -t list_config_options 'list keys' list_config_options
}

config_options=($boolean_config_options $list_config_options)

__conda_config(){
    _describe -t config_options 'conda configuration keys' config_options
}


__conda_commands(){
    local -a package maint environment help config special
    package=(
        search:'Search for packages and display their information.'
        install:'Install a list of packages into a specified conda environment.'
    )
    maint=(
        update:'Update conda packages.'
        clean:'Remove unused packages and caches.'
    )
    environment=(
        info:'Display information about current conda install.'
        create:'Create a new conda environment from a list of specified packages.'
        list:'List linked packages in a conda environment.'
        remove:'Remove a list of packages from a specified conda environment.'
        uninstall:'Alias for conda remove'
    )
    help=(
        help:'Displays a list of available conda commands and their help strings.'
    )
    config=(
        config:'Modify configuration values in .condarc.'
    )
    special=(
        run:'Launches an application installed with Conda.'
        init:'Initialize conda into a regular environment. (DEPRECATED)'
        package:'Low-level conda package utility. (EXPERIMENTAL)'
        bundle:'Create or extract a "bundle package" (EXPERIMENTAL)'
    )
    _describe -t package_commands "package commands" package
    _describe -t maint_commands "maint commands" maint
    _describe -t environment_commands "environment commands" environment
    _describe -t help_commands "help commands"  help
    _describe -t config_commands "config commands"  config
    _describe -t special_commands "special commands"  special
}

_conda_caching_policy() {
  local -a oldp
  oldp=( "$1"(Nmh+12) ) # 12 hour
  (( $#oldp ))
}

local -a opts
opts=(
    '(-h --help)'{-h,--help}'[show this help message and exit]'
    '(-V --version)'{-V,--version}'[show program''s version number and exit]'
)

_arguments -C $opts \
           ': :->command' \
           '*:: :->subcmd'

help_opts=(
    '(-h --help)'{-h,--help}'[show the help message and exit]' \
    )

json_opts=(
    '--json[report all output as json.]' \
    )

env_opts=(
    '(-n --name)'{-n,--name}'[name of environment]:environment:__conda_envs' \
    '(-p --prefix)'{-p,--prefix}'[full path to environment prefix]:path:_path_files' \
    )

channel_opts=(
    '(-c --channel)'{-c,--channel}'[additional channel to search for packages]:channel:__conda_channels'\
    '--override-channels [do not search default or .condarc channels]' \
    '--use-index-cache[use cache of channel index files]' \
    '--use-local[use locally built packages]' \
    )

install_opts=(
    '(-y --yes)'{-y,--yes}'[do not ask for confirmation]' \
    '--dry-run[only display what would have been done]' \
    '(-f --force)'{-f,--force}'[force install]' \
    '--file[read package versions from file]:file:_path_files' \
    '--no-deps[do not install dependencies]' \
    '(-m --mkdir)'{-m,--mkdir}'[create prefix directory if necessary]' \
    '--offline[offline mode, don''t connect to internet]' \
    '--no-pin[ignore pinned file]' \
    '(-q --quiet)'{-q,--quiet}'[do not display progress bar]'\
    '--copy[Install all packages using copies instead of hard- or soft-linking]' \
    '--alt-hint[Use an alternate algorithm to generate an unsatisfiable hint]' \
    )

case $state in
(command)
    __conda_commands
    ;;
(subcmd)
    case ${line[1]} in
    (info)
        _arguments -C $help_opts \
                      '--json[report all output as json.]' \
                      '(-a --all)'{-a,--all}'[show all information, (environments, license, and system information]' \
                      '(-e --envs)'{-e,--envs}'[list all known conda environments]' \
                      '(-l --license)'{-l,--license}'[display information about local conda licenses list]' \
                      '(-s --system)'{-s,--system}'[list environment variables]' \
                      '--root[display root environment path]'
        ;;
    (help)
        _arguments -C $help_opts \
                      '*:commands:__conda_commands' \
        ;;
    (list)
        _arguments -C $help_opts \
                      $env_opts \
                      $json_opts \
                      '(-c --canonical)'{-c,--canonical}'[output canonical names of packages only]' \
                      '(-e --export)'{-e,--export}'[output requirement string only]' \
                      '(-r --revisions)'{-r,--revision}'[list the revision history and exit]' \
                      '--no-pip[Do not include pip-only installed packages]' \
                      '*:regex:' \
        ;;
    (search)
        _arguments -C $help_opts \
                      $env_opts \
                      $json_opts \
                      $channel_opts \
                      '(-c --canonical)'{-c,--canonical}'[output canonical names of packages only]' \
                      '--unknown[use index metadata from the local package cache]' \
                      '(-o --outdated)'{-o,--outdated}'[only display installed but outdated packages]' \
                      '(-v --verbose)'{-v,--verbose}'[Show available packages as blocks of data]' \
                      '--platform[Search the given platform.]' \
                      '--spec[Treat regex argument as a package specification]' \
                      '*:regex:' \
        ;;
    (create)
        _arguments -C $help_opts \
                      $env_opts \
                      $install_opts \
                      $json_opts \
                      $channel_opts \
                      '--unknown[use index metadata from the local package cache]' \
                      '--clone[path to (or name of) existing local environment]' \
                      '--no-default-packages[ignore create_default_packages in condarc file]' \
                      '*:packages:__conda_packages_search' \
        ;;
    (install)
        _arguments -C $help_opts \
                      $env_opts \
                      $install_opts \
                      $json_opts \
                      $channel_opts \
                      '--revision[revert to the specified revision]:revision' \
                      '*:packages:__conda_packages_search' \
        ;;
    (update)
        _arguments -C $help_opts \
                      $env_opts \
                      $install_opts \
                      $json_opts \
                      $channel_opts \
                      '--unknown[use index metadata from the local package cache]' \
                      '--all[Update all installed packages in the environment]' \
                      '*:packages:__conda_packages_installed' \
        ;;
    (remove|uninstall)
        _arguments -C $help_opts \
                      $env_opts \
                      $json_opts \
                      $channel_opts \
                      '(-y --yes)'{-y,--yes}'[do not ask for confirmation]' \
                      '--dry-run[only display what would have been done]' \
                      '(-a --all)'{-a,--all}'[remove all packages, i.e. the entire environment]' \
                      '--features[remove features (instead of packages)]' \
                      '--no-pin[ignore pinned file]' \
                      '(-q --quiet)'{-q,--quiet}'[do not display progress bar]'\
                      '--offline[offline mode, don''t connect to internet]' \
                      '*:packages:__conda_packages_installed' \
        ;;
    (config)
        # this allows completing multiple keys whet --get is given
        local -a last_item get_opts
        last_item=$line[$CURRENT-1]
        if [[ ${line[(r)--get]} == --get ]]  && [[ ${config_options[(r)$last_item]} == $last_item ]] ; then
            get_opts=('*:keys:__conda_config')
        else
            get_opts=''
        fi
        _arguments -C $help_opts \
                      $json_opts \
                      '--system[write to the system .condarc file]' \
                      '--file[write to the given file.]:file:_path_files' \
                      '(      --add --set --remove --remove-key)--get[get the configuration value]:key:__conda_config' \
                      '(--get       --set --remove --remove-key)--add[add one configuration value to a list key]:list key:__list_conda_config:value:' \
                      '(--get --add       --remove --remove-key)--set[set a boolean key]:boolean key:__boolean_conda_config:value:__conda_boolean_config_value' \
                      '(--get --add --set          --remove-key)--remove[remove a configuration value from a list key]:list key:__list_conda_config:value:{__conda_config_values_describe '$last_item'}' \
                      '(--get --add --set --remove             )--remove-key[remove a configuration key (and all its values)]:key:__conda_config' \
                      '(-f --force)'{-f,--force}'[write to the config file using the yaml parser]' \
                      $get_opts
        ;;
    (init)
        _arguments -C $help_opts \
        ;;
    (clean)
        _arguments -C $help_opts \
                      $json_opts \
                      '(-y --yes)'{-y,--yes}'[do not ask for confirmation]' \
                      '--dry-run[only display what would have been done]' \
                      '(-i --index-cache)'{-i,--index-cache}'[remove index cache]' \
                      '(-l --lock)'{-l,--lock}'[remove all conda lock files]' \
                      '(-t --tarballs)'{-t,--tarballs}'[remove cached package tarballs]' \
                      '(-p --packages)'{-p,--packages}'[remove unused cached packages]' \
                      '(-s --source-cache)'{-s,--source-cache}'[remove files from the source cache of conda build]' \
        ;;
    (package)
        _arguments -C $help_opts \
                      $env_opts \
                      '(-w --which)'{-w,--which}'[given some path print which conda package the file came from]:path:_path_files' \
                      '(-L --ls-files)'{-L,--ls-files}'[list all files belonging to specified package]' \
                      '(-r --reset)'{-r,--reset}'[remove all untracked files and exit]' \
                      '(-u --untracked)'{-u,--untracked}'[display all untracked files and exit]' \
                      '--pkg-name[package name of the created package]:pkg_name:' \
                      '--pkg-version[package version of the created package]:pkg_version:' \
                      '--pkg-build[package build number of the created package]:pkg_build:' \
        ;;
    (bundle)
        _arguments -C $help_opts \
                      $env_opts \
                      $json_opts \
                      '(-c --create)'{-c,--create}'[create bundle]' \
                      '(-x --extract)'{-x,--extract}'[extact bundle located at path]:path:_path_files' \
                      '--metadump[dump metadata of bundle at path]:path:_path_files' \
                      '(-q --quiet)'{-q,--quiet}'[do not display progress bar]'\
                      '--bundle-name[name of bundle]:NAME:' \
                      '--data-path[path to data to be included in bundle]:path:_path_files' \
                      '--extra-meta[path to json file with additional meta-data no]:path:_path_files' \
                      '--no-env[no environment]' \
        ;;
    esac
    ;;
esac