From 60d779640bb97dc6730fc2298a46c660c0258936 Mon Sep 17 00:00:00 2001
From: Jonas Widen <jonas.widen@widens.eu>
Date: Mon, 17 Mar 2025 06:43:28 +0100
Subject: [PATCH] Added completion

---
 lua/gemini/completion.lua | 79 ++++++++++++++++++++++++---------------
 1 file changed, 48 insertions(+), 31 deletions(-)

diff --git a/lua/gemini/completion.lua b/lua/gemini/completion.lua
index 0374c29..817329d 100644
--- a/lua/gemini/completion.lua
+++ b/lua/gemini/completion.lua
@@ -34,11 +34,19 @@ local function show_suggestion(suggestion, start_col)
     local line = vim.api.nvim_win_get_cursor(0)[1] - 1
     local line_text = vim.api.nvim_buf_get_lines(0, line, line + 1, true)[1]
     
-    -- Validate column position
-    if start_col > #line_text then
-        return
+    -- Get text before cursor on current line
+    local prefix = string.sub(line_text, 1, start_col)
+    
+    -- Only show suggestion if it continues the current text
+    if not suggestion:find("^" .. vim.pesc(prefix), 1, true) then
+        suggestion = prefix .. suggestion
     end
     
+    -- Remove the prefix from the suggestion to avoid duplication
+    suggestion = suggestion:sub(#prefix + 1)
+    
+    if suggestion == "" then return end
+    
     current_suggestion.text = suggestion
     current_suggestion.start_col = start_col
     
@@ -50,29 +58,31 @@ local function show_suggestion(suggestion, start_col)
     })
 end
 
--- Function to accept current suggestion
 function M.accept_suggestion()
     if not (current_suggestion.text and current_suggestion.start_col) then
         return false
     end
 
     local line = vim.api.nvim_win_get_cursor(0)[1] - 1
-    local line_text = vim.api.nvim_buf_get_lines(0, line, line + 1, true)[1]
+    local col = vim.api.nvim_win_get_cursor(0)[2]
     
-    -- Validate column position
-    if current_suggestion.start_col > #line_text then
-        clear_suggestion()
+    -- Only accept if we're still at or before the suggestion start
+    if col < current_suggestion.start_col then
         return false
     end
     
     vim.api.nvim_buf_set_text(
         0,
         line,
-        current_suggestion.start_col,
+        col,
         line,
-        current_suggestion.start_col,
+        col,
         {current_suggestion.text}
     )
+    
+    -- Move cursor to end of inserted text
+    vim.api.nvim_win_set_cursor(0, {line + 1, col + #current_suggestion.text})
+    
     clear_suggestion()
     return true
 end
@@ -99,24 +109,35 @@ function M.trigger_completion()
             return
         end
         
-        -- Split into before and after cursor
-        local before_lines = {}
+        -- Get text before cursor on current line
+        local prefix = string.sub(current_line, 1, col)
         
-        -- Copy lines before current line
-        for i = 1, line - 1 do
-            table.insert(before_lines, lines[i])
+        -- Don't trigger on certain conditions
+        if prefix:match("^%s*$") or -- empty or whitespace
+           prefix:match("[%s%-]$") or -- ends with space or dash
+           #prefix < 3 then -- too short
+            clear_suggestion()
+            return
         end
         
-        -- Add current line (complete line for context)
-        table.insert(before_lines, current_line)
+        -- Construct context from buffer
+        local context = {}
+        
+        -- Add previous lines for context
+        for i = 1, line - 1 do
+            table.insert(context, lines[i])
+        end
+        
+        -- Add current line
+        table.insert(context, current_line)
         
         -- Combine all lines
-        local prefix = table.concat(before_lines, "\n")
+        local full_context = table.concat(context, "\n")
         
         -- Construct prompt for Gemini
         local prompt = string.format(
-            "Complete this code. Only provide the completion, no explanation:\n%s",
-            prefix
+            "Complete this code. Provide only the completion that would naturally follow, no explanation:\n%s",
+            full_context
         )
         
         -- Get completion from Gemini
@@ -124,17 +145,13 @@ function M.trigger_completion()
             if error then return end
             
             if type(response) == "string" then
-                -- Get first non-empty line as suggestion
-                local suggestion = response:match("^%s*(.-)%s*$")
-                if suggestion and #suggestion > 0 then
-                    vim.schedule(function()
-                        -- Check if cursor position is still the same
-                        local new_cursor = vim.api.nvim_win_get_cursor(0)
-                        if new_cursor[1] == line and new_cursor[2] == col then
-                            show_suggestion(suggestion, col)
-                        end
-                    end)
-                end
+                vim.schedule(function()
+                    -- Check if cursor position is still valid
+                    local new_cursor = vim.api.nvim_win_get_cursor(0)
+                    if new_cursor[1] == line and new_cursor[2] >= col then
+                        show_suggestion(response, col)
+                    end
+                end)
             end
         end)
     end)