From a3d15aa13f5b393b7b0f32482a33deaf01da271f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Sampaio?= <jpmelos@gmail.com>
Date: Sun, 30 Mar 2025 14:59:51 -0400
Subject: [PATCH] Detect visual blocks too

---
 autoload/augment.vim      |  2 +-
 autoload/augment/chat.vim | 22 ++++++++++++++--------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/autoload/augment.vim b/autoload/augment.vim
index b244711..091def0 100644
--- a/autoload/augment.vim
+++ b/autoload/augment.vim
@@ -173,7 +173,7 @@ function! s:CommandChat(range, args) abort
 
     " If range arguments were provided (when using :Augment chat) or in visual
     " mode, get the selected text
-    if a:range == 2 || mode() ==# 'v' || mode() ==# 'V'
+    if a:range == 2 || mode() ==# 'v' || mode() ==# 'V' || mode() ==# nr2char(22)
         let selected_text = augment#chat#GetSelectedText()
     else
         let selected_text = ''
diff --git a/autoload/augment/chat.vim b/autoload/augment/chat.vim
index a3a9acd..981b3d5 100644
--- a/autoload/augment/chat.vim
+++ b/autoload/augment/chat.vim
@@ -160,16 +160,26 @@ function! s:GetBufSelection(line_start, col_start, line_end, col_end) abort
         return getline(a:line_start)[a:col_start - 1:a:col_end - 1]
     endif
 
+    " mode() == 'n' means we're working with Ex mode using marks
+    if mode() ==# 'n' || mode() ==# 'v' || mode() ==# 'V'
+        let lines = []
+        call add(lines, getline(a:line_start)[a:col_start - 1:])
+        call extend(lines, getline(a:line_start + 1, a:line_end - 1))
+        call add(lines, getline(a:line_end)[0:a:col_end - 1])
+        return join(lines, "\n")
+    endif
+
+    " We're in visual block mode
     let lines = []
-    call add(lines, getline(a:line_start)[a:col_start - 1:])
-    call extend(lines, getline(a:line_start + 1, a:line_end - 1))
-    call add(lines, getline(a:line_end)[0:a:col_end - 1])
+    for lnum in range(a:line_start, a:line_end)
+        call add(lines, getline(lnum)[a:col_start - 1:a:col_end - 1])
+    endfor
     return join(lines, "\n")
 endfunction
 
 function! augment#chat#GetSelectedText() abort
     " If in visual mode use the current selection
-    if mode() ==# 'v' || mode() ==# 'V'
+    if mode() ==# 'v' || mode() ==# 'V' || mode() ==# nr2char(22)
         let [line_one, col_one] = getpos('.')[1:2]
         let [line_two, col_two] = getpos('v')[1:2]
 
@@ -197,10 +207,6 @@ function! augment#chat#GetSelectedText() abort
             endif
         endif
 
-        " . and v return column positions one lower than '< and '>
-        let col_start += 1
-        let col_end += 1
-
         " In visual line mode, the columns will be incorrect
         if mode() ==# 'V'
             let col_start = 1