#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

while (( $#state )); do
  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 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
  shift 1 state
  shift 1 context
done

return ret