From 4b7d5e9204eb37abd023639e5737fa671afe6c49 Mon Sep 17 00:00:00 2001
From: Jonas Widen <jonas.widen@widens.eu>
Date: Sun, 16 Mar 2025 18:23:27 +0100
Subject: [PATCH] Minor fix

---
 lua/gemini/init.lua | 36 ++++++++++++++++++++++++------------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/lua/gemini/init.lua b/lua/gemini/init.lua
index dabf264..49c3446 100644
--- a/lua/gemini/init.lua
+++ b/lua/gemini/init.lua
@@ -162,7 +162,7 @@ local function gemini_query(prompt, context)
     current_context = context
     
     -- Show initial message in chat window and ensure it's visible
-    local initial_content = "\nUser: " .. prompt .. "\n\nAssistant: Thinking..."
+    local initial_content = "User: " .. prompt .. "\n\nAssistant: Thinking..."
     update_chat_window(initial_content)
     
     -- Force Neovim to update the screen
@@ -170,22 +170,33 @@ local function gemini_query(prompt, context)
     
     local response = api.get_response(prompt, context)
     if response then
-        -- Replace "Thinking..." with the actual response
+        -- Make buffer modifiable
+        vim.api.nvim_buf_set_option(chat_bufnr, 'modifiable', true)
+        
+        -- Get all lines
         local lines = vim.api.nvim_buf_get_lines(chat_bufnr, 0, -1, false)
-        for i = #lines, 1, -1 do
-            if lines[i]:match("^Assistant: Thinking...") then
-                vim.api.nvim_buf_set_option(chat_bufnr, 'modifiable', true)
+        
+        -- Find and remove the "Thinking..." line
+        for i = 1, #lines do
+            if lines[i] == "Assistant: Thinking..." then
                 -- Remove the "Thinking..." line
-                vim.api.nvim_buf_set_lines(chat_bufnr, i, i + 1, false, {})
-                -- Split response into lines and insert "Assistant: " at the start of first line
+                table.remove(lines, i)
+                -- Insert response lines
                 local response_lines = vim.split(response, "\n")
                 response_lines[1] = "Assistant: " .. response_lines[1]
-                vim.api.nvim_buf_set_lines(chat_bufnr, i, i, false, response_lines)
-                vim.api.nvim_buf_set_option(chat_bufnr, 'modifiable', false)
+                for j = #response_lines, 1, -1 do
+                    table.insert(lines, i, response_lines[j])
+                end
                 break
             end
         end
         
+        -- Update buffer with modified lines
+        vim.api.nvim_buf_set_lines(chat_bufnr, 0, -1, false, lines)
+        
+        -- Make buffer unmodifiable
+        vim.api.nvim_buf_set_option(chat_bufnr, 'modifiable', false)
+        
         -- Move focus to chat window
         if chat_winnr and vim.api.nvim_win_is_valid(chat_winnr) then
             vim.api.nvim_set_current_win(chat_winnr)
@@ -196,10 +207,11 @@ local function gemini_query(prompt, context)
     else
         -- Replace "Thinking..." with error message
         local lines = vim.api.nvim_buf_get_lines(chat_bufnr, 0, -1, false)
-        for i = #lines, 1, -1 do
-            if lines[i]:match("^Assistant: Thinking...") then
+        for i = 1, #lines do
+            if lines[i] == "Assistant: Thinking..." then
+                lines[i] = "Assistant: Failed to get response from Gemini API"
                 vim.api.nvim_buf_set_option(chat_bufnr, 'modifiable', true)
-                vim.api.nvim_buf_set_lines(chat_bufnr, i, i + 1, false, {"Assistant: Failed to get response from Gemini API"})
+                vim.api.nvim_buf_set_lines(chat_bufnr, 0, -1, false, lines)
                 vim.api.nvim_buf_set_option(chat_bufnr, 'modifiable', false)
                 break
             end