_nix-common-options

#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_complete_includes () {
    local cur=${words[$CURRENT]}
    local -a nixpath=(${(s.:.)NIX_PATH})
    local -a path_names
    local p
    for p in ${nixpath[*]}; do
        [[ "$p" == *=* ]] && \
            path_names+=(${p%=*}:Path:_nix_complete_dotnix_files)
    done
    if [[ $cur == *=* ]]; then
        path_names+=(${cur%=*}:Path:_nix_complete_dotnix_files)
    fi
    _alternative \
        "nixpath:nixpath:_values -s = 'Nix path name' $path_names" \
        'file:Path:_files -/'
    return
}

_nix_generations () {
    # List of package names with version numbers stripped
    setopt extendedglob
    local -a generations=(${${${(f)"$(nix-env --list-generations)"}##  #}/  /:})
    _describe -V "Nix Generations" generations
}

_nix_installed_packages () {
    # List of package names with version numbers stripped
    # with the alternative to specify store paths
    local extra_help=$1
    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
    local -a packages=(${${(f)"$(nix-env -q)"}%%-[0-9]*})

    _alternative \
        "package:packages:_values 'Installed package $extra_help' $packages" \
        "file:Store path to package:_files ${prefix}"
}


# 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
}


# Complete attribute names using $1 as the toplevel expression NB: If calling
# this function from an action spec you need to prefix it with a space so that
# eg. _arguments won't pass it options which it doesn't expect, ie:
#     `: _nix_attr_paths`
_nix_attr_paths () {
    setopt local_options pipefail

    local defexpr=$1
    if [[ "$2" == false ]]; then
       _message "The URL isn't cached, download to get completion"
       return
    fi

    local attr_path=""
    if [[ $words[CURRENT] == *.* ]]; then
        attr_path=${words[CURRENT]%.*}
    fi
    local prefixed_path=top${attr_path:+.$attr_path}

    # Build up a modified NIX_PATH using -I and --include
    local i override=""
    for ((i=1; i < ${#words[*]}; i++)); do
        case "${words[i]}" in
            -I|--include)
                override+=${override:+:}${words[$((i+1))]}
                ;;
        esac
    done
    override+=${override:+:}${NIX_PATH}

    # Resolve channel: syntax
    while [[ "$override" == *(=|:)channel:* ]]; do
        local channel=${override#*channel:}
        channel="channel:"${channel%%:*}
        local url="https://nixos.org/channels/"${channel:8}"/nixexprs.tar.xz"
        # Replace the channel with its url
        override=${override/"$channel"/"$url"}
    done

    local cache_help=""
    # Resolve any url to a cache, else we might trigger a blocking download
    while [[ "$override" == *https://* ]]; do
        # Find the first url
        local url=${override#*https://}
        # Strip everything starting with the first colon
        url="https://"${url%%:*}
        local cache=$(_nix_resolve_url "$url")
        [[ -z "$cache" ]] \
            && cache_help=" (an URL in the NIX_PATH isn't cached, completion might not work)"
        # Replace the url with the cache
        override=${override/"$url"/"$cache"}
    done

    local -a result=($(
        NIX_PATH=$override nix-instantiate --eval -  2> /dev/null \
                <<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 $prefixed_path
NIX_FILE
    ))

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

    local -a prefix=()
    if [[ -n $attr_path ]]; then
        for i in ${=attr_path//./ }; do
            prefix+=("($i)" .)
        done
    fi

    _wanted package package "Attribute path$cache_help" \
            _sep_parts $prefix result \.
    return $?
}


# Generate the top level expression in all the various ways the different
# commands expects it to be built. Then generate completions by calling
# _nix_attr_paths $defexpr
_nix_complete_attr_paths () {

    local defexpr="" cached=true
    local file=$(_nix_get_file_arg)
    if [[ "$file" ]]; then
        # Extract --arg and --argstr into $args
        local i=1 args="" name="" value=""
        for ((i=1; i < ${#words[*]}; i++)); do
            case "${words[$i]}" in
                --arg)
                    name=${(Q)words[$((i+1))]}
                    value=${(Q)words[$((i+2))]}
                    args+="$name = $value;"
                    i=$((i+2))
                    ;;
                --argstr)
                    name=${(Q)words[$((i+1))]}
                    value=${(Q)words[$((i+2))]}
                    args+="$name = \"$value\";"
                    i=$((i+2))
                    ;;
            esac
        done

        defexpr="import $file {$args}"
    else
        if [[ $service == nix-env ]]; then
            defexpr=$(_nix_gen_defexpr ~/.nix-defexpr)

        elif [[ $service == nix ]]; then
            # Extract the channels from NIX_PATH and -I/--include
            local -a channels=(${(s.:.)NIX_PATH})
            # Add -I/--include afterwards, so they will shadow the NIX_PATH
            channels+=(${(s.:.)opt_args[-I]})
            channels+=(${(s.:.)opt_args[--include]})

            # Add the names in an associative array to avoid duplicates
            local -A names
            local channel name
            for channel in $channels; do
                name=${channel%%=*}
                nix_path=${channel#*=}
                if [[ $name != $channel ]]; then
                    # Only add paths with a name, not sure how they work
                    names[$name]=1
                fi
            done

            defexpr=$'{ '
            for name in ${(@k)names}; do
                # nixos-config isn't useful or possible to complete
                [[ $name == nixos-config ]] && continue
                defexpr+="$name = import <${name}> {}; "
            done
            defexpr+=' }'
        fi
    fi

    _nix_attr_paths $defexpr $cached 2>/dev/null
}

function _nix_resolve_url () {
    local -a sha
    local url=$1
    local version="$($service --version)"
    if [[ "${version##* }" == 1.11.* ]]; then
        # works for nix 1.11
        sha=($(printf "$url" | sha256sum -))
    else
        # works for nix 1.12
        local name="${url##*/}"
        sha=($(printf "$name\0$url" | sha256sum -))
    fi
    local base=$(nix-hash --to-base32 --type sha256 \
                          ${sha[*]:0:1})
    local cache=${XDG_CACHE_HOME:-~/.cache}/nix/tarballs
    local link="$cache"/"$base"-file
    if [[ -e "$link" ]]; then
        echo "$cache/$(basename $(readlink $link))-unpacked"
    fi
}

function _nix_get_file_arg () {

    local file=""
    if [[ "$service" == (nix-env|nix) ]]; then
        local i
        # Extract the last seen -f/--file argument
        for ((i=1; i < ${#words[*]}; i++)); do
            case "${words[i]}" in
                --file|-f)
                    file=${words[$((i+1))]}
                    ;;
            esac
        done
    elif [[ $line ]]; then
        file=$line[1]
    elif [[ -e shell.nix && $service == nix-shell ]]; then
        file=shell.nix
    elif [[ -e default.nix ]]; then
        file=default.nix
    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>`)
    file=${(Q)file}

    if [[ "file" ]]; then
        # Expand channel: syntax
        if [[ "$file" == channel:* ]]; then
            file="https://nixos.org/channels/"${file:8}"/nixexprs.tar.xz"
        fi

        if [[ -e $file ]]; then
            # If the path exist use the absolute path to make sure import will
            # accept it.
            # (Otherwise the path is likely a <nixpkgs> notation)
            file=${file:a}
        elif [[ "$file" == https://* ]]; then
            file=$(_nix_resolve_url $file)
            [[ -z "$file" ]] && cached=false
        fi
    fi
    print -n -- $file
}

function _nix_complete_function_arg () {
    local file=$(_nix_get_file_arg)
    local -a names=($(nix-instantiate --eval - <<NIX_FILE 2> /dev/null \
                 | tr '"[]' ' '
             let
             args = builtins.functionArgs (import $file);
             in
                 builtins.attrNames args
NIX_FILE
           ))
    _values "Argument name" $names
}

_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_includes'\
)

# 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=(
    $_nix_search_path_args \
  '(--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 \
  '(--attr -A)'{--attr,-A}'[Select an attribute from the top-level Nix expression being evaluated]:Packages: _nix_complete_attr_paths'\
  '*--arg[Argument to pass to the Nix function]:Name:_nix_complete_function_arg:Value: '\
  '*--argstr[Like --arg, but the value is a string]:Name:_nix_complete_function_arg: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