diff --git a/autosuggestions.zsh b/autosuggestions.zsh
index d90256c..dd69a4e 100644
--- a/autosuggestions.zsh
+++ b/autosuggestions.zsh
@@ -1,13 +1,4 @@
-# Fish-like autosuggestions for zsh. This is implemented on top of zsh's
-# builtin prediction feature by treating whatever is after the cursor
-# as 'unmaterialized'. The unmaterialized part is highlighted and ignored
-# when the user accepts the line. To materialize autosuggestions 'TAB' must
-# be pressed.
-#
-# Since predict-on doesn't work well on the middle of the line, many actions
-# that move the cursor to the left will pause autosuggestions, so it should be
-# safe enough to leave autosuggest enabled by default by adding the followingto
-# zshrc:
+# Fish-like autosuggestions for zsh based on the code for 'predict-on'
 #
 # ```zsh
 # zle-line-init() {
@@ -15,9 +6,24 @@
 # }
 # zle -N zle-line-init
 # ```
+# zmodload zsh/zpty
+#
 
+ZLE_AUTOSUGGEST_PAUSE_WIDGETS=(
+vi-cmd-mode vi-backward-char backward-char backward-word beginning-of-line
+history-search-forward history-search-backward up-line-or-history
+down-line-or-history
+)
+
+ZLE_AUTOSUGGEST_COMPLETION_WIDGETS=(
+complete-word expand-or-complete expand-or-complete-prefix list-choices
+menu-complete reverse-menu-complete menu-expand-or-complete menu-select
+accept-and-menu-complete
+)
 
 pause-autosuggestions() {
+	[[ -n $ZLE_AUTOSUGGESTING_PAUSED ]] && return
+	local widget
 	# When autosuggestions are disabled, kill the unmaterialized part
 	RBUFFER=''
 	unset ZLE_AUTOSUGGESTING
@@ -25,50 +31,40 @@ pause-autosuggestions() {
 	zle -A self-insert paused-autosuggest-self-insert
 	zle -A .magic-space magic-space
 	zle -A .backward-delete-char backward-delete-char
-	zle -A .delete-char-or-list delete-char-or-list
 	zle -A .accept-line accept-line
-	zle -A .vi-cmd-mode vi-cmd-mode
-	zle -A .vi-backward-char vi-backward-char
-	zle -A .backward-char backward-char
-	zle -A .backward-word backward-word
-	zle -A .beginning-of-line beginning-of-line
-	zle -A .history-search-forward history-search-forward
-	zle -A .history-search-backward history-search-backward
-	zle -A .up-line-or-history up-line-or-history
-	zle -A .down-line-or-history down-line-or-history
-	zle -A .complete-word complete-word
-	zle -A .expand-or-complete expand-or-complete
+	for widget in $ZLE_AUTOSUGGEST_PAUSE_WIDGETS; do
+		eval "zle -A autosuggest-${widget}-orig ${widget}"
+	done
+	for widget in $ZLE_AUTOSUGGEST_COMPLETION_WIDGETS; do
+		eval "zle -A autosuggest-${widget}-orig $widget"
+	done
 	highlight-suggested-text
 }
 
 enable-autosuggestions() {
+	[[ -n $ZLE_AUTOSUGGESTING ]] && return
+	local widget
 	unset ZLE_AUTOSUGGESTING_PAUSED
 	ZLE_AUTOSUGGESTING=1
 	# Replace prediction widgets by versions that will also highlight RBUFFER
 	zle -N self-insert autosuggest-self-insert
-	zle -N self-insert autosuggest-self-insert
+	zle -N magic-space autosuggest-self-insert
 	zle -N backward-delete-char autosuggest-delete
-	zle -N delete-char-or-list autosuggest-delete
-	# Replace some default widgets that should disable autosuggestion
+	zle -N accept-line autosuggest-accept-line
+	# Hook into some default widgets that should pause autosuggestion
 	# automatically 
-	zle -N accept-line execute-widget-and-pause
-	zle -N vi-cmd-mode execute-widget-and-pause
-	zle -N vi-backward-char execute-widget-and-pause
-	zle -N backward-char execute-widget-and-pause
-	zle -N backward-word execute-widget-and-pause
-	zle -N beginning-of-line execute-widget-and-pause
-	zle -N history-search-forward execute-widget-and-pause
-	zle -N history-search-backward execute-widget-and-pause
-	zle -N up-line-or-history execute-widget-and-pause
-	zle -N down-line-or-history execute-widget-and-pause
-	zle -N complete-word autosuggest-expand-or-complete
-	zle -N expand-or-complete autosuggest-expand-or-complete
+	for widget in $ZLE_AUTOSUGGEST_PAUSE_WIDGETS; do
+		eval "zle -A $widget autosuggest-${widget}-orig; \
+			zle -A autosuggest-pause $widget"
+	done
+	# Hook into completion widgets to handle suggestions after completions
+	for widget in $ZLE_AUTOSUGGEST_COMPLETION_WIDGETS; do
+		eval "zle -A $widget autosuggest-${widget}-orig; \
+			zle -A autosuggest-tab $widget"
+	done
 	if [[ $BUFFER != '' ]]; then
-		local cursor=$CURSOR
-		zle .expand-or-complete
-		CURSOR=$cursor
+		show-suggestion
 	fi
-	highlight-suggested-text
 }
 
 disable-autosuggestions() {
@@ -88,57 +84,17 @@ toggle-autosuggestions() {
 	fi
 }
 
-# TODO Most of the widgets here only override default widgets to disable
-# autosuggestion, find a way to do it in a loop for the sake of maintainability
-
-# When autosuggesting, ignore RBUFFER which corresponds to the 'unmaterialized'
-# section when the user accepts the line
-autosuggest-accept-line() {
-	RBUFFER=''
-	region_highlight=()
-	zle .accept-line
-}
-
-execute-widget-and-pause() {
-	pause-autosuggestions
-	zle .$WIDGET "$@"
-}
-
 highlight-suggested-text() {
 	if [[ -n $ZLE_AUTOSUGGESTING ]]; then
 		local color='fg=8'
 		[[ -n $AUTOSUGGESTION_HIGHLIGHT_COLOR ]] &&\
-		 	color=$AUTOSUGGESTION_HIGHLIGHT_COLOR
+			color=$AUTOSUGGESTION_HIGHLIGHT_COLOR
 		region_highlight=("$(( $CURSOR + 1 )) $(( $CURSOR + $#RBUFFER )) $color")
 	else
 		region_highlight=()
 	fi
 }
 
-paused-autosuggest-self-insert() {
-	if [[ $RBUFFER == '' ]]; then
-		# Resume autosuggestions when inserting at the end of the line
-		enable-autosuggestions
-		autosuggest-self-insert
-	else
-		zle .self-insert
-	fi
-}
-
-show-suggestion() {
-	local complete_word=$1
-	if ! zle .history-beginning-search-backward; then
-		RBUFFER=''
-		if [[ $LBUFFER[-1] != ' ' ]]; then
-			integer curs=$CURSOR
-			unsetopt automenu recexact
-			zle complete-word-orig
-			CURSOR=$curs
-		fi
-	fi
-	highlight-suggested-text
-}
-
 autosuggest-self-insert() {
 	setopt localoptions noshwordsplit noksharrays
 	if [[ ${RBUFFER[1]} == ${KEYS[-1]} ]]; then
@@ -151,15 +107,66 @@ autosuggest-self-insert() {
 	fi
 }
 
+# Taken from predict-on
 autosuggest-delete() {
-	zle .$WIDGET
-	show-suggestion
+	if (( $#LBUFFER > 1 )); then
+		setopt localoptions noshwordsplit noksharrays
+		# When editing a multiline buffer, it's unlikely prediction is wanted;
+		# or if the last widget was e.g. a motion, then probably the intent is
+		# to actually edit the line, not change the search prefix.
+		if [[ $LASTWIDGET != (self-insert|magic-space|backward-delete-char) ]]; then
+			LBUFFER="$LBUFFER[1,-2]"
+		else
+			((--CURSOR))
+			zle .history-beginning-search-forward || RBUFFER=""
+			return 0
+		fi
+	else
+		zle .kill-whole-line
+	fi
+	highlight-suggested-text
 }
 
-autosuggest-expand-or-complete() {
+# When autosuggesting, ignore RBUFFER which corresponds to the 'unmaterialized'
+# section when the user accepts the line
+autosuggest-accept-line() {
 	RBUFFER=''
-	zle .$WIDGET "$@"
-	show-suggestion
+	region_highlight=()
+	zle .accept-line
+}
+
+paused-autosuggest-self-insert() {
+	if [[ $RBUFFER == '' ]]; then
+		# Resume autosuggestions when inserting at the end of the line
+		enable-autosuggestions
+		autosuggest-self-insert
+	else
+		zle .self-insert
+	fi
+}
+
+autosuggest-first-completion() {
+	zle .complete-word || return 1
+}
+
+show-suggestion() {
+	[[ $LBUFFER == '' ]] && return
+	# TODO need a way to reset HISTNO so .history-beginning-search-backward
+	# will always retrieve the last matching history entry
+	# unset HISTNO
+	zle .history-beginning-search-backward || autosuggest-first-completion\
+		|| RBUFFER=''
+	highlight-suggested-text
+}
+
+autosuggest-pause() {
+	pause-autosuggestions
+	zle autosuggest-${WIDGET}-orig "$@"
+}
+
+autosuggest-tab() {
+	RBUFFER=''
+	zle autosuggest-${WIDGET}-orig "$@"
 }
 
 accept-suggested-small-word() {
@@ -173,7 +180,7 @@ accept-suggested-word() {
 }
 
 zle -N toggle-autosuggestions
-zle -N enable-autosuggestions
-zle -N disable-autosuggestions
 zle -N accept-suggested-small-word
 zle -N accept-suggested-word
+zle -N autosuggest-pause
+zle -N autosuggest-tab