#compdef yay

typeset -A opt_args
setopt extendedglob

# options for passing to _arguments: main pacman commands
_pacman_opts_commands=(
  {-D,--database}'[Modify database]'
  {-F,--files}'[Query the files database]'
  {-G,--getpkgbuild}'[Get PKGBUILD from ABS or AUR]'
  {-Q,--query}'[Query the package database]'
  {-R,--remove}'[Remove a package from the system]'
  {-S,--sync}'[Synchronize packages]'
  {-T,--deptest}'[Check if dependencies are installed]'
  {-U,--upgrade}'[Upgrade a package]'
  {-V,--version}'[Display version and exit]'
  '(-h --help)'{-h,--help}'[Display usage]'
)

# options for passing to _arguments: options common to all commands
_pacman_opts_common=(
  '--arch[Set an alternate architecture]'
  {-b,--dbpath}'[Alternate database location]:database_location:_files -/'
  '--color[colorize the output]:color options:(always never auto)'
  {-h,--help}'[Display syntax for the given operation]'
  {-r,--root}'[Set alternate installation root]:installation root:_files -/'
  {-v,--verbose}'[Be more verbose]'
  '--cachedir[Alternate package cache location]:cache_location:_files -/'
  '--config[An alternate configuration file]:config file:_files'
  '--confirm[Always ask for confirmation]'
  '--debug[Display debug messages]'
  '--gpgdir[Set an alternate directory for GnuPG (instead of /etc/pacman.d/gnupg)]: :_files -/'
  '--hookdir[Set an alternate hook location]: :_files -/'
  '--logfile[An alternate log file]:config file:_files'
  '--noconfirm[Do not ask for confirmation]'
  '--noprogressbar[Do not show a progress bar when downloading files]'
  '--noscriptlet[Do not execute the install scriptlet if one exists]'
  '--bottomup[Show AUR packages first]'
  '--topdown[Show repository packages first]'
)

# options for passing to _arguments: options for --upgrade commands
_pacman_opts_pkgfile=(
  '*-d[Skip dependency checks]'
  '*--nodeps[Skip dependency checks]'
  '*--assume-installed[Add virtual package to satisfy dependencies]'
  '--dbonly[Only remove database entry, do not remove files]'
  '--force[Overwrite conflicting files]'
  '--needed[Do not reinstall up to date packages]'
  '--asdeps[mark packages as non-explicitly installed]'
  '--asexplicit[mark packages as explicitly installed]'
  {-p,--print}'[Only print the targets instead of performing the operation]'
  '*--ignore[Ignore a package upgrade]:package: _pacman_completions_all_packages'
  '*--ignoregroup[Ignore a group upgrade]:package group:_pacman_completions_all_groups'
  '--print-format[Specify how the targets should be printed]'
  '*:package file:_files -g "*.pkg.tar*~*.sig(.,@)"'
)

# options for passing to _arguments: subactions for --query command
_pacman_opts_query_actions=(
  '(-Q --query)'{-Q,--query}
  {-g,--groups}'[View all members of a package group]:*:package groups:->query_group'
  {-o,--owns}'[Query the package that owns a file]:file:_files'
  {-p,--file}'[Package file to query]:*:package file:->query_file'
  {-s,--search}'[Search package names and descriptions]:*:search text:->query_search'
)

# options for passing to _arguments: options for --query and subcommands
_pacman_opts_query_modifiers=(
  {-c,--changelog}'[List package changelog]'
  {-d,--deps}'[List packages installed as dependencies]'
  {-e,--explicit}'[List packages explicitly installed]'
  {\*-i,\*--info}'[View package information]'
  {\*-k,\*--check}'[Check package files]'
  {-l,--list}'[List package contents]'
  {-m,--foreign}'[List installed packages not found in sync db(s)]'
  {-n,--native}'[List installed packages found in sync db(s)]'
  {-q,--quiet}'[Show less information for query and search]'
  {-t,--unrequired}'[List packages not required by any package]'
  {-u,--upgrades}'[List packages that can be upgraded]'
)

