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

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

diff --git a/lua/gemini/init.lua b/lua/gemini/init.lua
index 81bcd32..7e8928d 100644
--- a/lua/gemini/init.lua
+++ b/lua/gemini/init.lua
@@ -161,14 +161,25 @@ local function gemini_query(prompt, context)
     -- Store the context for subsequent queries
     current_context = context
     
-    -- Show initial message in chat window
+    -- Show initial message in chat window and ensure it's visible
     local initial_content = "\nUser: " .. prompt .. "\n\nAssistant: Thinking..."
     update_chat_window(initial_content)
     
+    -- Force Neovim to update the screen
+    vim.cmd('redraw')
+    
     local response = api.get_response(prompt, context)
     if response then
-        local formatted_content = "\nUser: " .. prompt .. "\n\nAssistant: " .. response
-        update_chat_window(formatted_content)
+        -- Replace "Thinking..." with the actual response
+        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)
+                vim.api.nvim_buf_set_lines(chat_bufnr, i, i + 1, false, {"Assistant: " .. response})
+                vim.api.nvim_buf_set_option(chat_bufnr, 'modifiable', false)
+                break
+            end
+        end
         
         -- Move focus to chat window
         if chat_winnr and vim.api.nvim_win_is_valid(chat_winnr) then
@@ -178,9 +189,16 @@ local function gemini_query(prompt, context)
         -- Clear the command line
         vim.cmd('echo ""')
     else
-        -- Update the chat window with error message
-        local error_content = "\nUser: " .. prompt .. "\n\nAssistant: Failed to get response from Gemini API"
-        update_chat_window(error_content)
+        -- 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
+                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_option(chat_bufnr, 'modifiable', false)
+                break
+            end
+        end
         vim.notify("Failed to get a response from Gemini API", vim.log.levels.ERROR)
     end
 end