autoload -U add-zsh-hook

# Reset up count on new command
ATUIN_UP_COUNT=0
_atuin_reset_up_count() {
	ATUIN_UP_COUNT=0
}
add-zsh-hook precmd _atuin_reset_up_count

# Define functions
_atuin_up_cond() { # first arg is mode (viins or vicmd)
	ATUIN_UP_COUNT=$(( ATUIN_UP_COUNT + 1 ))
	if [ $ATUIN_UP_COUNT -gt 3 ]
	then
		zle "atuin-up-search$1"
	else
		zle up-line-or-beginning-search
	fi
}
_atuin_up_cond_vicmd() {
    _atuin_up_cond "-vicmd"
}
_atuin_up_cond_viins() {
    _atuin_up_cond "-viins"
}

# Add functions to zle
zle -N atuin-up-cond _atuin_up_cond
zle -N atuin-up-cond-vicmd _atuin_up_cond_vicmd
zle -N atuin-up-cond-viins _atuin_up_cond_viins

# Bind keys
bindkey -M emacs '^[[A' atuin-up-cond
bindkey -M vicmd '^[[A' atuin-up-cond-vicmd
bindkey -M viins '^[[A' atuin-up-cond-viins

bindkey -M emacs '^[OA' atuin-up-cond
bindkey -M vicmd '^[OA' atuin-up-cond-vicmd
bindkey -M viins '^[OA' atuin-up-cond-viins

bindkey -M vicmd 'k' atuin-up-cond-vicmd