#!/usr/bin/env bash

# We're faking a `nix build` command-line to re-use Nix's own completion
# for the few options passed through to Nix.
_nixos-rebuild_pretend-nix() {
  COMP_LINE="nix build ${COMP_LINE}"
  # number of prepended chars
  (( COMP_POINT = COMP_POINT + 10))

  COMP_WORDS=(
    nix build
    "${COMP_WORDS[@]}"
  )
  # Add the amount of prepended words
  (( COMP_CWORD = COMP_CWORD + 2))
  _complete_nix "nix"
}

_nixos-rebuild() {
  local curr="$2"
  local prev="$3"
  local subcommandGiven=0
  local word
  local subcommand

  __load_completion nix

  # Arrays are re-ordered by the completion, so it's fine to sort them in logical chunks
  local all_args=(
    --verbose -v

    # nixos-rebuild options
    --fast
    --no-build-nix
    --profile-name -p # name
    --rollback
    --specialisation -c # name
    --use-remote-sudo
    --build-host # host
    --target-host # host
    # Used with list-generations
    --json

    # generation switching options
    --install-bootloader

    # nix-channel options
    --upgrade
    --upgrade-all

    # flakes options
    --commit-lock-file
    --flake # flake-uri
    --override-input # input-name flake-uri
    --recreate-lock-file
    --update-input
    --no-flake
    --no-registries
    --no-update-lock-file
    --no-write-lock-file

    # Nix-copy options
    --use-substitutes --substitute-on-destination -s

    # Nix options
    --option
    --impure
    --builders # builder-spec
    --show-trace
    --keep-failed -K
    --keep-going -k
    --max-jobs -j # number
    --log-format # format
    -I # NIX_PATH
  )

  local all_subcommands=(
    boot
    build
    build-vm
    build-vm-with-bootloader
    dry-activate
    dry-build
    edit
    list-generations
    switch
    test
  )

  # Suggest arguments that can be consumed under some conditions only
  for word in "${COMP_WORDS[@]}"; do
    for subcommand in "${all_subcommands[@]}"; do
      if [[ "$word" == "$subcommand" ]]; then
        subcommandGiven=1
      fi
    done
  done

  # Fake out a way to complete the second arg to some options
  case "${COMP_WORDS[COMP_CWORD-2]}" in
    "--override-input")
      prev="--override-input_2"
      ;;
    "--option")
      prev="--option_2"
      ;;
  esac

  case "$prev" in
    --max-jobs|-j)
      COMPREPLY=( )
      ;;

    --profile-name|-p)
      if [[ "$curr" == "" ]]; then
        COMPREPLY=( /nix/var/nix/profiles/* )
      else
        COMPREPLY=( "$curr"* )
      fi
      ;;

    --build-host|--target-host|-t|-h)
      _known_hosts_real "$curr"
    ;;

    --specialisation|-c)
      COMPREPLY=()
      ;;

    -I)
      _nixos-rebuild_pretend-nix
      ;;
    --builders)
      _nixos-rebuild_pretend-nix
      ;;
    --flake)
      _nixos-rebuild_pretend-nix
      ;;
    --override-input)
      _nixos-rebuild_pretend-nix
      ;;
    --override-input_2)
      _nixos-rebuild_pretend-nix
      ;;
    --log-format)
      _nixos-rebuild_pretend-nix
      ;;
    --option)
      _nixos-rebuild_pretend-nix
      ;;
    --option_2)
      _nixos-rebuild_pretend-nix
      ;;

    *)
      if [[ "$curr" == -* ]] || (( subcommandGiven )); then
        COMPREPLY=( $(compgen -W "${all_args[*]}" -- "$2") )
      else
        COMPREPLY=( $(compgen -W "${all_subcommands[*]}" -- "$2") )
      fi
    ;;
  esac
}

complete -F _nixos-rebuild nixos-rebuild