#!/bin/zsh
zmodload -F zsh/zpty b:zpty
zmodload -F zsh/parameter p:funcstack p:functions p:parameters
zmodload -F zsh/system b:sysopen p:sysparams
zmodload -F zsh/zselect b:zselect
zmodload -F zsh/terminfo b:echoti p:terminfo
zmodload -F zsh/zutil b:zparseopts
builtin autoload -RUz \
    add-zle-hook-widget \
    is-at-least

typeset -gU FPATH fpath=( ~autocomplete/z-async $fpath[@] )
builtin autoload -Uz z-async

typeset -g ZSH_AUTOSUGGEST_USE_ASYNC=yes

${0}:precmd() {
  [[ -v ZSH_AUTOSUGGEST_IGNORE_WIDGETS ]] &&
      ZSH_AUTOSUGGEST_IGNORE_WIDGETS+=(
          history-incremental-search-backward
          recent-paths
          .zasync.fd-callback
          .autocomplete:async:complete:callback
      )

  # Start names with `.` to avoid getting wrapped by syntax highlighting.
  builtin zle -N .autocomplete:async:pty:zle-widget
  builtin zle -C .autocomplete:async:pty:completion-widget list-choices .autocomplete:async:pty:completion-widget

  builtin zle -N .autocomplete:async:wait:callback
  builtin zle -N .autocomplete:async:complete:callback

  builtin zle -C ._list_choices list-choices .autocomplete:async:list-choices:completion-widget

  builtin zle -N history-incremental-search-backward .autocomplete:async:toggle-context
  builtin zle -N recent-paths .autocomplete:async:toggle-context

  add-zle-hook-widget line-init .autocomplete:async:reset-context
  add-zle-hook-widget line-pre-redraw .autocomplete:async:complete
  add-zle-hook-widget line-finish .autocomplete:async:clear

  add-zle-hook-widget isearch-update .autocomplete:async:isearch-update
  add-zle-hook-widget isearch-exit .autocomplete:async:isearch-exit
}

.autocomplete:async:toggle-context() {
  if [[ $curcontext == $WIDGET* ]]; then
    unset curcontext
  else
    typeset -g curcontext=${WIDGET}:::
  fi
  zle .autocomplete:async:complete -w
}

.autocomplete:async:reset-context() {
  .autocomplete:async:reset-state

  typeset -g curcontext=
  builtin zstyle -s :autocomplete: default-context curcontext

  .autocomplete:async:complete
  return 0
}

.autocomplete:async:isearch-update() {
  typeset -gi _autocomplete__isearch=1
}

.autocomplete:async:isearch-exit() {
  unset _autocomplete__isearch
}

.autocomplete:async:save-state() {
  typeset -g \
      _autocomplete__curcontext=$curcontext \
      _autocomplete__lbuffer="$LBUFFER" \
      _autocomplete__rbuffer="$RBUFFER"
}

.autocomplete:async:same-state() {
  [[ -v _autocomplete__curcontext && $_autocomplete__curcontext == $curcontext &&
     -v _autocomplete__lbuffer && $_autocomplete__lbuffer == $LBUFFER &&
     -v _autocomplete__rbuffer && $_autocomplete__rbuffer == $RBUFFER ]]
}

.autocomplete:async:reset-state() {
  unset \
      _autocomplete__curcontext \
      _autocomplete__lbuffer \
      _autocomplete__rbuffer
}

