bindkey -v
export KEYTIMEOUT=1

autoload -U add-zle-hook-widget

line_pre_redraw () {
    local previous_vi_keymap="${VI_KEYMAP}"

    case "${KEYMAP}" in
        vicmd)
            case "${REGION_ACTIVE}" in
                1)
                    VI_KEYMAP="VISUAL"
                    ;;
                2)
                    VI_KEYMAP="V-LINE"
                    ;;
                *)
                    VI_KEYMAP="NORMAL"
                    ;;
            esac
            ;;
        viins|main)
            if [[ "${ZLE_STATE}" == *overwrite* ]]; then
                VI_KEYMAP="REPLACE"
            else
                VI_KEYMAP="INSERT"
            fi
            ;;
    esac

    if [[ "${VI_KEYMAP}" != "${previous_vi_keymap}" ]]; then
        update_right_prompt
    fi
}

VI_KEYMAP="INSERT"

add-zle-hook-widget zle-line-pre-redraw line_pre_redraw