#compdef zle

local expl ret=1
local -a opts compwids state line context

compwids=(accept-and-menu-complete
	  complete-word
	  delete-char-or-list
	  expand-or-complete
	  expand-or-complete-prefix
	  list-choices
	  menu-complete
	  menu-expand-or-complete
	  reverse-menu-complete)

opts=(-A -C -D -L -M -N -R -U -a -c -l \* :)

_arguments -s -S \
    "($opts)-A[define widget alias]:old widget:->widget :new widget:->widget" \
    "($opts)-C[define completion widget]:new widget name:->comp-widget :completion widget:->builtin-comp-widget :widget shell function:->function" \
    "($opts)-D[delete widget]:*:widget:->widget" \
    \(${(j. .)opts:#-l}')-L[with -l, list as commands]' \
    \(${(j. .)opts:#-[La]}')-l+[list user-defined widgets]:*:-:->listing' \
    \(${(j. .)opts:#-l}')-a[with -l, list all widgets]' \
    "($opts)-M[display message]:message: " \
    "($opts)-N[define new widget]:widget name:->widget-or-function ::widget shell function:->function" \
    \(${(j. .)opts:#-c}')-R+[redisplay]:*:-:->redisplay' \
    \(${(j. .)opts:#-R}')-c[with -R, clear listing]' \
    "($opts)-U[unget to input stack]:string: " \
    '(-):widget:->widget' \
    '(-)*::widget args:->args' && ret=0

[[ $state == listing ]] &&
  _arguments -s -S '!-l' \
    "(-a)-L[list as commands]" \
    "(-L)-a[list all widgets]" \
    '*:widget name:->widget' && ret=0

case "$state[1]" in
  (args)
    _arguments \
      '(-N)-n[numeric prefix]:number:' \
      '(-n)-N[reset numeric prefix]' \
      '(-)*:widget arguments: ' && ret=0
    ;;
  (widget*)
    _wanted -C "$context[1]" widgets expl "${state_descr[1]:-widget}" compadd -k widgets && ret=0
    ;&
  (function)
    [[ $state[1] != *function ]] ||	# Handle fall-through
    _wanted -C "$context[1]" functions expl 'widget shell function' \
      compadd -k functions && ret=0
    ;;
  (comp-widget)
    _wanted -C "$context[1]" widgets expl 'completion widget' \
      compadd -k "widgets[(R)(*:|)(.|)(${(j(|))compwids})(|:*)]" && ret=0
    ;&
  (builtin-comp-widget)
    _wanted -C "$context[1]" widgets expl 'builtin completion widget' \
      compadd -k "widgets[(I)(.|)(${(j(|))compwids})]" && ret=0
    ;;
  (redisplay)
    _arguments -s -S '!-R' \
      "-c[clear listing]" \
      ":status line" "*:strings to list" && ret=0
    ;;
esac

return ret