_nix

#compdef nix
#autoload

# Most information is extractable from `nix --help` and `nix command --help`

# There's four different types of options and arguments which is passed to _arguments:
#
# 1. The main commands which we get from nix --help
#
# 2. The common options which are hardcoded for the moment
#
# 3. The regular arguments for the active command. The number and types of these
# arguments are taken from the "Usage:" line of `nix command --help`, $state
# is set to the name of the argument via a `->$arg` action, which are further
# handled in the case statement at the end of the file.
#
# 4. The options for the active command. These are taken from the "Flags:"
# section of `nix command --help`, and the option arguments are handled in the
# same way as in 3.

# Ensure access to eg. _nix_complete_attr_paths
_nix-common-options

setopt extendedglob

local context state state_descr line
typeset -A opt_args

type nix &> /dev/null || return

local -a common_options
common_options=("(--debug)"--debug"[enable debug output]"
                "(-h --help)"{-h,--help}"[show usage information]"
                "(--help-config)"--help-config"[show configuration options]"
                "*--builders[The machine to build on]:Machine eg. 'root@mac x86_64-darwin': "
                "--store[Select a nix store]:Store:->STORE-URI"
                "*--option[set a Nix configuration option (overriding nix.conf)]:Option:->nixoption:Value:->nixoptionvalue"
                "(--quiet)"--quiet"[decrease verbosity level]"
                "*"{-v,--verbose}"[increase verbosity level]"
                "(--version)"--version"[show version information]")

