updating code yank and replacement in popup

lua/ogpt/config.lua

@@ -196,6 +196,7 @@ function M.defaults()
         append = "a",
         prepend = "p",
         yank_code = "c",
+        replace_code = "r",
         yank_to_register = "y",
       },
     },

lua/ogpt/flows/actions/popup/init.lua

@@ -100,52 +100,116 @@ function PopupAction:run()
     self.provider.api:chat_completions(response, {
       custom_params = params,
       partial_result_fn = function(...)
+        opts = opts
         self:on_result(...)
       end,
       should_stop = nil,
     })
   end
 end
 
-function PopupAction:on_result(answer, usage)
-  self:set_loading(false)
-  local lines = utils.split_string_by_line(answer)
-  local _, start_row, start_col, end_row, end_col = self:get_visual_selection()
+function PopupAction:on_result(response)
+  local content = response:pop_content()
+  local answer = content[1]
+  local state = content[2]
+
   local bufnr = self:get_bufnr()
-  if self.strategy == STRATEGY_PREPEND then
-    answer = answer .. "\n" .. self:get_selected_text()
-    vim.api.nvim_buf_set_text(bufnr, start_row - 1, start_col - 1, end_row - 1, end_col, lines)
-  elseif self.strategy == STRATEGY_APPEND then
-    answer = self:get_selected_text() .. "\n\n" .. answer .. "\n"
-    vim.api.nvim_buf_set_text(bufnr, start_row - 1, start_col - 1, end_row - 1, end_col, lines)
-  elseif self.strategy == STRATEGY_REPLACE then
-    answer = answer
-    vim.api.nvim_buf_set_text(bufnr, start_row - 1, start_col - 1, end_row - 1, end_col, lines)
-  elseif self.strategy == STRATEGY_QUICK_FIX then
-    if #lines == 1 and lines[1] == "<OK>" then
-      vim.notify("Your Code looks fine, no issues.", vim.log.levels.INFO)
-      return
-    end
+  local _, start_row, start_col, end_row, end_col = self:get_visual_selection()
 
-    local entries = {}
-    for _, line in ipairs(lines) do
-      local lnum, text = line:match("(%d+):(.*)")
-      if lnum then
-        local entry = { filename = vim.fn.expand("%:p"), lnum = tonumber(lnum), text = text }
-        table.insert(entries, entry)
-      end
-    end
-    if entries then
-      vim.fn.setqflist(entries)
-      vim.cmd(Config.options.show_quickfixes_cmd)
-    end
+  if state == "ERROR" then
+    self:set_loading(false)
+    utils.log("An Error Occurred: " .. answer, vim.log.levels.ERROR)
+    return
   end
 
-  -- set the cursor onto the answer
-  if self.strategy == STRATEGY_APPEND then
-    local target_line = end_row + 3
-    vim.api.nvim_win_set_cursor(0, { target_line, 0 })
+  if state == "END" then
+    self:set_loading(false)
+    vim.api.nvim_buf_set_option(bufnr, "modifiable", true)
+    vim.api.nvim_buf_set_text(
+      bufnr,
+      start_row - 1,
+      start_col - 1,
+      end_row - 1,
+      end_col,
+      utils.split_string_by_line(answer)
+    )
+    return
+  end
+
+  if state == "START" then
+    self:set_loading(true)
+    vim.api.nvim_buf_set_option(bufnr, "modifiable", true)
+    -- Delete the text in the specified selection range
+    vim.api.nvim_buf_set_text(
+      bufnr,
+      opts.selection_idx.start_row - 1,
+      opts.selection_idx.start_col - 1,
+      opts.selection_idx.end_row - 1,
+      opts.selection_idx.end_col,
+      { "" }
+    )
+  end
+
+  if state == "START" or state == "CONTINUE" then
+    vim.api.nvim_buf_set_option(bufnr, "modifiable", true)
+    local lines = vim.split(answer, "\n", {})
+    local length = #lines
+
+    for i, line in ipairs(lines) do
+      if i > 1 then
+        vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { "" }) -- add in the new line here
+      end
+      local currentLine = vim.api.nvim_buf_get_lines(bufnr, -2, -1, false)[1]
+      if currentLine then
+        vim.api.nvim_buf_set_lines(bufnr, -2, -1, false, { currentLine .. line })
+      end
+    end
   end
 end
 