# options for passing to _arguments: options for --remove command
_pacman_opts_remove=(
  {-c,--cascade}'[Remove all dependent packages]'
  {-d,--nodeps}'[Skip dependency checks]'
  '*--assume-installed[Add virtual package to satisfy dependencies]'
  {-n,--nosave}'[Remove protected configuration files]'
  {-p,--print}'[Only print the targets instead of performing the operation]'
  {\*-s,\*--recursive}'[Remove dependencies not required by other packages]'
  {-u,--unneeded}'[Remove unneeded packages]'
  '--dbonly[Only remove database entry, do not remove files]'
  '--print-format[Specify how the targets should be printed]'
  '*:installed package:_pacman_completions_installed_packages'
)

_pacman_opts_database=(
  '--asdeps[mark packages as non-explicitly installed]'
  '--asexplicit[mark packages as explicitly installed]'
  '*:installed package:_pacman_completions_installed_packages'
)

_pacman_opts_files=(
  {-l,--list}'[List the files owned by the queried package]:package:_pacman_completions_all_packages'
  {-o,--owns}'[Query the package that owns]:files:_files'
  {-s,--search}'[Search package file names for matching strings]:files:_files'
  {-x,--regex}'[Enable searching using regluar expressions]:regex:'
  {-y,--refresh}'[Download fresh files databases from the server]'
  '--machinereadable[Produce machine-readable output]'
  {-q,--quiet}'[Show less information for query and search]'
)

# options for passing to _arguments: options for --sync command
_pacman_opts_sync_actions=(
  '(-S --sync)'{-S,--sync}
  {\*-c,\*--clean}'[Remove old packages from cache]:\*:clean:->sync_clean'
  {-g,--groups}'[View all members of a package group]:*:package groups:->sync_group'
  {-s,--search}'[Search package names and descriptions]:*:search text:->sync_search'
  '--dbonly[Only remove database entry, do not remove files]'
  '--needed[Do not reinstall up to date packages]'
  '--recursive[Reinstall all dependencies of target packages]'
)

# options for passing to _arguments: options for --sync command
_pacman_opts_sync_modifiers=(
  {\*-d,\*--nodeps}'[Skip dependency checks]'
  '*--assume-installed[Add virtual package to satisfy dependencies]'
  {\*-i,\*--info}'[View package information]'
  {-l,--list}'[List all packages in a repository]'
  {-p,--print}'[Print download URIs for each package to be installed]'
  {-q,--quiet}'[Show less information for query and search]'
  {\*-u,\*--sysupgrade}'[Upgrade all out-of-date packages]'
  {-w,--downloadonly}'[Download packages only]'
  {\*-y,\*--refresh}'[Download fresh package databases]'
  '*--ignore[Ignore a package upgrade]:package: _pacman_completions_all_packages'
  '*--ignoregroup[Ignore a group upgrade]:package group:_pacman_completions_all_groups'
  '--asdeps[Install packages as non-explicitly installed]'
  '--asexplicit[Install packages as explicitly installed]'
  '--force[Overwrite conflicting files]'
  '--print-format[Specify how the targets should be printed]'
)

# handles --help subcommand
_pacman_action_help() {
  _arguments -s : \
    "$_pacman_opts_commands[@]"
}

# handles cases where no subcommand has yet been given
_pacman_action_none() {
  _arguments -s : \
    "$_pacman_opts_commands[@]"
}

# handles --query subcommand
_pacman_action_query() {
  local context state line
  typeset -A opt_args

  case $state in
    query_file)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:package file:_files -g "*.pkg.tar*~*.sig(.,@)"'
      ;;
    query_group)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:groups:_pacman_completions_installed_groups'
      ;;
    query_owner)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:file:_files'
      ;;
    query_search)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:search text: '
      ;;
    *)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_actions[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:package:_pacman_completions_installed_packages'
      ;;
  esac
}

# handles --remove subcommand
_pacman_action_remove() {
  _arguments -s : \
    '(--remove -R)'{-R,--remove} \
    "$_pacman_opts_common[@]" \
    "$_pacman_opts_remove[@]"
}

