btrctl.bash-completion

# Run a btrctl subcommand that needs root, without ever prompting: btrctl
# would otherwise exec an interactive sudo from inside tab completion. With no
# cached sudo ticket this fails silently and callers get no candidates.
__btrctl_run_root() {
  local btrctl_path

  btrctl_path=$(type -P btrctl || true)
  [ -n "$btrctl_path" ] || return

  if [ "$(id -u)" -eq 0 ]; then
    command "$btrctl_path" "$@" 2>/dev/null || true
  elif [ -n "${BTRCTL_CONF:-}" ]; then
    command sudo -n BTRCTL_CONF="$BTRCTL_CONF" -- "$btrctl_path" "$@" 2>/dev/null || true
  else
    command sudo -n -- "$btrctl_path" "$@" 2>/dev/null || true
  fi
}

__btrctl_trunk_hosts() {
  __btrctl_run_root trunk hosts
}

__btrctl_local_sets() {
  __btrctl_run_root list | awk 'NR > 1 { print $2 }'
}

__btrctl() {
  local cur word i first_arg kind expect="" options="" set_candidates=""
  local host_seen=0 positional_count=0 dry_run_seen=0 all_seen=0
  cur="${COMP_WORDS[COMP_CWORD]}"

  local commands="snapshot list latest prune rm tag trunk config help -h --help"
  local trunk_commands="id backup hosts list prune rm tag"

  COMPREPLY=()

  if [ "$COMP_CWORD" -eq 1 ]; then
    COMPREPLY=($(compgen -W "$commands" -- "$cur"))
    return
  fi

  if [ "${COMP_WORDS[1]}" = "trunk" ] && [ "$COMP_CWORD" -eq 2 ]; then
    COMPREPLY=($(compgen -W "$trunk_commands" -- "$cur"))
    return
  fi

  if [ "${COMP_WORDS[1]}" = "trunk" ] && [ "${COMP_WORDS[2]}" = "list" ] &&
    [ "$COMP_CWORD" -eq 3 ]; then
    [[ "$cur" = -* ]] && return
    COMPREPLY=($(compgen -W "$(__btrctl_trunk_hosts)" -- "$cur"))
    return
  fi

  if [ "${COMP_WORDS[1]}" = "trunk" ] &&
    [[ "${COMP_WORDS[2]}" = rm || "${COMP_WORDS[2]}" = tag ]] &&
    [ "$COMP_CWORD" -eq 3 ]; then
    [[ "$cur" = -* ]] && return
    COMPREPLY=($(compgen -W "$(__btrctl_trunk_hosts | awk '{ print $0 "/" }')" -- "$cur"))
    compopt -o nospace
    return
  fi

  case "${COMP_WORDS[1]}" in
    list|latest)
      first_arg=2
      kind="host-only"
      ;;
    prune)
      first_arg=2
      kind="prune"
      ;;
    snapshot)
      first_arg=2
      kind="snapshot"
      ;;
    rm)
      first_arg=2
      kind="rm"
      ;;
    tag)
      first_arg=2
      kind="tag"
      ;;
    trunk)
      first_arg=3
      case "${COMP_WORDS[2]}" in
        backup) kind="trunk-backup" ;;
        tag) kind="trunk-tag" ;;
        prune) kind="trunk-prune" ;;
        *) return ;;
      esac
      ;;
    *) return ;;
  esac

  for ((i = first_arg; i < COMP_CWORD; i++)); do
    word="${COMP_WORDS[i]}"
    if [ -n "$expect" ]; then
      case "$expect" in
        host) host_seen=1 ;;
      esac
      expect=""
      continue
    fi

    case "$word" in
      -H|--host)
        [[ "$kind" != trunk-tag && "$kind" != trunk-prune ]] || return
        expect="host"
        ;;
      -n|--dry-run)
        [[ "$kind" = prune || "$kind" = trunk-prune ]] || return
        dry_run_seen=1
        ;;
      -a|--all)
        [ "$kind" = "trunk-prune" ] && [ "$positional_count" -eq 0 ] || return
        all_seen=1
        ;;
      -t|--tag)
        [ "$kind" = "snapshot" ] || return
        expect="tag"
        ;;
      -*) return ;;
      *)
        case "$kind" in
          host-only|snapshot|prune) return ;;
          trunk-prune) [ "$all_seen" -eq 0 ] || return ;;
        esac
        positional_count=$((positional_count + 1))
        ;;
    esac
  done

  case "$expect" in
    host)
      COMPREPLY=($(compgen -W "$(__btrctl_trunk_hosts)" -- "$cur"))
      return
      ;;
    tag) return ;;
  esac

  case "$kind" in
    host-only)
      [ "$host_seen" -eq 0 ] && options="--host"
      ;;
    prune)
      [ "$host_seen" -eq 0 ] && options="--host"
      [ "$dry_run_seen" -eq 1 ] || options="${options:+${options} }--dry-run"
      ;;
    trunk-prune)
      [ "$positional_count" -le 1 ] || return
      [ "$dry_run_seen" -eq 1 ] || options="--dry-run"
      if [ "$positional_count" -eq 0 ] && [ "$all_seen" -eq 0 ]; then
        options="${options:+${options} }--all"
        [[ "$cur" = -* ]] || set_candidates=$(__btrctl_trunk_hosts)
      fi
      ;;
    snapshot)
      [ "$host_seen" -eq 0 ] && options="--host"
      options="${options:+${options} }--tag"
      ;;
    rm)
      [ "$positional_count" -le 1 ] || return
      [ "$host_seen" -eq 0 ] && options="--host"
      if [ "$host_seen" -eq 0 ] && [ "$positional_count" -eq 0 ]; then
        set_candidates=$(__btrctl_local_sets)
      fi
      ;;
    tag)
      [ "$host_seen" -eq 0 ] && options="--host"
      if [ "$host_seen" -eq 0 ] && [ "$positional_count" -eq 0 ]; then
        set_candidates=$(__btrctl_local_sets)
      fi
      [ "$positional_count" -ne 1 ] || set_candidates="add rm"
      ;;
    trunk-tag)
      [ "$positional_count" -ne 1 ] || set_candidates="add rm"
      ;;
    trunk-backup)
      [ "$positional_count" -le 1 ] || return
      [ "$host_seen" -eq 0 ] && options="--host"
      if [ "$host_seen" -eq 0 ] && [ "$positional_count" -eq 0 ]; then
        set_candidates=$(__btrctl_local_sets)
      fi
      ;;
  esac

  COMPREPLY=($(compgen -W "${options}${options:+ }${set_candidates}" -- "$cur"))
}

complete -F __btrctl btrctl