+-- function PopupAction:on_result(response)
+--   local content = response:pop_content()
+--   local answer = content[1]
+--   local state = content[2]
+--
+--   self:set_loading(false)
+--   local lines = utils.split_string_by_line(answer)
+--   local _, start_row, start_col, end_row, end_col = self:get_visual_selection()
+--   local bufnr = self:get_bufnr()
+--   if self.strategy == STRATEGY_PREPEND then
+--     answer = answer .. "\n" .. self:get_selected_text()
+--     vim.api.nvim_buf_set_text(bufnr, start_row - 1, start_col - 1, end_row - 1, end_col, lines)
+--   elseif self.strategy == STRATEGY_APPEND then
+--     answer = self:get_selected_text() .. "\n\n" .. answer .. "\n"
+--     vim.api.nvim_buf_set_text(bufnr, start_row - 1, start_col - 1, end_row - 1, end_col, lines)
+--   elseif self.strategy == STRATEGY_REPLACE then
+--     answer = answer
+--     vim.api.nvim_buf_set_text(bufnr, start_row - 1, start_col - 1, end_row - 1, end_col, lines)
+--   elseif self.strategy == STRATEGY_QUICK_FIX then
+--     if #lines == 1 and lines[1] == "<OK>" then
+--       vim.notify("Your Code looks fine, no issues.", vim.log.levels.INFO)
+--       return
+--     end
+--
+--     local entries = {}
+--     for _, line in ipairs(lines) do
+--       local lnum, text = line:match("(%d+):(.*)")
+--       if lnum then
+--         local entry = { filename = vim.fn.expand("%:p"), lnum = tonumber(lnum), text = text }
+--         table.insert(entries, entry)
+--       end
+--     end
+--     if entries then
+--       vim.fn.setqflist(entries)
+--       vim.cmd(Config.options.show_quickfixes_cmd)
+--     end
+--   end
+--
+--   -- set the cursor onto the answer
+--   if self.strategy == STRATEGY_APPEND then
+--     local target_line = end_row + 3
+--     vim.api.nvim_win_set_cursor(0, { target_line, 0 })
+--   end
+-- end
+
 return PopupAction

lua/ogpt/flows/actions/popup/keymaps.lua

@@ -36,6 +36,40 @@ function M.apply_map(popup, opts)
     vim.cmd("q")
   end)
 
+  -- yank code and replace selection
+  popup:map("n", Config.options.popup.keymaps.replace_code, function()
+    -- Get the lines from the popup buffer
+    local _lines = vim.api.nvim_buf_get_lines(popup.bufnr, 0, -1, false)
+    local _code = utils.getSelectedCode(_lines)
+    local code = vim.split(_code, "\n")
+
+    -- Yank the lines to the specified register
+    vim.fn.setreg("yank_register", table.concat(_lines, "\n"))
+
+    -- Delete the text in the specified selection range
+    vim.api.nvim_buf_set_text(
+      opts.main_bufnr,
+      opts.selection_idx.start_row - 1,
+      opts.selection_idx.start_col - 1,
+      opts.selection_idx.end_row - 1,
+      opts.selection_idx.end_col,
+      { "" }
+    )
+
+    -- Paste the yanked lines over the selection
+    vim.api.nvim_buf_set_text(
+      opts.main_bufnr,
+      opts.selection_idx.start_row - 1,
+      opts.selection_idx.start_col - 1,
+      opts.selection_idx.start_row - 1,
+      opts.selection_idx.start_col - 1,
+      code
+    )
+
+    -- Close the popup
+    vim.cmd("q")
+  end)
+
   -- accept output and append
   popup:map("n", Config.options.popup.keymaps.append, function()
     local _lines = vim.api.nvim_buf_get_lines(popup.bufnr, 0, -1, false)