#autoload
# Intended to be called during accept-line or in zle-line-finish, but can
# be called by any widget to apply correction to all words in $BUFFER or
# used as a widget itself

# Compare modify-current-argument

setopt localoptions noksharrays multibyte norecexact
zmodload zsh/complist || return 1

local -a reply cmdline
local key REPLY REPLY2 MENUSELECT
integer pos posword poschar
unset MENUSELECT	# Bug with no-select plus yes=2 below

local curcontext="${curcontext:-}"
local widget="correct-${${WIDGET/correct-all-words/}:-all-words}"
if [[ -z "$curcontext" ]]; then
  curcontext="${widget}:::"
else
  curcontext="${widget}:${curcontext#*:}"
fi
local mycontext="${curcontext}"

# This breaks out of read-from-minibuffer if only the original matches
local shown_original
zstyle -e ":completion:${widget}:*" original \
       'if (( compstate[nmatches] == 0 )); \
       then reply=(false);
       elif [[ ${compstate[unambiguous]} = ${key} ]]; \
       then shown_original=unambiguous; zle -U n; reply=(false); \
       elif [[ -z ${shown_original} ]];
       then shown_original=original; reply=(true);
       else reply=(false); fi'

# Would be nice to make these conditional, but it's hard to test for
# these specifically if a more general context has the style defined
zstyle ":completion:${widget}:*" menu no-select yes=2
zstyle ":completion:${widget}:*" group-name ''
zstyle ":completion:${widget}:*" group-order original corrections
zstyle ":completion:${widget}:*" show-ambiguity true
zstyle ":completion:${widget}:*" show-completer true
zstyle ":completion:${widget}:*" accept-exact false

# Keep completion functions out of the results
zstyle ":completion:${widget}:*" ignored-patterns '_*'

# This shows earliest the words with the fewest necessary corrections
zstyle ":completion:${widget}:*" sort false

# Overload fake description to produce a prompt for multiple corrections
zstyle -e ":completion:${widget}:*:original" fake \
	 'if (( compstate[nmatches] > 1 )) && \
	     [[ -z ${shown_original} ]] ; \
	 then bindkey -M correctall y _correct_word; \
	 shown_original=fake-original; \
	 reply=("${key//:/\\:}:TAB to choose, ENTER to accept, n to skip"); \
	 fi'
zstyle -e ":completion:${widget}:*:corrections" fake \
	 'if [[ ${shown_original} = wanted ]] || \
	     ( (( compstate[nmatches] > 0 )) && \
	       [[ -z ${shown_original} ]] ); \
	 then bindkey -M correctall y _correct_word; \
	 shown_original=fake-corrections; \
	 reply+=("${key//:/\\:}:TAB to choose, ENTER to accept, n to skip"); \
	 fi'

# There's no good semantics for 'y' when there are multiple possible
# corrections.  Left as $'\t\n' it'll skip ahead after a correction is
# chosen from the list.  Changed to .accept-line it becomes equivalent
# to 'n' if no choice has been made yet.  The above makes it cycle the
# menu, but maybe it would be better just to have it do nothing?

autoload -Uz split-shell-arguments read-from-minibuffer mkshadow

split-shell-arguments
[[ ${#reply} -lt 2 ]] && return 1
(( posword = REPLY, poschar = REPLY2 ))
cmdline=("${reply[@]}")		# In case something else uses $reply ...

bindkey -N correctall
bindkey -M correctall $'\t' _correct_word
bindkey -M correctall ' ' _correct_word
for key in n a e \! $'\n' $'\r'
do
  bindkey -M correctall $key .accept-line
done
bindkey -M correctall u .undo
bindkey -M correctall '^_' .undo
bindkey -M correctall -s '^G' e
bindkey -M correctall -s '^U' a		# Should copy from main XXX
bindkey -M correctall -s y $'\t\n'

# Work around a bug with recursive-edit from zle-line-finish
[[ $WIDGET = zle-line-finish ]] && zle recursive-edit	# Fails

local lmini rmini
{
  mkshadow -s all-words _original_file _correct_word

  # Force unedited original to precede _path_files additions.
  # Otherwise the "original" style above handles this.
  function _original_file {
    [[ ${key} = */* ]] && shown_original=wanted
    return 1	# Force call to _correct
  }

  # This is to avoid having _correct_word stomp on $curcontext,
  # plus break out of read-from-minibuffer when nothing matches
  function _correct_word {
    local shown_original ret REPLY
    _main_complete _original_file _correct
    ret=$?
    [[ ${compstate[nmatches]} -eq 0 ]] && zle -U n
    return ret
  }

  # "Real" words from reply[2], reply[4], etc., see split-shell-arguments
  for pos in {2..${#cmdline}..2}
  do
    # Don't try to correct numbers and non-syntax punctuation
    [[ ${cmdline[pos]} = *[A-Za-z]* ]] || continue
    
    {
      key=${cmdline[pos]}
      curcontext="${mycontext}"

      # Can't use -K KEYMAP because read-from-minibuffer always resets
      bindkey -A main llatcerroc
      bindkey -A correctall main

      # lmini="${PREBUFFER}${(j::)cmdline[1,pos-1]}"	# Too much?
      lmini=${(j::)cmdline[1,pos-1]}
      rmini=${(j::)cmdline[pos+1,-1]}

      zle -U $'\t'	# Start minibuffer in completion
      read-from-minibuffer "Correct $key [nyae!]: " "${lmini}${key}" "${rmini}"
      REPLY=${${REPLY#${lmini}}%${rmini}}
    } always { bindkey -A llatcerroc main }
    case ${KEYS} in
      (a) zle send-break; break;;
      (n) continue;;
      (y|$'\n'|$'\r') cmdline[pos]=${REPLY:-${cmdline[pos]}};;
      (e) [[ ${WIDGET} = zle-line-finish ]] && {
	    print -z "${(j::)cmdline}"
	    cmdline=()
	    (( posword = 0, poschar = 0 ))
	  }
	  ;&
      (*) break;;
    esac
  done
} always {
  rmshadow
  zle -R -c
}

BUFFER="${(j::)cmdline}"
CURSOR=$(( ${#${(j::)cmdline[1,posword-1]}} + poschar ))

[[ "${KEYS}" = \! ]] && zle .accept-line