#!/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 -g ZSH_AUTOSUGGEST_USE_ASYNC=yes

.autocomplete:async:precmd() {
  [[ -v ZSH_AUTOSUGGEST_IGNORE_WIDGETS ]] &&
      ZSH_AUTOSUGGEST_IGNORE_WIDGETS+=(
          history-incremental-search-backward
          .autocomplete:async:complete:fd-widget
      )

  # 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:complete:fd-widget
  builtin zle -N .autocomplete:async:wait:fd-widget

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

  builtin zle -N history-incremental-search-backward .autocomplete:async:history-incremental-search

  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:history-incremental-search() {
  if [[ $curcontext == $WIDGET* ]]; then
    unset curcontext
  else
    typeset -g curcontext=${WIDGET}:::
  fi
  .autocomplete:async:complete
}

.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() {
  .autocomplete:zle-flags $LASTWIDGET
  unset _autocomplete__isearch
}

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

.autocomplete:async:same-state() {
  [[ -v _autocomplete__curcontext && $_autocomplete__curcontext == $context &&
      -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() {
  .autocomplete:zle-flags $LASTWIDGET ||
      return 0

  (( KEYS_QUEUED_COUNT || PENDING )) &&
      return
  {
    if [[ -v DEBUG ]]; then
      set -x -o localoptions
      local +h PS4=$_autocomplete__ps4
    fi

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

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

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

    # Completion already started
    [[ -v _autocomplete__lbuffer || -v _autocomplete__rbuffer ]] &&
        return 0

    # Don't get triggered by asynchronous widgets.
    if [[ $LASTWIDGET == (autosuggest-suggest|.autocomplete:async:*:fd-widget) ]]; then
      return 0
    fi

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

    [[ $LASTWIDGET ==
          (_complete_help|(|.)(describe-key-briefly|(|.)(|reverse-)menu-complete|what-cursor-position|where-is)) ]] &&
        return

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

    .autocomplete:async:wait

    set +x
  } 2>>| $_autocomplete__log

  return 0
}

.autocomplete:async:clear() {
  unset curcontext _autocomplete__isearch
  .autocomplete:async:reset-context
  builtin zle -Rc
  return 0
}

.autocomplete:async:wait() {
  local fd=

  .autocomplete:async:save-state

  sysopen -r -o cloexec -u fd <(
    local -F seconds=
    builtin zstyle -s :autocomplete: delay seconds ||
        builtin zstyle -s :autocomplete: min-delay seconds ||
        (( seconds = 0.04 ))

    # 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 * max( 0, seconds ) ))
    zselect -t $timeout

    print
  )
  builtin zle -Fw "$fd" .autocomplete:async:wait:fd-widget

  return 0
}

.autocomplete:async:wait:fd-widget() {
  {
    local -i fd=$1
    builtin zle -F $fd # Unhook ourselves immediately, so we don't get called more than once.
    exec {fd}<&-
    .autocomplete:zle-flags

    {
      if [[ -v DEBUG ]]; then
        set -x -o localoptions
        local +h PS4=$_autocomplete__ps4
      fi

      if .autocomplete:async:same-state; then
        .autocomplete:async:start
      else
        .autocomplete:async:wait
      fi
      set +x
    } 2>>| $_autocomplete__log
  }
  return 0
}

.autocomplete:async:start() {
  .autocomplete:async:save-state

  local fd=
  sysopen -r -o cloexec -u fd <(
    .autocomplete:async:start:inner 2>>| $_autocomplete__log
  )
  builtin zle -Fw "$fd" .autocomplete:async:complete:fd-widget

  # 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:inner() {
  {
    if [[ -v DEBUG ]]; then
      set -x
      PS4=$_autocomplete__ps4
    fi

    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=

      builtin zstyle -s :autocomplete: timeout seconds ||
          (( seconds = 0.4 ))

      # 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 * max( 0, seconds - 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 we always reach the callback, so we can close
    # the fd.
    print -rNC1 -- "${text%$'\C-B'}"
  }
}

.autocomplete:async:pty() {
  if [[ -v DEBUG ]]; then
    functions -T .autocomplete:async:pty:zle-widget
    set -x -o localoptions
    PS4=$_autocomplete__ps4
  fi

  # Make sure this shell dies after it times out.
  local -F seconds=
  builtin zstyle -s :autocomplete: timeout seconds ||
      seconds=0.5
  TMOUT=$(( [#10] 1 + seconds ))

  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() {
  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

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

    [[ -v DEBUG ]] &&
        functions -T .autocomplete:async:pty:completion-widget

    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() {
  {
    if ! .autocomplete:async:sufficient-input; then
      return
    fi
    {
      unset 'compstate[vared]'
      .autocomplete:async:list-choices:main-complete
    } always {
      _autocomplete__list_lines=$compstate[list_lines]
    }
  } 2>>| $_autocomplete__log
}

.autocomplete:async:complete:fd-widget() {
  {
    if [[ -v DEBUG ]]; then
      local +h PS4=$_autocomplete__ps4
      set -x -o localoptions
      functions -T .autocomplete:async:list-choices:completion-widget
    fi

    local -i fd=$1
    {
      builtin zle -F $fd # Unhook ourselves immediately, so we don't get called more than once.

      if ! .autocomplete:zle-flags; then
        .autocomplete:async:reset-state
        return 0
      fi

      if ! .autocomplete:async:same-state; then
        .autocomplete:async:start
        return 0
      fi

      .autocomplete:async:reset-state

      local -a reply=()
      IFS=$'\0' read -rAu $fd
      shift -p reply
    } always {
      exec {fd}<&-
    }

    setopt $_autocomplete__comp_opts[@]
    [[ -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.
    if ! builtin zle ._list_choices -w "$reply[@]" 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

      # Refresh if and only if our widget got called. Otherwise, Zsh will crash (eventually).
      builtin zle -R
    fi
    .autocomplete:async:reset-state
    set +x
    return 0
  } 2>>| $_autocomplete__log
}

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

  local ignored=
  builtin zstyle -s ":autocomplete:${curcontext}:" 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() {
  if [[ $1 != <1-> ]]; then
    compstate[list]=
    return
  fi

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

  compstate[insert]=
  compstate[old_list]=
  compstate[pattern_insert]=
  .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{black}%K{12}(MORE)%f%k'
    _lastcomp[list_lines]=$compstate[list_lines]
  fi

  [[ -v DEBUG ]] &&
      typeset -m _lastcomp >&2

  return 2  # Don't return 1, to prevent beeping.
} 2>>| $_autocomplete__log

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

.autocomplete:async:list-choices:main-complete() {
  local -i _autocomplete__max_lines=0 _autocomplete__described_lines=0
  if [[ $curcontext == *history-* ]]; then
    setopt $_autocomplete__func_opts[@]
    autocomplete:_main_complete:new - history-lines _autocomplete__history_lines
  else
    {
      () {
        emulate -L zsh
        setopt $_autocomplete__func_opts[@]

        local curcontext=list-choices:::
        [[ -v functions[compadd] ]] &&
            functions[autocomplete:async:compadd:old]="$functions[compadd]"
        functions[compadd]="$functions[.autocomplete:async:compadd]"

        [[ -v VERBOSE ]] &&
            functions -T compadd _describe
      } "$@"

      .autocomplete:async:list-choices:max-lines 16
      autocomplete:_main_complete:new "$@"
    } always {
      unfunction compadd comptags 2> /dev/null
      if [[ -v functions[autocomplete:async:compadd:old] ]]; then
        functions[compadd]="$functions[autocomplete:async:compadd:old]"
        unfunction autocomplete:async:compadd:old
      fi
    }
  fi
}

.autocomplete:async:compadd() {
  local -Pi _ret_=1
  local -P _displ_name_= _matches_name_=

  local -A _opts_=()
  local -a _displ_=() _dopt_=() _groupname_=() _matches_=()
  zparseopts -A _opts_ -E -- D: E: O: X: x: d+:=_dopt_ l

  if [[ -v _opts_[-x] && $# -eq 2 ]]; then
    # Adding a message only.
    builtin compadd "$@"
    return
  fi

  # -X and -x are titles.
  local -Pi _avail_list_lines_=$((
      max( 0, _autocomplete__max_lines - 1 - ${${_opts_[(i)-[Xx]]}:+1} - compstate[list_lines] )
  ))

  if (( funcstack[(I)_describe] )); then

    if [[ -n $_opts_[-D] ]]; then
      _displ_name_=$_opts_[-D]
      _matches_name_=$_opts_[-O]

      if (( _autocomplete__described_lines >= _autocomplete__max_lines )); then
        set -A $_displ_name_
        return 1

      else
        builtin compadd "$@"
        _ret_=$?

        _displ_=( ${(@PA)_displ_name_} )
        _matches_=( ${(@PA)_matches_name_} )

        if (( $#_displ_[@] )); then
          local -PaU _uniques_=( ${_displ_[@]#*:} )
          local -P _lines_too_many=$(( -1 * ( _avail_list_lines_ - _autocomplete__described_lines - $#_uniques_[@] ) ))

          (( _lines_too_many > 0 )) &&
              shift -p $_lines_too_many $_displ_name_ $_matches_name_

          (( _autocomplete__described_lines += ${(@PA)#_displ_name_} ))
        fi

        return _ret_
      fi

    elif [[ -n $_opts_[-E] ]]; then
      # End of describe loop. It will now add matches in a special way. Can't interfere with this or it will break.
      (( _autocomplete__described_lines = 0 ))  # Done with this _describe loop.
      builtin compadd "$@"
      return
    fi

    # If neither of those are true, then _describe is doing normal completion.
    (( _autocomplete__described_lines = 0 ))
  fi

  # Deleting or collecting completions, but not actually adding them.
  if [[ -n $_opts_[-D]$_opts_[-O] ]]; then
    builtin compadd "$@"
    return
  fi

  [[ -v _autocomplete__partial_list ]] &&
      return 1

  local -Pi _resulting_list_lines="$(
      builtin compadd "$@"
      print -nr -- $compstate[list_lines]
  )"
  local -Pi _new_list_lines_=$(( _resulting_list_lines - $compstate[list_lines] ))
  local -Pi _remaining_list_lines=$(( _avail_list_lines_ - _new_list_lines_ ))

  (( _remaining_list_lines <= 0 )) &&
      .autocomplete:async:compadd:disable

  if (( _remaining_list_lines >= 0 )); then
    builtin compadd "$@"
    return
  fi

  _displ_name_=$_dopt_[2]
  [[ -n $_displ_name_ ]] &&
      local -a _Dopt_=( -D $_displ_name_ )

  # Collect all matching completions and, if present, remove the display strings of those that don't match.
  builtin compadd -O _matches_ $_Dopt_ "$@"

  if [[ -z $_displ_name_ ]]; then
    _displ_=( "$_matches_[@]" )
    _dopt_=( -d _displ_ )
    _displ_name_=_displ_
  fi

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

  if (( _nmatches_per_line_ < 1 )); then
    # If we need more than one line per match, then make each match fit exactly one line.
    _nmatches_per_line_=1
    _dopt_=( -l "$_dopt_[@]" )

    # WORKAROUND: Display strings that are exactly $COLUMNS wide are mistakenly counted as spanning two lines, causing
    # the dreaded `zsh: do you wish to see all XX possibilities (YY lines)?`
    set -A $_displ_name_ ${(@r:COLUMNS-1:)${(PA)_displ_name_}[@]//$'\n'/\n}
  fi

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

  (( _nmatches_to_remove_ > 0 )) &&
      shift -p $_nmatches_to_remove_ _matches_ $_displ_name_

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

.autocomplete:async:compadd:disable() {
  # Disable all further completions.
  typeset -g _autocomplete__partial_list=$curtag
  comptags() { false }
}
