#compdef pgrep pkill

# Notes:
# - We assume that Linux systems use procps-ng — specifically, procps-ng >=3.3.4
#   (which changed the behaviour of -f and added -a)
# - We don't really need to keep pgopts and pkopts separate, but it seems like
#   it should make things a bit easier to follow
# - @todo We could complete log-in classes given to -c
# - @todo We could complete routing tables given to -T

local curcontext="$curcontext" state line ret=1 expl pgopts pkopts no i
typeset -A opt_args
typeset -a arguments sig_arguments aopts

# These arguments (a) are common to all variants (like -x), (b) are the most
# common amongst all variants (like -a), or (c) have a single unambiguous
# meaning amongst all variants (like --help). Many of them are filtered out or
# overridden below
arguments=(
  '-a[include process ancestors in match list]'
  '-c+[match only on specified login class]:login class'
  '(-F --pidfile)'{-F+,--pidfile=}'[match only processes in specified PID file]:PID file:_files'
  '(-f --full)'{-f,--full}'[match against full command line]'
  '(-G --group)'{-G+,--group=}'[match only on specified real group IDs]: :_sequence _groups'
  '(-g --pgroup)'{-g+,--pgroup=}'[match only on specified process group IDs]: :->pgid'
  '(: * -)'{-h,--help}'[display help information]'
  '-I[request confirmation before signalling each process]'
  '-i[ignore case distinctions]'
  '-J+[match only on specified project IDs]: :->projid'
  '-j+[match only on specified jail IDs]:jail:_sequence _jails -0 -o jid'
  '(-L --logpidfile)'{-L,--logpidfile}'[fail if PID file not locked (with -F)]'
  '(-N)-M+[extract name list from specified core]:core file:_files'
  '(-M)-N+[extract name list from specified system]:system file:_files'
  '(-o -n --oldest --newest)'{-n,--newest}'[match newest process]'
  '(-o -n --oldest --newest)'{-o,--oldest}'[match oldest process]'
  '(-P --parent)'{-P+,--parent=}'[match only on specified parent process IDs]: :->ppid'
  '(-l)-q[suppress normal output]'
  '-S[search also in system processes (kernel threads)]'
  '(-s --session)'{-s+,--session=}'[match only on specified process session IDs]: :->sid'
  '(ss)--signal=[specify signal to send to process]:signal:(${signals\:#(EXIT|ZERR|DEBUG)})'
  '-T+[match only on specified routing table]:routing table'
  '(-t --terminal)'{-t+,--terminal=}'[match only on specified controlling terminals]:terminal device:_sequence _ttys -d'
  '(-U --uid)'{-U+,--uid=}'[match only on specified real user IDs]: :_sequence _users'
  '(-u --euid)'{-u+,--euid=}'[match only on specified effective user IDs]: :_sequence _users'
  '(-v --inverse)'{-v,--inverse}'[negate matching]'
  '(-x --exact)'{-x,--exact}'[match process name or command line (with -f) exactly]'
  '--ns=[match only on same namespaces as specified PID]: :_pids'
  '--nslist=[match only on specified namespaces (with --ns)]:namespace:(ipc mnt net pid user uts)'
  '(: * -)'{-V,--version}'[display version information]'
  '-z+[match only on specified zone IDs]:zone:_sequence _zones'
)
[[ $service == pgrep ]] && arguments+=(
  '(-d --delimiter)'{-d+,--delimiter=}'[specify output delimiter]:delimiter:compadd ${(s<>)IFS}'
  '(-q)-l[display process name (and arguments with -f)]'
  '(-w --lightweight)'{-w,--lightweight}'[show all thread IDs instead of PID]'
)
[[ $service == pkill ]] && arguments+=(
  '(-e --echo)'{-e,--echo}'[display signalled process]'
  '-l[display kill command]'
)

case $OSTYPE in
  linux*)
    # Note: We deliberately exclude -v but not --inverse from pkill
    pgopts=acdFfGghLlnoPstUuVvwx-
    pkopts=ceFfGghLnoPstUuVx-
    arguments=(
      ${arguments:#((#s)|*\))(\*|)-[acl]*}
      '(-c --count)'{-c,--count}'[display count of matching processes]'
    )
    [[ $service == pgrep ]] && arguments+=(
      '(-a -l --list-full --list-name)'{-a,--list-full}'[display full command line]'
      '(-a -l --list-full --list-name)'{-l,--list-name}'[display process name]'
    )
    ;;
  dragonfly*|freebsd*)
    pgopts=acdFfGgijLlMNnoPqSstUuvx
    pkopts=acFfGgIijLlMNnoPstUuvx
    ;;
  openbsd*)
    pgopts=dfGglnoPqsTtUuvx
    pkopts=fGgIlnoPqsTtUuvx
    ;;
  darwin*)
    pgopts=adFfGgiLlnoPqtUuvx
    pkopts=aFfGgIiLlnoPtUuvx
    ;;
  solaris*)
    pgopts=cdfGgJlnoPsTtUuvxz
    pkopts=cfGgJnoPsTtUuvxz
    arguments=(
      ${arguments:#((#s)|*\))(\*|)-[cT]*}
      '-c+[match only on specified contract IDs]: :->contract'
      '-T+[match only on specified task IDs]: :->task'
    )
    ;;
  *)
    pgopts=dfGgilnPstUuvx
    pkopts=fGgilnPstUuvx
    ;;