.autocomplete:async:complete() {
  if [[ -v _autocomplete__inserted ]]; then
    unset _autocomplete__inserted
    typeset -g curcontext=
    builtin zstyle -s :autocomplete: default-context curcontext
  fi
  .autocomplete:async:save-state

  .autocomplete__zle-flags ||
      return 0

  (( KEYS_QUEUED_COUNT || PENDING )) &&
      return

  [[ -v ZSH_AUTOSUGGEST_IGNORE_WIDGETS ]] && (( ZSH_AUTOSUGGEST_IGNORE_WIDGETS[(I)$LASTWIDGET] )) &&
      unset POSTDISPLAY

  # Don't get activated by asynchronous widgets.
  [[ $LASTWIDGET == (autosuggest-suggest|.autocomplete:async:*:callback|.zasync.fd-callback) ]] &&
      return 0

  {
    if (( REGION_ACTIVE )) ||
        [[ -v _autocomplete__isearch && $LASTWIDGET == *(incremental|isearch)* ]]; then
      builtin zle -Rc
      return 0
    fi

    if [[ ${LASTWIDGET##.} == (up-line-or-search|history-search-backward) &&
        _lastcomp[nmatches] -lt 1 ]]; then
      return 0
    fi

    builtin zstyle -t ":autocomplete:${LASTWIDGET}:" ignore &&
        return 0

    local -Pa ignored=(
      '_complete_help'
      '(copy|insert)-*-word'
      'describe-key-briefly'
      '(|reverse-)menu-complete'
      'what-cursor-position'
      'where-is'
    )
    [[ ${LASTWIDGET##.} == (${(~j:|:)~ignored}) ]] &&
        return 0

    [[ $KEYS == ([\ -+*]|$'\e\t') ]] &&
        builtin zle -Rc

    # WORKAROUND: #549 Bug in zdharma/fast-syntax-highlighting.
    [[ -v _FAST_MAIN_CACHE ]] &&
        _zsh_highlight

    typeset -ga _autocomplete__region_highlight=( "$region_highlight[@]" )

    if  [[ -v ZSH_AUTOSUGGEST_IGNORE_WIDGETS ]] &&
        (( ZSH_AUTOSUGGEST_IGNORE_WIDGETS[(I)$LASTWIDGET] )); then
      unset POSTDISPLAY
    fi

    z-async cancel complete
    z-async start wait .autocomplete:async:wait-worker .autocomplete:async:wait:callback
  }

  return 0
}

.autocomplete:async:clear() {
  z-async cancel wait
  z-async cancel complete
  unset curcontext _autocomplete__isearch

  .autocomplete:async:reset-context
  builtin zle -Rc
  return 0
}

.autocomplete:async:wait-worker() {
  local -F seconds=
  builtin zstyle -s :autocomplete: delay seconds ||
      builtin zstyle -s :autocomplete: min-delay seconds ||
      (( seconds = 0.05 ))

  (( seconds = max( 0, seconds ) ))

  # Convert to 100ths of a second for `zselect -t`.
  # WORKAROUND: #441 Directly using $(( [#10] … max( … ) )) leads to 0 in Zsh 5.9, as the result
  # of max() gets converted to an integer _before_ being multiplied.
  local -i timeout=$(( 100 * seconds ))

  zselect -t $timeout
}

.autocomplete:async:wait:callback() {
  (( YANK_ACTIVE )) &&
      return 0

  (( KEYS_QUEUED_COUNT || PENDING )) &&
      return

  .autocomplete:async:same-state &&
      .autocomplete:async:start
  return 0
}

.autocomplete:async:start() {
  z-async start complete .autocomplete:async:start-worker .autocomplete:async:complete:callback

  # WORKAROUND: https://github.com/zsh-users/zsh-autosuggestions/issues/364
  # There's a weird bug in Zsh < 5.8, where ^C stops working unless we force a fork.
  command true
}

.autocomplete:async:start-worker() {
  local +h PS4=$_autocomplete__ps4
  .autocomplete:async:start:inner 2>>| $_autocomplete__log
}

.autocomplete:async:start:inner() {
  {
    typeset -F SECONDS=0

    local -P hooks=( chpwd periodic precmd preexec zshaddhistory zshexit )
    builtin unset ${^hooks}_functions &> /dev/null
    $hooks[@] () { : }

    local -P hook=
    for hook in \
        zle-{isearch-{exit,update},line-{pre-redraw,init,finish},history-line-set,keymap-select}
    do
      builtin zle -N $hook .autocomplete:async:pty:no-op
    done
    {
      local REPLY=
      zpty AUTOCOMPLETE .autocomplete:async:pty
      local -Pi fd=$REPLY

      zpty -w AUTOCOMPLETE $'\C-@'

      local header=
      zpty -r AUTOCOMPLETE header $'*\C-A'

      local -a reply=()
      local text=

      local -F seconds=0.0
      builtin zstyle -s ":autocomplete:${curcontext}" timeout seconds ||
          (( seconds = 1.0 ))

      (( seconds = max( 0, seconds - SECONDS ) ))

      # Convert to 100ths of a second for `zselect -t`.
      # WORKAROUND: #441 Directly using $(( [#10] … max( … ) )) leads to 0 in Zsh 5.9, as the result
      # of max() gets converted to an integer _before_ being multiplied.
      local -i timeout=$(( 100 * seconds ))

      if zselect -rt $timeout "$fd"; then
        zpty -r AUTOCOMPLETE text $'*\C-B'
      else
        # Press ^C twice: Once to abort completion, then once to abort the command line.
        # Then exit the shell with ^D.
        zpty -wn AUTOCOMPLETE $'\C-C\C-C\C-D'
      fi
    } always {
      zpty -d AUTOCOMPLETE
    }
  } always {
    # Always produce output, so z-async always fires the callback.
    # Strip the leading terminal output; keep only the trailing integer (list_lines count).
    print -r -- "${${text%$'\C-B'}##*[^0-9]}"
  }
}

.autocomplete:async:pty() {
  typeset -g +h PS4=pty:$_autocomplete__ps4

  # Force the shell to exit on timeout.
  local -i seconds=
  builtin zstyle -s ":autocomplete:${curcontext}" timeout seconds ||
      seconds=1
  TMOUT=$(( [#10] 1 + seconds ))
  TRAPALRM() {
    builtin exit
  }

  builtin bindkey $'\C-@' .autocomplete:async:pty:zle-widget
  local __tmp__=
  builtin vared __tmp__
} 2>>| $_autocomplete__log

.autocomplete:async:pty:no-op() {
  :
}

.autocomplete:async:pty:zle-widget() {
  setopt localoptions NO_banghist

  local -a _autocomplete__comp_mesg=()
  local -i _autocomplete__list_lines=0
  local _autocomplete__mesg=
  {
    # The completion widget sometimes returns without calling its function. So, we need to print all
    # our control characters here, to ensure we don't end up waiting endlessly to read them.
    print -n -- '\C-A'
    LBUFFER=$_autocomplete__lbuffer
    RBUFFER=$_autocomplete__rbuffer

    [[ -n $curcontext ]] &&
        setopt $_autocomplete__ctxt_opts[@]

    builtin zle .autocomplete:async:pty:completion-widget -w 2>>| $_autocomplete__log
  } always {
    print -rNC1 -- ${_autocomplete__list_lines:-0}$'\C-B'
    builtin exit
  }
} 2>>| $_autocomplete__log

.autocomplete:async:pty:completion-widget() {
  setopt localoptions banghist
  {
    if ! .autocomplete:async:sufficient-input; then
      return
    fi
    {
      unfunction compadd 2> /dev/null
      unset 'compstate[vared]'
      .autocomplete:async:list-choices:main-complete
    } always {
      _autocomplete__list_lines=$compstate[list_lines]
    }
  }
} 2>>| $_autocomplete__log

.autocomplete:async:complete:callback() {
  setopt localoptions NO_banghist

  (( YANK_ACTIVE )) &&
      return 0

  (( KEYS_QUEUED_COUNT || PENDING )) &&
      return

  .autocomplete:async:same-state ||
      return 0

  local -i list_lines="$( z-async reply complete )"

  [[ -n $curcontext ]] &&
      setopt $_autocomplete__ctxt_opts[@]

  # If a widget can't be called, zle always returns true.
  # Thus, we return false on purpose, so we can check if our widget got called.
  local +h PS4=zle:$_autocomplete__ps4
  if ! builtin zle ._list_choices -w "$list_lines" 2>>| $_autocomplete__log; then

    typeset -g region_highlight=( "$_autocomplete__region_highlight[@]" )

    # Need to call this here, because on line-pre-redraw, $POSTDISPLAY is empty.
    [[ -v functions[_zsh_autosuggest_highlight_apply] ]] &&
        _zsh_autosuggest_highlight_apply


    local clear=
    [[ -z $_lastcomp[list] ]] &&
      clear=c  # Force-clear the old list if the new one is empty.

    # Refresh if and only if our widget got called. Otherwise, Zsh will crash (eventually).
    builtin zle -R$clear
  fi
  .autocomplete:async:reset-state

  return 0
}

.autocomplete:async:sufficient-input() {
  local min_input=
  if ! builtin zstyle -s ":autocomplete:${curcontext:-list-choices}:" min-input min_input; then
    if [[ -n $curcontext ]]; then
      min_input=0
    else
      min_input=1
    fi
  fi

  local ignored=
  builtin zstyle -s ":autocomplete:${curcontext:-list-choices}:" ignored-input ignored

  if (( ${#words[@]} == 1 && ${#words[CURRENT]} < min_input )) ||
      [[ -n $ignored && $words[CURRENT] == $~ignored ]]; then
    compstate[list]=
    false
  else
    true
  fi
}

.autocomplete:async:list-choices:completion-widget() {
  local +h PS4=$_autocomplete__ps4

  setopt localoptions banghist

  if [[ $1 != <1-> ]]; then
    compstate[list]=
    return
  fi

  .autocomplete:async:sufficient-input ||
      return 2

  .autocomplete:async:list-choices:main-complete

  # Workaround: In Zsh <= 5.9.0, comppostfuncs don't get called after completing subscripts.
  unset MENUSELECT MENUMODE
  compstate[insert]=
  _lastcomp[insert]=
  compstate[pattern_insert]=
  _lastcomp[pattern_insert]=
  if [[ -v _autocomplete__partial_list ]]; then
    builtin compadd -J -last- -x '%F{0}%K{12}(MORE)%f%k'
    _lastcomp[list_lines]=$compstate[list_lines]
  fi

  if [[ -n $compstate[exact_string] && -z $_lastcomp[tags] &&
        compstate[nmatches] -eq 1 && compstate[list_lines] -eq 1 ]]; then
    # WORKAROUND: _arguments adds $compstate[exact_string] as a completion w/out description.
    compstate[list]=
    _lastcomp[list]=
  fi

  return 2  # Don't return 1, to prevent beeping.
}

.autocomplete:async:list-choices:max-lines() {
  local -Pi max_lines
  builtin zstyle -s ":autocomplete:${curcontext}:" list-lines max_lines ||
      max_lines=16
  _autocomplete__max_lines=$(( min( max_lines, LINES - BUFFERLINES - 1 ) ))
}

.autocomplete:async:list-choices:main-complete() {
  local -i _autocomplete__max_lines

  case $curcontext in
  *history-* )
    autocomplete:_main_complete:new - history-lines _autocomplete__history_lines
  ;;
  recent-paths:* )
    autocomplete:_main_complete:new - recent-paths _autocomplete__recent_paths
  ;;
  * )
    {
      local curcontext=list-choices:::
      () {
        emulate -L zsh
        setopt $_autocomplete__func_opts[@]

        .autocomplete:async:shadow compadd

        autoload -Uz +X _describe
        .autocomplete:async:shadow _describe
      } "$@"

      .autocomplete:async:list-choices:max-lines
      autocomplete:_main_complete:new "$@"
    } always {
      unfunction compadd comptags 2> /dev/null
      .autocomplete:async:unshadow compadd
      .autocomplete:async:unshadow _describe
    }
  ;;
  esac
}

.autocomplete:async:shadow() {
  [[ -v functions[$1] ]] &&
      functions[autocomplete:async:${1}:old]="$functions[$1]"
  functions[$1]="$functions[.autocomplete:async:$1]"
}

.autocomplete:async:unshadow() {
  if [[ -v functions[autocomplete:async:${1}:old] ]]; then
    functions[$1]="$functions[autocomplete:async:${1}:old]"
    unfunction autocomplete:async:${1}:old
  fi
}

.autocomplete:async:_describe() {
  local -i _autocomplete__described_lines=1  # Assume we'll add a title.
  autocomplete:async:_describe:old "$@"
}

.autocomplete:async:compadd() {
  setopt localoptions multibyte

  local -A _opts_=()
  local -a _xopts_=() _displ_=() _matches_=()
  local -P _displ_name_= _matches_name_=

  zparseopts -A _opts_ -E -- D: E: O:

  local -Pi _unused_lines_=$(( _autocomplete__max_lines - compstate[list_lines] ))

  # If $_grp is not set, then _describe is adding completions in a normal way and we don't need to
  # do all this.
  if [[ -v _autocomplete__described_lines && -n $_grp ]]; then

    # We cannot interfere when _describe is actually adding the completions or we risk breaking the
    # layout.
    if [[ -z $_opts_[-D] ]]; then
      builtin compadd "$@"
      return
    fi

    _displ_name_=$_opts_[-D]
    _matches_name_=$_opts_[-O]

    # We already ran out of space.
    if [[ -v _autocomplete__partial_list ]]; then
      set -A $_displ_name_
      [[ -n $_matches_name_ ]] &&
          set -A $_matches_name_
    fi

    builtin compadd "$@"
    local -Pi _ret_=$?

    local -i _ndisplay_=${(PA)#_displ_name_}
    local -i _lines_left_for_describe_=$(( _unused_lines_ - _autocomplete__described_lines ))

    # The number of lines that would be added is equal to the number of unique display strings.
    if (( ${#${(u)${(PA)_displ_name_}[@]#*:}} > _lines_left_for_describe_ )); then
      local -Pi _matches_to_remove=$(( _ndisplay_ - max( 0, _lines_left_for_describe_ ) ))
      if (( _matches_to_remove < _ndisplay_ )); then
        shift -p $_matches_to_remove $_displ_name_ $_matches_name_
      else
        set -A $_displ_name_
        [[ -n $_matches_name_ ]] &&
            set -A $_matches_name_
      fi
      _ndisplay_=${(PA)#_displ_name_}
      .autocomplete:async:compadd:disable
    fi
    (( _autocomplete__described_lines += _ndisplay_ ))

    return _ret_
  fi

  # We already ran out of space.
  if [[ -v _autocomplete__partial_list ]]; then
    [[ -n $_opts_[-D] ]] &&
        set -A $_opts_[-D]
    [[ -n $_opts_[-O] ]] &&
        set -A $_opts_[-O]
    return 1
  fi

  # If
  if [[ -n $_opts_[-D]$_opts_[-O] ]]; then
    builtin compadd "$@"
    return
  fi

  local -i _old_total_lines=$compstate[list_lines]

  # +: Keep all occurences. Otherwise, only last occurence is kept.
  zparseopts -a _xopts_ -D -E -- X+: x+:

  if (( $#_xopts_ )); then
    # Use only the first one, because subsequent ones are ignored by compadd anyway.
    local original="$_xopts_[2]"
    local formatted="${(%g:ce:)original}"

    # Trim text to one line and omit ANSI sequences to avoid crash in command substitution.
    local ansi_start=${(M)formatted##($'\C-[['[;[:digit:]]#m)#}
    local ansi_end=${(M)formatted%%($'\C-[['[;[:digit:]]#m)#}
    local text=${${formatted#$ansi_start}%$ansi_end}
    _xopts_=(
      "$_xopts_[1]"
      "${(mr:COLUMNS-1:)text%%[[:space:]]#}"
    )
  fi

  local -Pi _total_new_lines_="$(
      builtin compadd "$_xopts_[@]" "$@"
      print -nr -- $(( $compstate[list_lines] - _old_total_lines ))
  )"

  if (( $#_xopts_ )); then
    # Replace original text with trimmed text.
    _xopts_[2]="${original/$text/$_xopts_[-1]}"
  fi

  # Everything fits.
  if (( _total_new_lines_ + $compstate[list_lines] <= _autocomplete__max_lines )); then
    builtin compadd "$_xopts_[@]" "$@"
    return
  fi

  local -Pi _new_completion_lines_="$(
      builtin compadd "$@"
      print -nr -- $(( $compstate[list_lines] - _old_total_lines ))
  )"

  local -a _dopt_=()
  zparseopts -a _dopt_ -D -E -- d: ld:

  _displ_name_=$_dopt_[2]

  # Collect all matching completions and their display strings (if any).
  local -a _Dopt_=()
  [[ -n $_displ_name_ ]] &&
      _Dopt_=( -D $_displ_name_ )
  builtin compadd -O _matches_ $_Dopt_ "$@"

  # If we don't have an array with display strings, then create one.
  if [[ -z $_displ_name_ ]]; then
    _displ_name_=_displ_
    _displ_=( "$_matches_[@]" )
    _dopt_=( -d $_displ_name_ )
  fi

  local -Pi _nmatches_per_line_=$(( 1.0 * $#_matches_ / _new_completion_lines_ ))

  # If we need more than one line per match, then make each match fit exactly one line.
  if (( _nmatches_per_line_ < 1 )); then
    # WORKAROUND: Zsh mistakenly treats display strings that are exactly $COLUMNS wide as not
    # fitting on one line.
    set -A $_displ_name_ ${(@mr:COLUMNS-1:)${(PA)_displ_name_}[@]//$'\n'/\n}

    _dopt_=( -ld $_displ_name_ )
    (( _nmatches_per_line_ = 1 ))
  fi

  local -Pi _new_heading_lines_=$(( _total_new_lines_ - _new_completion_lines_ ))

  # Need to round this down _before_ subtracting it or it will effectively be rounded up.
  local -Pi _nmatches_that_fit_=$((
      ( _unused_lines_ - _new_heading_lines_ ) * _nmatches_per_line_
  ))

  local -Pi _nmatches_to_remove_=$(( $#_matches_ - max( 0, _nmatches_that_fit_ ) ))

  if (( _nmatches_to_remove_ > 0 )); then
    # If we're going to remove anything, then we need to make room for the `(MORE)` prompt.
    (( _nmatches_to_remove_ = min( ++_nmatches_to_remove_, $#_matches_ ) ))

    .autocomplete:async:compadd:disable

    # Make sure we always add a headline, even when we don't add any matches, because a single
    # '(MORE)' without other content doesn't make sense.
    (( _nmatches_to_remove_ >= $#_matches_ )) && _xopts_[1]=-x

    shift -p $_nmatches_to_remove_ _matches_ $_displ_name_
  fi

  _autocomplete__compadd_opts_len "$@"
  builtin compadd "$_xopts_[@]" "$_dopt_[@]" -a "$@[1,?]" _matches_
}

.autocomplete:async:compadd:disable() {
  typeset -g _autocomplete__partial_list=$curtag
  comptags() { false }  # Stop completion from trying more tags.
}
