#compdef sudo sudoedit

local curcontext="$curcontext" environ e cmd cpp ret=1
local -a context state state_descr line args _comp_priv_prefix
local -A opt_args

zstyle -a ":completion:${curcontext}:" environ environ

for e in "${environ[@]}"
do local -x "$e"
done

args=(
  '(-A --askpass)'{-A,--askpass}'[use a helper program for password prompting]'
  '(-B --bell)'{-B,--bell}'[ring bell when prompting]'
  '(-C --close-from)'{-C+,--close-from=}'[close file descriptors]:lowest fd to close'
  '(-D --chdir)'{-D+,--chdir=}'[change the working directory before running command]:directory:_directories'
  '(-g --group)'{-g+,--group=}'[run command as the specified group name or ID]:group:_groups'
  '(-)'{-h,--help}'[display help message and exit]'
  '(-h --host)'{-h+,--host=}'[run command on host]:host:_hosts'
  '(-k --reset-timestamp -K --remove-timestamp -N --no-update)'{-K,--remove-timestamp}'[remove timestamp file completely]'
  '(-k --reset-timestamp -K --remove-timestamp -N --no-update)'{-k,--reset-timestamp}'[invalidate timestamp file]'
  \*{-l,--list}"[list user's privileges or check a specific command]"
  '(-n --non-interactive)'{-n,--non-interactive}'[non-interactive mode, no prompts are used]'
  '(-k --reset-timestamp -K --remove-timestamp -N --no-update)'{-N,--no-update}"[don't update user's cached credentials]"
  '(-p --prompt)'{-p+,--prompt=}'[use the specified password prompt]:prompt'
  '(-R --chroot)'{-R+,--chroot=}'[change the root directory before running command]:directory:_directories'
  '(-r --role)'{-r+,--role=}'[create SELinux security context with specified role]: :_selinux_roles'
  '(-S --stdin)'{-S,--stdin}'[read password from standard input]'
  '(-t --type)'{-t+,--type=}'[create SELinux security context with specified type]: :_selinux_types'
  '(-T --command-timeout)'{-T+,--command-timeout=}'[terminate command after specified time limit]:timeout'
  '(-U --other-user -v --validate)'{-U+,--other-user=}'[in list mode, display privileges for user]:user:_users'
  '(-u --user)'{-u+,--user=}'[run command (or edit file) as specified user]:user:_users'
  '(-)'{-V,--version}'[display version information and exit]'
  '(-v --validate -U --other-user *)'{-v,--validate}"[update user's timestamp without running a command]"
)

# Does -e appears before the first word that doesn't begin with a hyphen?
# The way (i) works in subscripts, the test will always be true if all the
# words begin with a hyphen.
# 
# TODO: use _arguments' $opt_args to detect the cases '-u jrandom -e' and '-Ae'
if [[ $service = sudoedit ]] || (( $words[(i)-e] < $words[(i)^(*sudo|-[^-]*)] ))  ; then
  args+=( '!(-V --version -h --help)-e' '*: :_files' )
  _arguments -s -S -A '-?*' : $args && ret=0
else
  cmd="$words[1]"
  args+=(
    '(-e --edit 1 *)'{-e,--edit}'[edit files instead of running a command]' \
    '(-s --shell)'{-s,--shell}'[run shell as the target user; a command may also be specified]' \
    '(-i --login)'{-i,--login}'[run login shell as the target user; a command may also be specified]' \
    '(-b --background -i --login -s --shell -e --edit)'{-b,--background}'[run command in the background]' \
    '(--preserve-env -i --login -s --shell -e --edit)-E[preserve user environment when running command]' \
    '(-E -i --login -s --shell -e --edit)--preserve-env=-[preserve user environment when running command]::environment variable:_sequence _parameters -g "*export*"' \
    '(-H --set-home -i --login -s --shell -e --edit)'{-H,--set-home}"[set HOME variable to target user's home dir]" \
    '(-P --preserve-groups -i --login -s --shell -e --edit)'{-P,--preserve-groups}"[preserve group vector instead of setting to target's]" \
    '*: :->normal'
  )
  # sudo allows environment-variable assignments to be interspersed with
  # options. this pattern represents the lax method it uses to determine whether
  # an arg is an assignment
  _arguments -s -S -A '(-?*|[^/=]*=*)' : $args && ret=0
fi

if [[ $state = normal ]]; then
  _comp_priv_prefix=(
    $cmd -n
    ${(kv)opt_args[(I)(-[ugHEP]|--(user|group|set-home|preserve-env|preserve-groups))]}
  )

  # there's no special handling for assignments with -l. they're parsed as
  # normal, just ignored. we'll ignore them too
  if (( $+opt_args[-l] || $+opt_args[--list] )); then
    CURRENT=$#line words=( "${(@)line}" )
    _normal -p $service
    return
  fi

  # $line excludes words that match the -A pattern, except when it's the current
  # word. so if $#line == 1 we haven't seen a command yet
  if (( $#line == 1 )) && [[ $PREFIX == *=* ]]; then
    compstate[parameter]="${PREFIX%%\=*}"
    compset -P 1 '*='
    _value && ret=0
  elif (( $#line == 1 )); then
    curcontext="${curcontext%:*:*}:-command-:"
    _alternative \
      'commands:: _command_names -e' \
      'parameters: :_parameters -g "*export*~*readonly*" -qS=' && ret=0
  else
    CURRENT=$#line words=( "${(@)line}" )
    _normal && ret=0
  fi
fi

return ret
