autoload -Uz add-zsh-hook

zmodload zsh/zpty

ZSH_AUTOSUGGEST_HIGHLIGHT_STYLE='fg=8'
ZSH_AUTOSUGGEST_ORIGINAL_WIDGET_PREFIX=autosuggest-orig-
ZSH_AUTOSUGGEST_STRATEGY=default

ZSH_AUTOSUGGEST_CLEAR_WIDGETS=(
	history-search-forward
	history-search-backward
	history-beginning-search-forward
	history-beginning-search-backward
	history-substring-search-up
	history-substring-search-down
	up-line-or-history
	down-line-or-history
	accept-line
)

ZSH_AUTOSUGGEST_ACCEPT_WIDGETS=(
	forward-char
	end-of-line
	vi-forward-char
	vi-end-of-line
	vi-add-eol
)

ZSH_AUTOSUGGEST_EXECUTE_WIDGETS=(
)

ZSH_AUTOSUGGEST_PARTIAL_ACCEPT_WIDGETS=(
	forward-word
	vi-forward-word
	vi-forward-word-end
	vi-forward-blank-word
	vi-forward-blank-word-end
)

ZSH_AUTOSUGGEST_IGNORE_WIDGETS=(
	orig-\*
	beep
	run-help
	set-local-history
	which-command
	yank
)

ZSH_AUTOSUGGEST_BUFFER_MAX_SIZE=
ZSH_AUTOSUGGEST_ASYNC_PTY_NAME=zsh_autosuggest_pty

_zsh_autosuggest_escape_command() {
	setopt localoptions EXTENDED_GLOB
	echo -E "${1//(#m)[\"\'\\()\[\]|*?~]/\\$MATCH}"
}

_zsh_autosuggest_feature_detect_zpty_returns_fd() {
	typeset -g _ZSH_AUTOSUGGEST_ZPTY_RETURNS_FD
	typeset -h REPLY

	zpty zsh_autosuggest_feature_detect '{ zshexit() { kill -KILL $$; sleep 1 } }'

	if (( REPLY )); then
		_ZSH_AUTOSUGGEST_ZPTY_RETURNS_FD=1
	else
		_ZSH_AUTOSUGGEST_ZPTY_RETURNS_FD=0
	fi

	zpty -d zsh_autosuggest_feature_detect
}

_zsh_autosuggest_incr_bind_count() {
	if ((${+_ZSH_AUTOSUGGEST_BIND_COUNTS[$1]})); then
		((_ZSH_AUTOSUGGEST_BIND_COUNTS[$1]++))
	else
		_ZSH_AUTOSUGGEST_BIND_COUNTS[$1]=1
	fi

	bind_count=$_ZSH_AUTOSUGGEST_BIND_COUNTS[$1]
}

_zsh_autosuggest_get_bind_count() {
	if ((${+_ZSH_AUTOSUGGEST_BIND_COUNTS[$1]})); then
		bind_count=$_ZSH_AUTOSUGGEST_BIND_COUNTS[$1]
	else
		bind_count=0
	fi
}

_zsh_autosuggest_bind_widget() {
	typeset -gA _ZSH_AUTOSUGGEST_BIND_COUNTS

	local widget=$1
	local autosuggest_action=$2
	local prefix=$ZSH_AUTOSUGGEST_ORIGINAL_WIDGET_PREFIX

	local -i bind_count

	case $widgets[$widget] in
		# Already bound
		user:_zsh_autosuggest_(bound|orig)_*);;

		user:*)
			_zsh_autosuggest_incr_bind_count $widget
			zle -N $prefix${bind_count}-$widget ${widgets[$widget]#*:}
			;;

		builtin)
			_zsh_autosuggest_incr_bind_count $widget
			eval "_zsh_autosuggest_orig_${(q)widget}() { zle .${(q)widget} }"
			zle -N $prefix${bind_count}-$widget _zsh_autosuggest_orig_$widget
			;;

		completion:*)
			_zsh_autosuggest_incr_bind_count $widget
			eval "zle -C $prefix${bind_count}-${(q)widget} ${${(s.:.)widgets[$widget]}[2,3]}"
			;;
	esac

	_zsh_autosuggest_get_bind_count $widget

	eval "_zsh_autosuggest_bound_${bind_count}_${(q)widget}() {
		_zsh_autosuggest_widget_$autosuggest_action $prefix$bind_count-${(q)widget} \$@
	}"

	zle -N $widget _zsh_autosuggest_bound_${bind_count}_$widget
}

