From 9d72ba46c4848c341f279a01c16a3b021c7d0235 Mon Sep 17 00:00:00 2001
From: Dominik Polakovics <dominik.polakovics@cloonar.com>
Date: Sun, 9 Feb 2025 00:02:40 +0100
Subject: [PATCH] feat: add file appending to prompt

---
 lua/chatgpt_nvim/init.lua | 71 +++++++++++++++++++++++++++++++++------
 1 file changed, 60 insertions(+), 11 deletions(-)

diff --git a/lua/chatgpt_nvim/init.lua b/lua/chatgpt_nvim/init.lua
index 2a2c15a..d0d1009 100644
--- a/lua/chatgpt_nvim/init.lua
+++ b/lua/chatgpt_nvim/init.lua
@@ -54,15 +54,6 @@ local function read_file(path)
   return data
 end
 
-local function close_existing_buffer_by_name(pattern)
-  for _, b in ipairs(vim.api.nvim_list_bufs()) do
-    local name = vim.api.nvim_buf_get_name(b)
-    if name:match(pattern) then
-      vim.api.nvim_buf_delete(b, { force = true })
-    end
-  end
-end
-
 ------------------------------------------------------------------------------
 -- PROMPT CONSTRUCTION
 ------------------------------------------------------------------------------
@@ -101,7 +92,7 @@ local function build_prompt(user_input, dirs, conf)
   task_lines[#task_lines+1] = "</task>\n"
   table.insert(final_sections, table.concat(task_lines, "\n"))
 
-  -- 4) <file_content path="...">
+  -- 4) <file_content path="..."> from initial_files
   local file_content_blocks = {}
   for _, file_path in ipairs(initial_files) do
     local full_path = root .. "/" .. file_path
@@ -118,6 +109,37 @@ local function build_prompt(user_input, dirs, conf)
     table.insert(final_sections, table.concat(file_content_blocks, "\n\n"))
   end
 
+  -- 4.1) Dynamic file inclusion via @ operator in user_input
+  local dynamic_files = {}
+  for file in user_input:gmatch("@([^%s]+)") do
+    local already_included = false
+    for _, existing in ipairs(initial_files) do
+      if existing == file then
+        already_included = true
+        break
+      end
+    end
+    if not already_included then
+      table.insert(dynamic_files, file)
+    end
+  end
+
+  local dynamic_file_blocks = {}
+  for _, file in ipairs(dynamic_files) do
+    local full_path = root .. "/" .. file
+    if is_subpath(root, full_path) then
+      local fdata = read_file(full_path)
+      if fdata then
+        dynamic_file_blocks[#dynamic_file_blocks+1] = string.format(
+          "<file_content path=\"%s\">\n%s\n</file_content>", file, fdata
+        )
+      end
+    end
+  end
+  if #dynamic_file_blocks > 0 then
+    table.insert(final_sections, table.concat(dynamic_file_blocks, "\n\n"))
+  end
+
   -- 5) <environment_details>
   local env_lines = {}
   env_lines[#env_lines+1] = "<environment_details>"
@@ -133,7 +155,6 @@ local function build_prompt(user_input, dirs, conf)
   env_lines[#env_lines+1] = os.date("%x, %X (%Z)")
   env_lines[#env_lines+1] = ""
   env_lines[#env_lines+1] = "# Current Working Directory (" .. root .. ") Files"
-  -- Using the new get_project_prompt function instead of get_project_structure.
   env_lines[#env_lines+1] = context.get_project_prompt(dirs, conf) or ""
   env_lines[#env_lines+1] = ""
   env_lines[#env_lines+1] = "# Current Mode"
@@ -172,6 +193,8 @@ local function run_chatgpt_command()
   local bufnr = vim.api.nvim_create_buf(false, false)
   vim.api.nvim_buf_set_name(bufnr, "ChatGPT_Prompt.md")
   vim.api.nvim_buf_set_option(bufnr, "filetype", "markdown")
+  -- Set omnifunc for file name auto-completion
+  vim.api.nvim_buf_set_option(bufnr, "omnifunc", "v:lua.chatgpt_file_complete")
   vim.api.nvim_buf_set_option(bufnr, "bufhidden", "wipe")
   vim.api.nvim_buf_set_option(bufnr, "buftype", "")
   vim.api.nvim_buf_set_option(bufnr, "modifiable", true)
@@ -179,6 +202,7 @@ local function run_chatgpt_command()
   vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, {
     "# Enter your main user prompt (task) below.",
     "",
+    "You can include files by typing @filename in your prompt.",
     "Save & close with :wq, :x, or :bd to finalize your prompt."
   })
 
@@ -416,4 +440,29 @@ M.run_chatgpt_command = run_chatgpt_command
 M.run_chatgpt_paste_command = run_chatgpt_paste_command
 M.run_chatgpt_current_buffer_command = run_chatgpt_current_buffer_command
 
+-- New: Global function for file name auto-completion in ChatGPT prompt
+function _G.chatgpt_file_complete(findstart, base)
+  if findstart == 1 then
+    local line = vim.fn.getline('.')
+    local col = vim.fn.col('.')
+    local start = line:sub(1, col):find("@[^%s]*$")
+    if start then
+      return start - 1
+    else
+      return -1
+    end
+  else
+    local conf = config.load()
+    local files = context.get_project_files({'.'}, conf)
+    local completions = {}
+    local esc_base = base:gsub("([^%w])", "%%%1")
+    for _, f in ipairs(files) do
+      if f:match("^" .. esc_base) then
+        table.insert(completions, f)
+      end
+    end
+    return completions
+  end
+end
+
 return M