# handles --database subcommand
_pacman_action_database() {
  _arguments -s : \
    '(--database -D)'{-D,--database} \
    "$_pacman_opts_common[@]" \
    "$_pacman_opts_database[@]"
}

# handles --files subcommand
_pacman_action_files() {
  _arguments -s : \
    '(--files -F)'{-F,--files} \
    "$_pacman_opts_common[@]" \
    "$_pacman_opts_files[@]"
}

_pacman_action_deptest () {
  _arguments -s : \
    '(--deptest)-T' \
    "$_pacman_opts_common[@]" \
    ":packages:_pacman_all_packages"
}


# handles --sync subcommand
_pacman_action_sync() {
  local context state line
  typeset -A opt_args
  if (( $+words[(r)--clean] )); then
    state=sync_clean
  elif (( $+words[(r)--groups] )); then
    state=sync_group
  elif (( $+words[(r)--search] )); then
    state=sync_search
  fi

  case $state in
    sync_clean)
      _arguments -s : \
        {\*-c,\*--clean}'[Remove old packages from cache]' \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_sync_modifiers[@]"
        ;;
    sync_group)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_sync_modifiers[@]" \
        '(-g --group)'{-g,--groups} \
        '*:package group:_pacman_completions_all_groups'
      ;;
    sync_search)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_sync_modifiers[@]" \
        '*:search text: '
      ;;
    *)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_sync_actions[@]" \
        "$_pacman_opts_sync_modifiers[@]" \
        '*:package:_pacman_completions_all_packages'
      ;;
  esac
}

# handles --upgrade subcommand
_pacman_action_upgrade() {
  _arguments -s : \
    '(-U --upgrade)'{-U,--upgrade} \
    "$_pacman_opts_common[@]" \
    "$_pacman_opts_pkgfile[@]"
}

# handles --version subcommand
_pacman_action_version() {
  # no further arguments
  return 0
}

# provides completions for package groups
_pacman_completions_all_groups() {
  local -a cmd groups
  _pacman_get_command
  groups=( $(_call_program groups $cmd[@] -Sg) )
  typeset -U groups

  if [[ ${words[CURRENT-1]} == '--ignoregroup' ]]; then
    _sequence compadd -S ',' "$@" -a groups
  else
    compadd "$@" -a groups
  fi
}