esac

if [[ $service == pgrep ]]; then
  arguments=( ${(M)arguments:#((#s)|*\))(\*|)-[$pgopts]*} )
else
  arguments=( ${(M)arguments:#((#s)|*\))(\*|)-[$pkopts]*} )

  # Signals on non-Linux systems can only be completed as the first argument
  (( CURRENT != 2 )) && [[ $OSTYPE != linux* ]] && no='!'

  # This is used for exclusion with --signal
  sig_arguments=( + '(ss)' )

  # Make sure to exclude zsh-specific signals from our signal list. This
  # presumes that EXIT is always listed first and ZERR/DEBUG last — see
  # signames.c
  for (( i = 1; i <= $#signals; i++ )); do
    [[ $signals[i] == (EXIT|ZERR|DEBUG) ]] && continue
    sig_arguments+=( $no'(--signal)-'$signals[i] )
    # Sometimes zsh will only have the number for a signal
    [[ $signals[i] == <-> ]] || sig_arguments+=(
      '!(--signal)-SIG'$signals[i]
      '!(--signal)-'$(( i - 1 ))
    )
  done
fi

arguments+=( $sig_arguments + o '*: :->pname' )

[[ $OSTYPE == linux* ]] || aopts+=( -A '*-' )
_arguments -C -s -S $aopts : $arguments && ret=0

case $state in
  (sid)
    if [[ $OSTYPE == openbsd* ]]; then
      break
    fi

    compset -P '*,'

    local -a used sid
    used=(${(s:,:)IPREFIX})
    if [[ $OSTYPE == freebsd* ]]; then
      sid=(${(uon)$(ps -ax -o sid=)})
    else
      sid=(${(uon)$(ps -A -o sid=)})
    fi

    _wanted sid expl 'session ID' compadd -S ',' -q -F used $sid
    ;;

  (ppid)
    compset -P '*,'

    local -a used ppid
    used=(${(s:,:)IPREFIX})
    if [[ $OSTYPE == (freebsd|openbsd|darwin)* ]]; then
      ppid=(${(uon)$(ps -ax -o ppid=)})
    else
      ppid=(${(uon)$(ps -A -o ppid=)})
    fi

    _wanted ppid expl 'parent process ID' compadd -S ',' -q -F used $ppid
    ;;

  (pgid)
    compset -P '*,'

    local -a used pgid
    used=(${(s:,:)IPREFIX})
    if [[ $OSTYPE == (freebsd|openbsd|darwin)* ]]; then
      pgid=(${(uon)$(ps -ax -o pgid=)})
    else
      pgid=(${(uon)$(ps -A -o pgid=)})
    fi

    _wanted pgid expl 'process group ID' compadd -S ',' -q -F used $pgid
    ;;

  (projid)
    compset -P '*,'

    local -a used projid
    used=(${(s:,:)IPREFIX})
    projid=(${(uon)$(ps -A -o project=)})

    _wanted projid expl 'project ID' compadd -S ',' -q -F used $projid
    ;;

  (contract)
    compset -P '*,'

    local -a used ctid
    used=(${(s:,:)IPREFIX})
    ctid=(${(uon)$(ps -A -o ctid=)})

    _wanted ctid expl 'contract ID' compadd -S ',' -q -F used $ctid
    ;;

  (task)
    compset -P '*,'

    local -a used taskid
    used=(${(s:,:)IPREFIX})
    taskid=(${(uon)$(ps -A -o project=)})

    _wanted taskid expl 'task ID' compadd -S ',' -q -F used $taskid
    ;;

  (pname)
    local ispat="pattern matching "
    if (( ${+opt_args[-x]} )); then
      ispat=""
    fi

    local command
    if (( ${+opt_args[-f]} )); then
      if [[ "$OSTYPE" == freebsd* ]] && (( ${+opt_args[-S]} )); then
        command="$(ps -axH -o command=)"
      elif [[ "$OSTYPE" == (freebsd|openbsd|darwin)* ]]; then
        command="$(ps -ax -o command=)"
      elif [[ "$OSTYPE" == solaris* ]]; then
        command="$(ps -A -o args=)"
      else
        command="$(ps -A o cmd=)"
      fi
      _wanted pname expl $ispat'process command line' compadd ${(u)${(f)${command}}}
    else
      if [[ "$OSTYPE" == freebsd* ]] && (( ${+opt_args[-S]} )); then
        command="$(ps -axcH -o command=)"
      elif [[ "$OSTYPE" == (freebsd|openbsd|darwin)* ]]; then
        command="$(ps -axc -o command=)"
      elif [[ "$OSTYPE" == solaris* ]]; then
        command="$(ps -A -o comm=)"
      else
        command="$(ps -A co cmd=)"
      fi
      _wanted pname expl $ispat'process name' compadd ${(u)${(f)${command}}}
    fi
    ;;

esac && ret=0

return ret