#compdef nix-env
#autoload

emulate -L zsh
setopt extendedglob

local context state state_descr line
typeset -A opt_args

_nix-common-options

local -a _1st_arguments
_1st_arguments=(
    {--install,-i}"[Install package]"
    {--upgrade,-u}"[Upgrade package]"
    {--uninstall,-e}"[Uninstall package]"
    --set-flag'[Modify meta attribute of installed package]:Name:->flag_name:Value:->flag_value:*:Packages: _nix_installed_packages'
    {--query,-q}"[List information about derivations]"
    {--switch-profile,-S}"[Set the current profile path]"
    "--list-generations[Print a list of all generations in the active profile]"
    "--delete-generations[Delete specified generations]"
    {--switch-generation,-G}"[Activate specified generation]"
    "--rollback[Switch to the previous generation of active profile]"
)

local -a _nix_env_common_opts=(
  $__nix_common_opts
  '(--profile -p)'{--profile,-p}'[Specify the profile to use]:Path:_nix_profiles'
  $__nix_dry_run
  '--system-filter[Only show derivations matching the specified platform]:system:_nix_systems'
)

local -a _nix_env_b=(
    '(--prebuilt-only -b)'{--prebuilt-only,-b}'[Fail if there is no pre-built binary available]'
)

local _nix_env_from_profile
_nix_env_from_profile='--from-profile[Fetch store paths from another profile]:Profile:_nix_profiles'

# Workaround:
# opt_args doesn't record plain options, so we can't use that
local expect_attr_paths=false
local opt
for opt in $words; do
    case $opt in
        --attr|-[^-]#A[^-]#)
            expect_attr_paths=true
            ;;
    esac
done

local -a command_options=()
# Look for a main option in the input and setup command_options appropriately
for opt in $words; do
case "$opt" in
    --install|-[^-]#i[^-]#)
        command_options=(
            $_nix_env_common_opts
            '(--attr -A)'{--attr,-A}'[Specify packages by attribute path instead of name]'
            $_nix_env_b
            $_nix_env_from_profile
            '(--preserve-installed -P)'{--preserve-installed,-P}'[Do not remove derivations with the same name]'
            '(--remove-all -r)'{--remove-all,-r}'[Remove all previously installed packages prior to installing]'
            '*:Package:{if $expect_attr_paths; then
                           _nix_complete_attr_paths;
                        else
                           _nix_installed_packages "(use -A/--attr to access all packages)"
                        fi}')
        break
        ;;
    --upgrade|-[^-]#u[^-]#)
        command_options=(
            $_nix_env_common_opts
            $_nix_env_b
            $_nix_env_from_profile
            '(-lt -leq -eq --always)--lt[Upgrade derivations with newer versions (default)]'
            '(-lt -leq -eq --always)--leq[Upgrade derivations with the same or newer version]'
            '(-lt -leq -eq --always)--eq[Upgrade derivations with equivalent versions]'
            '(-lt -leq -eq --always)--always[Upgrade even if version number decreases]'
            '*:Packages: _nix_installed_packages')
        break
        ;;
    --uninstall|-[^-]#e[^-]#)
        command_options=(
            ${_nix_env_common_opts:#'(--attr -A)'*}
            '*::Packages: _nix_installed_packages')
        break
        ;;
    --set-flag)
        break
        ;;
    --query|-[^-]#q[^-]#)
        command_options=(
            $_nix_env_common_opts
            '(--available -a)'{--available,-a}'[display all installable derivations]'
            $_nix_env_b
            '(--status -s)'{--status,-s}'[print status of derivation]'
            '(--attr-path -P)'{--attr-path,-P}'[print attribute path of derivations]'
            '--no-name[suppress printing of name attribute]'
            '(--compare-versions -c)'{--compare-versions,-c}'[compare installed and available version]'
            '--system[print system attribute]'
            '--drv-path[print store derivation path]'
            '--out-path[print output path]'
            '--description[print description]'
            '--xml[print output as xml]'
            '--json[print output as json]'
            '--meta[Print all meta attributes: only available with --xml]')
        break
        ;;
    --switch-profile|-[^-]#S[^-]#)
        command_options=($_nix_env_common_opts ':Profile:_nix_profiles')
        break
        ;;
    --delete-generations)
        command_options=($_nix_env_common_opts '*::Generations:_nix_generations')
        break
        ;;
    --switch-generation|-[^-]#G[^-]#)
        command_options=($_nix_env_common_opts '::Generations:_nix_generations')
        break
        ;;
    --list-generations)
        command_options=($_nix_env_common_opts)
        break
        ;;
    --[^-]*|-[^-]#)
        # Complete common options if the user has started writing something else
        command_options=($_nix_env_common_opts)
        ;;
esac
done

# Let _arguments handle the rest, with _1st_arguments being mutually exclusive
# under the main_options group
_arguments -s $command_options \
           '*'{--file,-f}'[Specify Nix expression used to obtain derivations]:Path to file:_nix_path'\
           '-f.[evaluate ./default.nix rather than the default]' \
           $__nix_boilerplate_opts \
           - '(main_options)' \
           $_1st_arguments

# Handle the --set-flag option
case $state in
    flag_name)
        local -a _set_flag_attrs=(
            'priority:Resolve package name conflicts'
            'keep:Prevent package from being upgraded'
            'active:Package is symlinked to profile'
        )
        _describe -t commands "Package flag" _set_flag_attrs
        break
        ;;
    flag_value)
        local flag_name=$words[$(($CURRENT - 1))]
        case $flag_name in
            priority)
                _message "Number, lower priority values denote a higher priority"
                break
                ;;
            keep)
                _values "TF" true false
                break
                ;;
            active)
                _values "TF" true false
                break
                ;;
        esac
        ;;
esac