#compdef dnf dnf-2 dnf-3

_dnf_helper() {
  compadd $($python_exec $helper "$@" -d 0 -q -C 2>/dev/null)
}

_dnf_query_db() {
  sqlite3 -batch -init /dev/null "$cache_file" "$1"
}

_dnf_disabled_repos() {
  _dnf_helper repolist disabled ""
}

_dnf_enabled_repos() {
  _dnf_helper repolist enabled ""
}

_dnf_available_packages() {
  if [ -r $cache_file ]; then
    compadd $(_dnf_query_db "select pkg from available WHERE pkg LIKE \"$1%\"")
  else
    _dnf_helper install "$1"
  fi
}

_dnf_installed_packages() {
  if [ -r $cache_file ]; then
    compadd $(_dnf_query_db "select pkg from installed WHERE pkg LIKE \"$1%\"")
  else
    _dnf_helper remove "$1"
  fi
}

_dnf_local_packages() {
  _files -/ -g '(#i)*.rpm(-.)'
}

_dnf() {
  if [[ "$(readlink /usr/bin/dnf)" == "dnf-2" ]]; then
    local python_exec="/usr/bin/python2"
  else
    local python_exec="/usr/bin/python3"
  fi
  local helper=$(${python_exec} -c "import dnf.cli; print('{}/completion_helper.py'.format(dnf.cli.__path__[0]))")
  local cache_file="/var/cache/dnf/packages.db"

  _arguments -s \
    '(- *)'{-h,--help}'[show the help message]' \
    '--version[show dnf version]' \
    '(-v --verbose)'{-v,--verbose}'[set verbose, show debug messages]' \
    '(-q --quiet)'{-q,--quiet}'[show just the relevant content]' \
    '--allowerasing[allow erasing of installed  packages]' \
    '(-y --assumeyes)'{-y,--assumeyes}'[answer yes for all questions]' \
    '(-C --cacheonly)'{-C,--cacheonly}'[run entirely from cache]' \
    '(-c --config)'{-c,--config=}'[config file location]:config file:_files' \
    '(-R --randomwait)'{-R,--randomwait=}'[maximum command wait time (in minutes)]:max wait time' \
    '--releasever=[configure DNF for another release]:release' \
    '--refresh[set metadata as expired before running the command]' \
    '--nogpgcheck[skip checking GPG signatures on package]' \
    '--installroot=[set install root]:install root:_files -/' \
    '*--enablerepo=[enable one or more repositories]:repos to enable:_dnf_disabled_repos' \
    '*--disablerepo=[disable one or more repositories]:disable repos:_dnf_enabled_repos' \
    '*::dnf command:_dnf_command'
}

_dnf_command() {
  local -a _dnf_cmds
  _dnf_cmds=(
    "autoremove:automatically remove no longer required packages"
    "check-update:check for available package upgrades"
    "clean:remove cached data"
    "distro-sync:synchronize installed packages to the latest available versions"
    "downgrade:downgrade a package"
    "erase:deprecated alias for remove"
    "group:display, or use, the groups information"
    "help:display a helpful usage message"
    "history:display, or use, the transaction history"
    "info:display details about a package or group of packages"
    "install:install a package or packages on your system"
    "list:list a package or groups of packages"
    "makecache:generate the metadata cache"
    "mark:mark or unmark installed packages as installed by user"
    "provides:find what package provides the given value"
    "reinstall:reinstall a package"
    "remove:remove a package or packages from your system"
    "repolist:display the configured software repositories"
    "repository-packages:run commands on top of all packages in given repository"
    "search:search package details for the given string"
    "update:deprecated alias for upgrade"
    "updateinfo:display advisories about packages"
    "upgrade:upgrade a package or packages on your system"
    "upgrade-to:upgrade a package on your system to the specified version"
  )

  if (( CURRENT == 1 )); then
    _describe -t commands 'dnf command' _dnf_cmds || compadd "$@"
  else
    local command="${${_dnf_cmds[(r)$words[1]:*]%%:*}}"
    # Deal with any aliases
    case $command in
      erase) command="remove";;
      whatprovides) command="provides";;
      update) command="upgrade";;
    esac

    _is_path() {
      [[ "$1" == *\/* ]] || [[ "$1" == \~* ]]
    }

    local cur=$words[CURRENT]
    local prev=""
    [[ $CURRENT > 2 ]] && prev=$words[$((CURRENT - 1))]

    case $command in
      install|upgrade|reinstall|info|check-update|distro-sync)
        if ! _is_path "$cur"; then
          _dnf_available_packages "$cur"
        else
          _dnf_local_packages
        fi
        ;;
      remove|downgrade)
        if ! _is_path "$cur"; then
          _dnf_installed_packages "$cur"
        elif [[ "$command" == downgrade ]]; then
          _dnf_local_packages
        fi
        ;;
      list|clean)
        _dnf_helper $command "$prev" "$cur"
        ;;
      group)
        local -a _dnf_group_cmds
        _dnf_group_cmds=(
          "summary:display groups overview"
          "info:display package lists of a group"
          "install:install packages from a group"
          "list:list all matching groups"
          "remove:mark the group removed"
          "upgrade:upgrades the group and its packages"
          "mark:mark a group for installation or removal"
        )
        if (( CURRENT == 2 )); then
          _describe -t commands 'dnf group command' _dnf_group_cmds
        fi
        ;;
      help)
        if (( CURRENT == 2 )); then
          _dnf_helper '_cmds' ''
        fi
        ;;
      history)
        local -a _dnf_history_cmds
        _dnf_history_cmds=(
          "list:list transactions"
          "info:describe the given transactions"
          "redo:repeat the specified transaction"
          "rollback:undo all since the given transaction"
          "undo:undo transactions"
          "userinstalled:list names of all packages installed by a user"
        )
        if (( CURRENT == 2 )); then
          _describe -t commands 'dnf history command' _dnf_history_cmds
        else
          _dnf_helper $command "$prev" "$cur"
        fi
        ;;
      makecache)
        if (( CURRENT == 2 )); then
          _values 'make cache' 'timer'
        fi
        ;;
      mark)
        if (( CURRENT == 2 )); then
          _values 'mark' 'install' 'remove'
        else
          _dnf_installed_packages "$cur"
        fi
        ;;
      provides)
        _files
        ;;
      repolist)
        if (( CURRENT == 2 )); then
          _values 'repolist' 'enabled' 'disabled' 'all'
        fi
        ;;
      search)
        if (( CURRENT == 2 )); then
          _values 'search' 'all'
        fi
        ;;
    esac
  fi
}

_dnf "$@"