#autoload

if [[ $_NIX_SHELL_COMPLETION_LOADED ]]; then
    # No point in re-defining these functions each time we do a completion
    return 0
fi

# Simple completion function to select a system
# List gathered from: https://github.com/NixOS/nixpkgs/blob/master/lib/platforms.nix
_nix_systems () {
    _values 'Systems' \
        i686-linux x86_64-linux \
        armv5tel-linux armv6l-linux armv7l-linux mips64el-linux \
        x86_64-darwin \
        i686-freebsd x86_64-freebsd \
        i686-openbsd x86_64-openbsd \
        i686-netbsd x86_64-netbsd \
        i686-cygwin x86_64-cygwin
}

# Completion function to select an angle-bracket expression from the nix path
# Assumptions: No '=' in the actual path components in NIX_PATH
# TODO: Complete files in /some/path for expressions like <nixpkgs/pkgs/...>
# IMPROVEMENT: Remove '<nixos-config>' since that seems rather useless(?)
_nix_shortcuts () {
    local nix_path=(${(s.:.)NIX_PATH})
    local named=(${(M)nix_path:#*=*})
    local dirs=(${nix_path:#*=*})
    local valid_dir_globs=($^dirs"/*/default.nix(N:h:t)")
    local valid_dirs=(${~valid_dir_globs})
    local names=(${named%%=*})
    _values shortcuts "<"${^valid_dirs}">" "<"${^names}">"
}

_nix_path() {
  _alternative \
    'nixpkgs:Nixpkgs:_nix_shortcuts' \
    'path:File Path:_nix_complete_dotnix_files'
}

_nix_complete_dotnix_files () {
    _path_files '-g' '*.nix(N) *(N-/)'
}

_nix_generations () {
    # List of package names with version numbers stripped
    local generations=($(nix-env --list-generations | awk '{print $1":"$2}'))
    _describe "Nix Generations" generations
}

_nix_installed_packages () {
    # List of package names with version numbers stripped
    local -a packages
    packages=($(nix-env -q | perl -pe 's/-[0-9.-].*$//'))
    _values "Nix Packages" $packages
}


# Generate nix code creating the default expression used by 'nix-env -iA'
_nix_gen_defexpr () {
    setopt local_options null_glob
    local -a result
    local -a queue=($1)
    while [[ ${#queue} > 0 ]]; do
        local current=$queue[1]
        shift queue
        if [[ -e $current/default.nix ]]; then
            result+=($current)
        else
            queue+=($current/*)
        fi
    done

    local nix_expr="{\n"
    for p in $result; do
        nix_expr+="$(basename $p) = import $p {};"
        nix_expr+="\n"
    done
    nix_expr+="}"

    echo $nix_expr
}


# Prints the attribute names (whitespace separated) of the given path
# Arguments: attribute_path [top_level_expr_path]
#   top_level_expr is interpreted as 'nix-env --file'
#   if top_level_expr is omitted it defaults to what 'nix-env -iA' uses
# Eg. _nix_attr_names nixos.python2Packages mypkgs.nix
_nix_attr_names () {
    setopt local_options pipefail

    local attr_path=top${1:+.$1}
    local defexpr_path=$2

    local defexpr
    if [[ -n $defexpr_path ]]; then
        if [[ -e $defexpr_path ]]; then
            # If the path exist use the absolute path to make sure import will
            # accept it.
            # (Otherwise the path is likely a <nixpkgs> notation)
            defexpr_path=${defexpr_path:a}
        fi
        defexpr="import $defexpr_path"
    elif [[ $command_name == nix-env ]]; then
        defexpr=$(_nix_gen_defexpr ~/.nix-defexpr)
    fi

    nix-instantiate --eval - \
          <<NIX_FILE | tr '[]"' ' '
                let
                  top_gen = $defexpr;
                  # --file arguments can be a lambda producing a record too
                  top = if builtins.typeOf top_gen == "lambda" then top_gen {} else top_gen ;
                in
                  builtins.attrNames $attr_path
NIX_FILE

    return $?
}


# Complete attribute paths
# NB: this function expects an optional argument so the user must make sure
#     zsh completion functions doesn't supply arguments. eg. when used in a
#     action spec to _arguments the function should be prefixed with a space:
#     `: _nix_attr_paths`
#     Without the space prefix _arguments will inject lots of arguments. These
#     are available(?) in the `expl` array anyway
_nix_attr_paths () {
    local attr_path=""
    if [[ $words[CURRENT] == *.* ]]; then
        attr_path=${words[CURRENT]%.*}
    fi

    local defexpr_path=$1
    if [[ -z $defexpr_path ]]; then
        if [[ $command_name == nix-env ]]; then
            defexpr_path=$opt_args[-f]
            if [[ -z $defexpr_path ]]; then
                defexpr_path=$opt_args[--file]
            fi
        else
            if [[ $line ]]; then
                defexpr_path=$line[1]
            elif [[ -e shell.nix && $command_name == nix-shell ]]; then
                defexpr_path=shell.nix
            elif [[ -e default.nix ]]; then
                defexpr_path=default.nix
            fi
        fi
    fi

    # Remove one level of shell quoting to make sure we see the same value as
    # the nix-* program will see.
    # ($opt_args and $line contain the verbatim string:
    #  eg. given `nix-shell '<nixpkgs>' -A ` $line[1] will be `'<nixpkgs>'` while
    #  nix-shell will see `<nixpkgs>`)
    defexpr_path=${(Q)defexpr_path}

    local -a result # must be on separate line to capture error code..
    result=( $(_nix_attr_names "$attr_path" $defexpr_path) ) 2>/dev/null

    if [[ $? > 0 ]]; then
        # Probably because of invalid attribute path or error in top-level expr
        return 1
    fi

    local prefix
    if [[ -n $attr_path ]]; then
        prefix=${attr_path}.
    fi

    _values -s . "Packages" ${prefix}${^result}
}

_nix_profiles () {
    local -a profiles
    profiles=($(find /nix/var/nix/profiles))
    _values "Nix Profiles" $profiles
}

# Used in: nix-build, nix-env, nix-instantiate, nix-shell, nixops
_nix_boilerplate_opts=(
  '(- *)--help[Print help message and exit]' \
  '(- *)--version[Print version number and exit]'
)

# Used in: nix-collect-garbage, nix-env, nix-store, nixops
_nix_dry_run='--dry-run[Show what would be done without doing it]'

# Used in: nix-collect-garbage, nix-store
_nix_gc_common=(
  '(- --print* --delete)--print-roots[Print roots used by garbage collector]' \
  '(- --print* --delete)--print-live[Print store paths reachable from roots]' \
  '(- --print* --delete)--print-dead[Print store paths not reachable from roots]' \
  '(- --print* --delete)--delete[Garbage collect all dead paths from the store]' \
)

# Used in: nixos-install, nix_common_opts
_nix_search_path_args=(
  '*-I+[Add path to Nix expression search path]:Include path:_nix_complete_dotnix_files'\
)

# Either true or false: useful for completing many Nix options
_nix_options_bool () {
    _values true false
}
# List gathered from: https://nixos.org/nix/manual/#sec-conf-file
# TODO: Complete the value as well, not just the key
_nix_options () {
    _values \
        'gc-keep-outputs' \
        'gc-keep-derivations' \
        'env-keep-derivations' \
        'build-max-jobs' \
        'build-cores' \
        'build-max-silent-time' \
        'build-timeout' \
        'build-max-log-size' \
        'build-users-group' \
        'build-use-chroot' \
        'build-chroot-dirs' \
        'build-extra-chroot-dirs' \
        'build-use-substitutes' \
        'build-fallback' \
        'build-cache-failure' \
        'build-keep-log' \
        'build-compress-log' \
        'use-binary-caches' \
        'binary-caches' \
        'binary-caches-files' \
        'trusted-binary-caches' \
        'extra-binary-caches' \
        'signed-binary-caches' \
        'binary-cache-public-keys' \
        'binary-caches-parallel-connections' \
        'verify-https-binary-caches' \
        'force-manifest' \
        'system' \
        'fsync-metadata' \
        'auto-optimise-store' \
        'connect-timeout' \
        'log-servers' \
        'trusted-users' \
        'allowed-users' \
        'restrict-eval' \
        'pre-build-hook'
}

_nix_repair='--repair[Fix corrupted or missing store paths by redownloading or rebuilding]';

# Misc Nix options accepted by nixos-rebuild
_nix_common_nixos_rebuild=(
  '(--verbose -v)*'{--verbose,-v}'[Increase verbosity of diagnostic messages]'\
  '(--no-build-output -Q)'{--no-build-output,-Q}'[Silence output to stdout and stderr]'\
  '(--max-jobs -j)'{--max-jobs,-j}'[Set the maximum number of build jobs that Nix will perform in parallel]'\
  '--cores[Set the parallelism of the individual builders (e.g. -j argument to make)]'\
  '(--keep-going -k)'{--keep-going,-k}'[Keep going in case of failed builds, to the greatest extent possible]'\
  '(--keep-failed -K)'{--keep-failed,-K}'[Do not delete the build directory if build fails]'\
  '--fallback[If binary download fails, fall back on building from source]'\
  '--show-trace[Print stack trace of evaluation errors]'\
  '--option[Set Nix configuration option]:Options:_nix_options:Value:( )'\
  $_nix_repair
  )

# Used in: nix-build, nix-env, nix-instantiate, nix-shell
_nix_common_opts=(
  $_nix_common_nixos_rebuild \
  $_nix_search_path_args \
  '(--attr -A)'{--attr,-A}'[Select an attribute from the top-level Nix expression being evaluated]:Packages: _nix_attr_paths'\
  '(--expr -E)'{--expr,-E}'[Interpret command line args as Nix expressions]:*:Files:_files'\
  '*--arg[Argument to pass to the Nix function]:Name:( ):Value:( )'\
  '--argstr[Like --arg, but the value is a string]: :'\
  '--max-silent-time[Builder times out after not producing stdout/stderr for x seconds]:Seconds:( )'\
  '--timeout[Timeout builder after given number of seconds]:Seconds:( )'\
  '--readonly-mode[Do not open Nix database]'\
  '--log-type[Configure how output is formatted]:Output format:((pretty\:"Default" escapes\:"Indicate nesting with escape codes" flat\:"Remove all nesting"))'\
)

_NIX_SHELL_COMPLETION_LOADED=1