#!/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
typeset -g _autocomplete__overhead=0

${0}:precmd() {
  [[ -v ZSH_AUTOSUGGEST_IGNORE_WIDGETS ]] &&
      ZSH_AUTOSUGGEST_IGNORE_WIDGETS+=(
          history-incremental-search-backward
          recent-paths
          .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: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
}

# function -T
.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:*:fd-widget) ]] &&
      return 0

  {
    if (( REGION_ACTIVE )) ||
        [[ -v _autocomplete__isearch && $LASTWIDGET == *(incremental|isearch)* ]]; then
      builtin zle -Rc
      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

    .autocomplete:async:wait
  }

  return 0
}

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

# function -t
.autocomplete:async:wait() {
  local fd=

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

    (( seconds_delay = max( 0, seconds_delay - _autocomplete__overhead ) ))

    # 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

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

  return 0
}

# function -T
.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}<&-

    if [[ -n $_autocomplete__zle_flags ]]; then
      builtin zle -f $_autocomplete__zle_flags

      [[ $_autocomplete__zle_flags == yank* ]] &&
          return 0
    fi

    (( KEYS_QUEUED_COUNT || PENDING )) &&
        return

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

# function -T
.autocomplete:async:start() {
  local fd=
  sysopen -r -o cloexec -u fd <(
    local +h PS4=$_autocomplete__ps4
    .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
}

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

# function -T
.autocomplete:async:pty() {
  local +h PS4=$_autocomplete__ps4

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

  # WORKAROUND: #654 Inside the PTY, formatting can lead to incorrect list trimming.
  eval ${${${"$( zstyle -L '*' format )"}//zstyle /zstyle -d }// -e / }

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

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

# function -T
.autocomplete:async:pty:zle-widget() {
  local +h PS4=$_autocomplete__ps4

  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[@]

    local +h PS4=$_autocomplete__ps4
    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

# function -T
.autocomplete:async:pty:completion-widget() {
  local +h PS4=$_autocomplete__ps4

  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
}

# function -T
.autocomplete:async:complete:fd-widget() {
  setopt localoptions NO_banghist

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

      if [[ -n $_autocomplete__zle_flags ]]; then
        builtin zle -f $_autocomplete__zle_flags

        [[ $_autocomplete__zle_flags == yank* ]] &&
            return 0
      fi

      (( KEYS_QUEUED_COUNT || PENDING )) &&
          return

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

      local -a reply=()
      IFS=$'\0' read -rAu $fd
      shift -p reply
      (( SECONDS += reply[2] ))

    } always {
      exec {fd}<&-
    }

    [[ -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=$_autocomplete__ps4
    if ! builtin zle ._list_choices -w "$reply[1]" 2>>| $_autocomplete__log; then

      typeset -g _autocomplete__overhead=$SECONDS

      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

    return 0
  }
}

.autocomplete:async:sufficient-input() {
  local min_input=
  if ! builtin zstyle -s ":autocomplete:${curcontext}:" 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}:" 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

  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{0}%K{12}(MORE)%f%k'
    _lastcomp[list_lines]=$compstate[list_lines]
  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-* )
    setopt $_autocomplete__func_opts[@]
    autocomplete:_main_complete:new - history-lines _autocomplete__history_lines
  ;;
  recent-paths:* )
    setopt $_autocomplete__func_opts[@]
    autocomplete:_main_complete:new - recent-paths _autocomplete__recent_paths
  ;;
  * )
    {
      () {
        emulate -L zsh
        setopt $_autocomplete__func_opts[@]

        local curcontext=list-choices:::

        .autocomplete:async:shadow compadd

        autoload -Uz +X _describe
        .autocomplete:async:shadow _describe

        # functions -T compadd _describe _description
      } "$@"

      .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() {
  local -A _opts_=()
  local -a _xopts_=() _displ_=() _matches_=()
  local -P _displ_name_= _matches_name_=

  zparseopts -A _opts_ -E -- D: E: O: X:=_xopts_ x:=_xopts_

  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]

  # Need to remove prompt escape codes or compadd might crash.
  local -Pi _total_new_lines_="$(
      zparseopts -a _xopts_ -D -E -- X: x:
      builtin compadd "$@"
      print -nr -- $(( $compstate[list_lines] - _old_total_lines ))
  )"

  local -Pi _new_completion_lines_="$(
      zparseopts -a _xopts_ -D -E -- X: x:
      builtin compadd "$@"
      print -nr -- $(( $compstate[list_lines] - _old_total_lines ))
  )"

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

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

  # We're apparently trying to add a message that spans more than one line, but not any completions.
  if (( _new_heading_lines_ > 1 && _new_completion_lines_ < 1 )); then
    .autocomplete:async:compadd:disable
    return 1
  fi

  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_ ${(@r:COLUMNS-1:)${(PA)_displ_name_}[@]//$'\n'/\n}

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

  # 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_++ ))

    if (( _nmatches_to_remove_ < $#_matches_ )); then
      shift -p $_nmatches_to_remove_ _matches_ $_displ_name_
    else
      set -A _matches_
      set -A $_displ_name_
    fi
    .autocomplete:async:compadd:disable
  fi

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

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