local -a main_commands
# Extract the commands with descriptions
# like ('command:some description' 'run:run some stuff'):
#
# Split the help output on newlines, and remove all lines which doesn't start
# with two spaces and a lower case character,
main_commands=(${${(f)"$(nix --help)"}:#^(  [a-z]*)})
# Strip the leading spaces
main_commands=(${main_commands/#  /})
# And replace the whitespace separating the command and description with ':'
main_commands=(${main_commands/  ##/:})

# Add commands to an associative array for easy lookup
local -A command_lookup
local command_description
for command_description in $main_commands; {
    local command=${command_description%%\:*}
    command_lookup[$command]=1
}

local -a command_options=()
local -a command_arguments=()
# Setup the correct command_arguments and command_options depending on which
# command we've typed
local word
for word in $words; do

    # Check if we're in a valid command
    if [[ $command_lookup[$word] == 1 ]]; then
        local -a help=(${(f)"$(nix $word --help)"})
        # Extract an array describing the possible arguments to the command eg.
        # (NAR PATH) for cat-nar or (INSTALLABLES) for run
        local -a args=(${=help[1]//(*<FLAGS>|.|<|>)/})
        # And add the appropriate _argument spec
        local arg
        for arg in $args; do
            # Prefix * if $arg ends in 'S', ie. is plural
            local plural="${${(M)arg:#*S}:+*}"
            command_arguments+=("$plural:$arg:->$arg")
        done

        # Extract the lines containing the option descriptions
        local -a option_descriptions
        option_descriptions=(${help:#^(  -*|  ??? -*)})

        local option_description
        local option
        for option_description in $option_descriptions; do

            # Extract the description by stripping everything up to the the last
            # two consecutive spaces
            local description=${option_description/(*  )/}
            # Remove the description from the help line
            local option_spec=${option_description%$description}

            # Extract the options by stripping everything starting at '<' (and ',')
            local -a option_group=(${=option_spec//(<*|,)/})
            # Extract any arguments, by stripping the options (and <|>)
            local -a option_args=(${=option_spec//(*$option_group[${#option_group}]|<|>)/})

            local ACTIONS=""
            for arg in $option_args; do
                # Prefix with ':*' if $arg ends in 'S', ie. is plural
                local plural="${${(M)arg:#*S}:+:*}"
                ACTIONS+="$plural:$arg:->$arg"
            done

            for option in $option_group; do
                # Handle `run --keep/--unset` manually as there's ambiguity the NAME argument
                if [[ $word == run && -z ${option:#(-k|--keep|-u|--unset)} ]]; then
                    command_options+=("*${option}[$description]:Environment Variable:_parameters")
                elif [[ "$word" == add-to-store \
                            && "$option" == (-n|--name) ]]; then
                    # Another <NAME> ambiguity
                    command_options+=("($option_group)"$option"[$description]:->store-name")
                elif [[ $option == (-I|--include) ]]; then
                    # Special handling of --include due to type ambiguity
                    command_options+=("*${option}[$description]:Includes:->INCLUDE")
                elif [[ -z ${option:#(--arg|--argstr|-f|--file)} ]]; then
                    # Repeatable options
                    command_options+=("*${option}[$description]"$ACTIONS)
                else
                    # Default to mutually exclusive non-repeatable options
                    command_options+=("($option_group)"$option"[$description]"$ACTIONS)
                fi
            done
        done

        break
    fi
done

_arguments -s \
           ":Command: _describe -t main_commands Command main_commands"\
           $command_arguments \
           $common_options \
           $command_options && return 0

# Handle arguments to commands and options
case "${context%%-*} $state" in
    'argument '(INSTALLABLES|INSTALLABLE|PACKAGE|DEPENDENCY))
        # Complete attribute paths and files starting with either "/" or "./"

        # We need to prefix relative paths with ./ so nix will evaluate it as a
        # path
        local prefix='-P ./'
        local current_word=$words[$CURRENT]
        # When referencing an absolute path we can't prefix with ./
        if [[ -z ${current_word:##(/*|\~/*)} && -n $current_word ]]; then
            prefix=""
        fi
        _alternative \
            'path:Attribute path: _nix_complete_attr_paths' \
            "file:File path to package:_files ${prefix}"
        return
        ;;
    'argument '(PATH|PATHS|FILES))
        _files
        return
        ;;
    'argument NAR')
        _files -X "NAR"
        return
        ;;
    'argument REGEX')
        _message 'Regex to search for'
        return
        ;;
    'argument STRINGS')
        _message 'HASH to convert to a different base'
        return
        ;;
    'option NAME')
        _nix_complete_function_arg
        return
        ;;
    'option EXPR')
        _message 'Expression argument'
        return
        ;;
    'option STRING')
        _message 'String argument'
        return
        ;;
    'option STORE-URI')
        local -a others=(
            'local:Use /nix/store'
            'remote:Go via the Nix daemon'
            'https\://:HTTPS'
            's3\://:S3 binary cache'
            'ssh\://:Nix store over ssh')
        _alternative \
            'file:Use a chroot store:_directories' \
            'url:Other: _describe "URLS and special options" others' \
            'cache:Local binary cache (adds file url prefix):_files -P "file://" -/'
        return
        ;;
    'option FILE')
        _alternative \
            'file:Local file:_nix_complete_dotnix_files' \
            'shortcuts:Shortcuts:_nix_shortcuts' \
            'channel:Channel:(channel:
                       channel:nixos-13.10           channel:nixos-16.09
                       channel:nixos-14.04           channel:nixos-16.09-small
                       channel:nixos-14.04-small     channel:nixos-17.03
                       channel:nixos-14.12           channel:nixos-17.03-small
                       channel:nixos-14.12-small     channel:nixos-17.09
                       channel:nixos-15.09           channel:nixos-17.09-small
                       channel:nixos-15.09-small     channel:nixos-unstable
                       channel:nixos-16.03           channel:nixos-unstable-small
                       channel:nixos-16.03-small     channel:nixpkgs-17.09-darwin
                       channel:nixos-16.03-testing   channel:nixpkgs-unstable)' \
                           'url:URL:(https:// http://)'
        return
        ;;
    'option FILES'|'option PATH')
        _files
        return
        ;;
    'option COMMAND')
        _command_names
        return
        ;;
    'option ARGS')
        _message 'Arguments to command'
        return
        ;;
    'option TYPE')
        _values 'Hash type' md5 sha1 sha256 sha512
        return
        ;;
    'option INCLUDE')
        # --include <PATH> completion
        _nix_complete_includes
        return
        ;;
    'option nixoption')
        # Complete nix options with descriptions
        local -a nix_options
        # Strip the header line, remove leading spaces and replace separating
        # whitespace with ':'
        nix_options=(${${${${(f)"$(nix --help-config)"}:1:-1}/#  /}/  ##/:})
        _describe -t nix_options "Option" nix_options
        return
        ;;
    'option nixoptionvalue')
        # Print the description of the option we're setting
        local OPTION=$words[$(($CURRENT - 1))]
        # Remove lines not starting with " $OPTION " and strip eveything up to the
        # last two consecutive spaces
        local description=${${${(f)"$(nix --help-config)"}:#^(  $OPTION *)}/*  /}
        _message $description
        return
        ;;
    *)
        # Fallback to argument name description
        _message $state_descr
esac