_zsh_autosuggest_bind_widgets() {
	local widget
	local ignore_widgets

	ignore_widgets=(
		.\*
		_\*
		zle-\*
		autosuggest-\*
		$ZSH_AUTOSUGGEST_ORIGINAL_WIDGET_PREFIX\*
		$ZSH_AUTOSUGGEST_IGNORE_WIDGETS
	)

	for widget in ${${(f)"$(builtin zle -la)"}:#${(j:|:)~ignore_widgets}}; do
		if [ ${ZSH_AUTOSUGGEST_CLEAR_WIDGETS[(r)$widget]} ]; then
			_zsh_autosuggest_bind_widget $widget clear
		elif [ ${ZSH_AUTOSUGGEST_ACCEPT_WIDGETS[(r)$widget]} ]; then
			_zsh_autosuggest_bind_widget $widget accept
		elif [ ${ZSH_AUTOSUGGEST_EXECUTE_WIDGETS[(r)$widget]} ]; then
			_zsh_autosuggest_bind_widget $widget execute
		elif [ ${ZSH_AUTOSUGGEST_PARTIAL_ACCEPT_WIDGETS[(r)$widget]} ]; then
			_zsh_autosuggest_bind_widget $widget partial_accept
		else
			_zsh_autosuggest_bind_widget $widget modify
		fi
	done
}

_zsh_autosuggest_invoke_original_widget() {
	[ $# -gt 0 ] || return

	local original_widget_name="$1"

	shift

	if [ $widgets[$original_widget_name] ]; then
		zle $original_widget_name -- $@
	fi
}

_zsh_autosuggest_highlight_reset() {
	typeset -g _ZSH_AUTOSUGGEST_LAST_HIGHLIGHT

	if [ -n "$_ZSH_AUTOSUGGEST_LAST_HIGHLIGHT" ]; then
		region_highlight=("${(@)region_highlight:#$_ZSH_AUTOSUGGEST_LAST_HIGHLIGHT}")
		unset _ZSH_AUTOSUGGEST_LAST_HIGHLIGHT
	fi
}

_zsh_autosuggest_highlight_apply() {
	typeset -g _ZSH_AUTOSUGGEST_LAST_HIGHLIGHT

	if [ $#POSTDISPLAY -gt 0 ]; then
		_ZSH_AUTOSUGGEST_LAST_HIGHLIGHT="$#BUFFER $(($#BUFFER + $#POSTDISPLAY)) $ZSH_AUTOSUGGEST_HIGHLIGHT_STYLE"
		region_highlight+=("$_ZSH_AUTOSUGGEST_LAST_HIGHLIGHT")
	else
		unset _ZSH_AUTOSUGGEST_LAST_HIGHLIGHT
	fi
}

_zsh_autosuggest_disable() {
	typeset -g _ZSH_AUTOSUGGEST_DISABLED
	_zsh_autosuggest_clear
}

_zsh_autosuggest_enable() {
	unset _ZSH_AUTOSUGGEST_DISABLED

	if [ $#BUFFER -gt 0 ]; then
		_zsh_autosuggest_fetch
	fi
}

_zsh_autosuggest_toggle() {
	if [ -n "${_ZSH_AUTOSUGGEST_DISABLED+x}" ]; then
		_zsh_autosuggest_enable
	else
		_zsh_autosuggest_disable
	fi
}

_zsh_autosuggest_clear() {
	unset POSTDISPLAY

	_zsh_autosuggest_invoke_original_widget $@
}

_zsh_autosuggest_modify() {
	local -i retval

	local -i KEYS_QUEUED_COUNT

	local orig_buffer="$BUFFER"
	local orig_postdisplay="$POSTDISPLAY"

	unset POSTDISPLAY

	_zsh_autosuggest_invoke_original_widget $@
	retval=$?

	if [[ $PENDING > 0 ]] || [[ $KEYS_QUEUED_COUNT > 0 ]]; then
		return $retval
	fi

	if [ $#BUFFER -gt $#orig_buffer ]; then
		local added=${BUFFER#$orig_buffer}

		if [ "$added" = "${orig_postdisplay:0:$#added}" ]; then
			POSTDISPLAY="${orig_postdisplay:$#added}"
			return $retval
		fi
	fi

	if [ "$BUFFER" = "$orig_buffer" ]; then
		POSTDISPLAY="$orig_postdisplay"
		return $retval
	fi

	if [ -n "${_ZSH_AUTOSUGGEST_DISABLED+x}" ]; then
		return $?
	fi

	if [ $#BUFFER -gt 0 ]; then
		if [ -z "$ZSH_AUTOSUGGEST_BUFFER_MAX_SIZE" -o $#BUFFER -le "$ZSH_AUTOSUGGEST_BUFFER_MAX_SIZE" ]; then
			_zsh_autosuggest_fetch
		fi
	fi

	return $retval
}

_zsh_autosuggest_fetch() {
	if zpty -t "$ZSH_AUTOSUGGEST_ASYNC_PTY_NAME" &>/dev/null; then
		_zsh_autosuggest_async_request "$BUFFER"
	else
		local suggestion
		_zsh_autosuggest_strategy_$ZSH_AUTOSUGGEST_STRATEGY "$BUFFER"
		_zsh_autosuggest_suggest "$suggestion"
	fi
}

_zsh_autosuggest_suggest() {
	local suggestion="$1"

	if [ -n "$suggestion" ] && [ $#BUFFER -gt 0 ]; then
		POSTDISPLAY="${suggestion#$BUFFER}"
	else
		unset POSTDISPLAY
	fi
}

_zsh_autosuggest_accept() {
	local -i max_cursor_pos=$#BUFFER

	if [ "$KEYMAP" = "vicmd" ]; then
		max_cursor_pos=$((max_cursor_pos - 1))
	fi

	if [ $CURSOR -eq $max_cursor_pos ]; then
		BUFFER="$BUFFER$POSTDISPLAY"
		unset POSTDISPLAY
		CURSOR=${#BUFFER}
	fi

	_zsh_autosuggest_invoke_original_widget $@
}

_zsh_autosuggest_execute() {
	BUFFER="$BUFFER$POSTDISPLAY"
	unset POSTDISPLAY
	_zsh_autosuggest_invoke_original_widget "accept-line"
}

_zsh_autosuggest_partial_accept() {
	local -i retval
	local original_buffer="$BUFFER"
	BUFFER="$BUFFER$POSTDISPLAY"
	_zsh_autosuggest_invoke_original_widget $@
	retval=$?

	if [ $CURSOR -gt $#original_buffer ]; then
		POSTDISPLAY="$RBUFFER"
		BUFFER="$LBUFFER"
	else
		BUFFER="$original_buffer"
	fi

	return $retval
}

for action in clear modify fetch suggest accept partial_accept execute enable disable toggle; do
	eval "_zsh_autosuggest_widget_$action() {
		local -i retval

		_zsh_autosuggest_highlight_reset

		_zsh_autosuggest_$action \$@
		retval=\$?

		_zsh_autosuggest_highlight_apply

		zle -R

		return \$retval
	}"
done

zle -N autosuggest-fetch _zsh_autosuggest_widget_fetch
zle -N autosuggest-suggest _zsh_autosuggest_widget_suggest
zle -N autosuggest-accept _zsh_autosuggest_widget_accept
zle -N autosuggest-clear _zsh_autosuggest_widget_clear
zle -N autosuggest-execute _zsh_autosuggest_widget_execute
zle -N autosuggest-enable _zsh_autosuggest_widget_enable
zle -N autosuggest-disable _zsh_autosuggest_widget_disable
zle -N autosuggest-toggle _zsh_autosuggest_widget_toggle

_zsh_autosuggest_strategy_default() {
	emulate -L zsh
	setopt EXTENDED_GLOB
	local prefix="${1//(#m)[\\*?[\]<>()|^~#]/\\$MATCH}"
	suggestion="${history[(r)$prefix*]}"
}

_zsh_autosuggest_strategy_match_prev_cmd() {
	local prefix="${1//(#m)[\\()\[\]|*?~]/\\$MATCH}"
	local history_match_keys
	history_match_keys=(${(k)history[(R)$prefix*]})
	local histkey="${history_match_keys[1]}"
	local prev_cmd="$(_zsh_autosuggest_escape_command "${history[$((HISTCMD-1))]}")"

	for key in "${(@)history_match_keys[1,200]}"; do
		[[ $key -gt 1 ]] || break

		if [[ "${history[$((key - 1))]}" == "$prev_cmd" ]]; then
			histkey="$key"
			break
		fi
	done

	suggestion="$history[$histkey]"
}

_zsh_autosuggest_async_server() {
	emulate -R zsh

	zshexit() {
		kill -KILL $$
		sleep 1
	}

	stty -onlcr

	exec 2>/dev/null

	local strategy=$1
	local last_pid

	while IFS='' read -r -d $'\0' query; do
		kill -KILL $last_pid &>/dev/null

		(
			local suggestion
			_zsh_autosuggest_strategy_$ZSH_AUTOSUGGEST_STRATEGY "$query"
			echo -n -E "$suggestion"$'\0'
		) &

		last_pid=$!
	done
}

_zsh_autosuggest_async_request() {
	zpty -w -n $ZSH_AUTOSUGGEST_ASYNC_PTY_NAME "${1}"$'\0'
}

_zsh_autosuggest_async_response() {
	setopt LOCAL_OPTIONS EXTENDED_GLOB

	local suggestion

	zpty -rt $ZSH_AUTOSUGGEST_ASYNC_PTY_NAME suggestion '*'$'\0' 2>/dev/null
	zle autosuggest-suggest -- "${suggestion%%$'\0'##}"
}

_zsh_autosuggest_async_pty_create() {
	typeset -h REPLY

	if [ $_ZSH_AUTOSUGGEST_ZPTY_RETURNS_FD -eq 0 ]; then
		integer -l zptyfd
		exec {zptyfd}>&1
		exec {zptyfd}>&-
	fi

	zpty -b $ZSH_AUTOSUGGEST_ASYNC_PTY_NAME "_zsh_autosuggest_async_server _zsh_autosuggest_strategy_$ZSH_AUTOSUGGEST_STRATEGY"

	if (( REPLY )); then
		_ZSH_AUTOSUGGEST_PTY_FD=$REPLY
	else
		_ZSH_AUTOSUGGEST_PTY_FD=$zptyfd
	fi

	zle -F $_ZSH_AUTOSUGGEST_PTY_FD _zsh_autosuggest_async_response
}

_zsh_autosuggest_async_pty_destroy() {
	if zpty -t $ZSH_AUTOSUGGEST_ASYNC_PTY_NAME &>/dev/null; then
		zle -F $_ZSH_AUTOSUGGEST_PTY_FD &>/dev/null

		zpty -d $ZSH_AUTOSUGGEST_ASYNC_PTY_NAME &>/dev/null
	fi
}

_zsh_autosuggest_async_pty_recreate() {
	_zsh_autosuggest_async_pty_destroy
	_zsh_autosuggest_async_pty_create
}

_zsh_autosuggest_async_start() {
	typeset -g _ZSH_AUTOSUGGEST_PTY_FD

	_zsh_autosuggest_feature_detect_zpty_returns_fd
	_zsh_autosuggest_async_pty_recreate

	add-zsh-hook precmd _zsh_autosuggest_async_pty_recreate
}

_zsh_autosuggest_start() {
	add-zsh-hook -d precmd _zsh_autosuggest_start

	_zsh_autosuggest_bind_widgets

	add-zsh-hook precmd _zsh_autosuggest_bind_widgets

	if [ -n "${ZSH_AUTOSUGGEST_USE_ASYNC+x}" ]; then
		_zsh_autosuggest_async_start
	fi
}

add-zsh-hook precmd _zsh_autosuggest_start