# provides completions for packages available from repositories
# these can be specified as either 'package' or 'repository/package'
_pacman_completions_all_packages() {
  local -a seq sep cmd packages repositories packages_long

  if [[ ${words[CURRENT-1]} == '--ignore' ]]; then
    seq='_sequence'
    sep=(-S ',')
  else
    seq=
    sep=()
  fi

  if compset -P1 '*/*'; then
    packages=( $(_call_program packages yay -Pc ${words[CURRENT]%/*}) )
    typeset -U packages
    ${seq} _wanted repo_packages expl "repository/package" compadd ${sep[@]} ${(@)packages}
  else
    packages=( $(_call_program packages yay -Pc ) )
    typeset -U packages
    ${seq} _wanted packages expl "packages" compadd ${sep[@]} - "${(@)packages}"

    repositories=(${(o)${${${(M)${(f)"$(</etc/pacman.conf)"}:#\[*}/\[/}/\]/}:#options})
    typeset -U repositories
    _wanted repo_packages expl "repository/package" compadd -S "/" $repositories
  fi
}

# provides completions for package groups
_pacman_completions_installed_groups() {
  local -a cmd groups
  _pacman_get_command
  groups=(${(o)${(f)"$(_call_program groups $cmd[@] -Qg)"}% *})
  typeset -U groups
  compadd "$@" -a groups
}

# provides completions for installed packages
_pacman_completions_installed_packages() {
  local -a cmd packages packages_long
  packages_long=(/var/lib/pacman/local/*(/))
  packages=( ${${packages_long#/var/lib/pacman/local/}%-*-*} )
  compadd "$@" -a packages
}

_pacman_all_packages() {
  _alternative : \
    'localpkgs:local packages:_pacman_completions_installed_packages' \
    'repopkgs:repository packages:_pacman_completions_all_packages'
}

# provides completions for repository names
_pacman_completions_repositories() {
  local -a cmd repositories
  repositories=(${(o)${${${(M)${(f)"$(</etc/pacman.conf)"}:#\[*}/\[/}/\]/}:#options})
  # Uniq the array
  typeset -U repositories
  compadd "$@" -a repositories
}

# builds command for invoking pacman in a _call_program command - extracts
# relevant options already specified (config file, etc)
# $cmd must be declared by calling function
_pacman_get_command() {
  # this is mostly nicked from _perforce
  cmd=( "pacman" "2>/dev/null")
  integer i
  for (( i = 2; i < CURRENT - 1; i++ )); do
    if [[ ${words[i]} = "--config" || ${words[i]} = "--root" ]]; then
      cmd+=( ${words[i,i+1]} )
    fi
  done
}

# main dispatcher
_pacman_zsh_comp() {
  local -a args cmds;
  local tmp
  args=( ${${${(M)words:#-*}#-}:#-*} )
  for tmp in $words; do
    cmds+=("${${_pacman_opts_commands[(r)*$tmp\[*]%%\[*}#*\)}")
  done
  case $args in #$words[2] in
    h*)
      if (( ${(c)#args} <= 1 && ${(w)#cmds} <= 1 )); then
        _pacman_action_help
      else
        _message "no more arguments"
      fi
      ;;
    *h*)
      _message "no more arguments"
      ;;
    D*)
      _pacman_action_database
      ;;
    F*)
      _pacman_action_files
      ;;
    Q*g*) # ipkg groups
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:groups:_pacman_completions_installed_groups'
      ;;
    Q*o*) # file
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:package file:_files'
      ;;
    Q*p*) # file *.pkg.tar*
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_query_modifiers[@]" \
        '*:package file:_files -g "*.pkg.tar*~*.sig(.,@)"'
      ;;
    T*)
      _pacman_action_deptest
      ;;
    Q*)
      _pacman_action_query
      ;;
    R*)
      _pacman_action_remove
      ;;
    S*c*) # no completion
      _arguments -s : \
        '(-c --clean)'{\*-c,\*--clean}'[Remove all files from the cache]' \
        "$_pacman_opts_common[@]"
      ;;
    S*l*) # repos
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_sync_modifiers[@]" \
        '*:package repo:_pacman_completions_repositories' \
      ;;
    S*g*) # pkg groups
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_sync_modifiers[@]" \
        '*:package group:_pacman_completions_all_groups'
      ;;
    S*s*)
      _arguments -s : \
        "$_pacman_opts_common[@]" \
        "$_pacman_opts_sync_modifiers[@]" \
        '*:search text: '
        ;;
    S*)
      _pacman_action_sync
      ;;
    T*)
       _arguments -s : \
        '-T' \
        "$_pacman_opts_common[@]" \
        ":packages:_pacman_all_packages"
      ;;
    U*)
      _pacman_action_upgrade
      ;;
    V*)
      _pacman_action_version
      ;;
    *)

      case ${(M)words:#--*} in
        *--help*)
          if (( ${(w)#cmds} == 1 )); then
            _pacman_action_help
          else
            return 0;
          fi
          ;;
        *--sync*)
          _pacman_action_sync
          ;;
        *--query*)
          _pacman_action_query
          ;;
        *--remove*)
          _pacman_action_remove
          ;;
        *--deptest*)
          _pacman_action_deptest
          ;;
        *--database*)
          _pacman_action_database
          ;;
        *--files*)
          _pacman_action_files
          ;;
        *--version*)
          _pacman_action_version
          ;;
        *--upgrade*)
          _pacman_action_upgrade
          ;;
        *)
          _pacman_action_none
          ;;
      esac
      ;;
  esac
}
_pacman_comp() {
  case "$service" in
    yay)
      _pacman_zsh_comp "$@"
      ;;
    *)
      _message "Error"
      ;;
  esac
}

_pacman_comp "$@"
