From 514a1e08acb007a2baf04a2b549d72aa17e4b3cb Mon Sep 17 00:00:00 2001
From: Jonas Widen <jonas.widen@widens.eu>
Date: Sun, 16 Mar 2025 14:39:43 +0100
Subject: [PATCH] Move windows to the right and added color coding

---
 README.md           |  1 +
 lua/gemini/init.lua | 47 +++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index d1c7cc0..9ad7385 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,7 @@ A Neovim plugin for interacting with the Google AI API (formerly Gemini). Get AI
 * Neovim 0.8 or higher
 * A Google AI API key (obtainable from [https://makersuite.google.com/app/apikey](https://makersuite.google.com/app/apikey))
 * `curl` installed on your system
+* Treesitter installed with the markdown parser (automatically installed during setup)
 
 ## Installation
 
diff --git a/lua/gemini/init.lua b/lua/gemini/init.lua
index fba6647..6491ad4 100644
--- a/lua/gemini/init.lua
+++ b/lua/gemini/init.lua
@@ -8,6 +8,16 @@ local function get_current_buffer_content()
     return table.concat(lines, "\n")
 end
 
+local function setup_treesitter_highlight(bufnr)
+    -- Enable treesitter for the buffer
+    vim.api.nvim_buf_set_option(bufnr, 'syntax', '')
+    local success = pcall(vim.treesitter.start, bufnr, "markdown")
+    if not success then
+        -- Fallback to basic markdown syntax if treesitter fails
+        vim.api.nvim_buf_set_option(bufnr, 'syntax', 'markdown')
+    end
+end
+
 local function gemini_query(prompt, context)
     local response = api.get_response(prompt, context)
 
@@ -19,19 +29,33 @@ local function gemini_query(prompt, context)
         -- Set buffer options
         vim.api.nvim_buf_set_option(new_buf, 'modifiable', false)
         vim.api.nvim_buf_set_option(new_buf, 'buftype', 'nofile')
+        vim.api.nvim_buf_set_option(new_buf, 'filetype', 'markdown')
+        
+        -- Calculate dimensions
+        local width = math.floor(vim.o.columns / 3)
+        local height = vim.o.lines - 2  -- Account for status line and command line
         
         -- Create the window
         local new_win = vim.api.nvim_open_win(new_buf, true, {
             relative = "editor",
-            width = 80,
-            height = 20,
-            row = 5,
-            col = vim.o.columns / 2 - 40,
+            width = width,
+            height = height,
+            row = 0,
+            col = vim.o.columns - width,
             border = "rounded",
             title = "Google AI Response",
+            title_pos = "center",
             style = "minimal"
         })
 
+        -- Set window-local options
+        vim.api.nvim_win_set_option(new_win, 'wrap', true)
+        vim.api.nvim_win_set_option(new_win, 'linebreak', true)
+        vim.api.nvim_win_set_option(new_win, 'breakindent', true)
+
+        -- Setup treesitter highlighting
+        setup_treesitter_highlight(new_buf)
+
         -- Set window-local keymaps
         local close_keys = {'q', '<Esc>', '<CR>'}
         for _, key in ipairs(close_keys) do
@@ -50,6 +74,9 @@ local function gemini_query(prompt, context)
             end,
             once = true,
         })
+
+        -- Return focus to the previous window
+        vim.cmd('wincmd p')
     else
         vim.notify("Failed to get a response from Gemini API", vim.log.levels.ERROR)
     end
@@ -59,6 +86,18 @@ end
 M.gemini_query = gemini_query
 
 function M.setup()
+    -- Ensure markdown parser is installed
+    local parser_installed = pcall(vim.treesitter.language.require_language, "markdown")
+    if not parser_installed then
+        vim.notify("Installing markdown parser for treesitter...", vim.log.levels.INFO)
+        vim.fn.system({
+            "nvim",
+            "--headless",
+            "-c", "TSInstall markdown",
+            "-c", "q"
+        })
+    end
+
     -- Create the user command
     vim.api.nvim_create_user_command("Gemini", function(opts)
         local prompt = opts.args