From 6f53a8d1d27391937d0f0044eaa2e14f1a02173e Mon Sep 17 00:00:00 2001
From: Frank Roeder <fr.coding@proton.me>
Date: Thu, 29 Aug 2024 20:48:54 +0200
Subject: [PATCH 1/4] refactor: make code nicer

---
 lua/parrot/chat_handler.lua | 122 +++++++++++++-------------
 lua/parrot/chat_utils.lua   | 165 +++++++++++++++++-------------------
 lua/parrot/config.lua       |   4 +-
 lua/parrot/state.lua        |  40 ++++-----
 lua/parrot/utils.lua        |  69 ++++++++-------
 5 files changed, 189 insertions(+), 211 deletions(-)

diff --git a/lua/parrot/chat_handler.lua b/lua/parrot/chat_handler.lua
index dbd4b2e..0763eb3 100644
--- a/lua/parrot/chat_handler.lua
+++ b/lua/parrot/chat_handler.lua
@@ -12,7 +12,6 @@ local Job = require("plenary.job")
 local pft = require("plenary.filetype")
 
 local ChatHandler = {}
-
 ChatHandler.__index = ChatHandler
 
 function ChatHandler:new(options, providers, available_providers, available_models, commands)
@@ -22,20 +21,17 @@ function ChatHandler:new(options, providers, available_providers, available_mode
     _plugin_name = "parrot.nvim",
     options = options,
     providers = providers,
-    current_provider = {
-      chat = nil,
-      command = nil,
-    },
+    current_provider = { chat = nil, command = nil },
     pool = Pool:new(),
     queries = Queries:new(),
     commands = commands,
     state = state,
     _toggle = {},
     _toggle_kind = {
-      unknown = 0, -- unknown toggle
-      chat = 1, -- chat toggle
-      popup = 2, -- popup toggle
-      context = 3, -- context toggle
+      unknown = 0,
+      chat = 1,
+      popup = 2,
+      context = 3,
     },
     available_providers = available_providers,
     available_models = available_models,
@@ -60,64 +56,55 @@ function ChatHandler:set_provider(selected_prov, is_chat)
   local endpoint = self.providers[selected_prov].endpoint
   local api_key = self.providers[selected_prov].api_key
   local _prov = init_provider(selected_prov, endpoint, api_key)
-  if is_chat then
-    self.current_provider.chat = _prov
-  else
-    self.current_provider.command = _prov
-  end
+  self.current_provider[is_chat and "chat" or "command"] = _prov
   self.state:set_provider(_prov.name, is_chat)
   self.state:refresh(self.available_providers, self.available_models)
   self:prepare_commands()
 end
 
 function ChatHandler:get_provider(is_chat)
-  local current_prov = nil
-  if is_chat then
-    current_prov = self.current_provider.chat
-  else
-    current_prov = self.current_provider.command
-  end
-
+  local current_prov = self.current_provider[is_chat and "chat" or "command"]
   if not current_prov then
     local prov = self.state:get_provider(is_chat)
+    if not prov then
+      logger.error("No provider found for " .. (is_chat and "chat" or "command"))
+      return nil
+    end
     self:set_provider(prov, is_chat)
+    current_prov = self.current_provider[is_chat and "chat" or "command"]
   end
-
-  if is_chat then
-    return self.current_provider.chat
-  else
-    return self.current_provider.command
-  end
+  return current_prov
 end
 
 function ChatHandler:buf_handler()
   local gid = utils.create_augroup("PrtBufHandler", { clear = true })
-
   utils.autocmd({ "BufEnter" }, nil, function(event)
     local buf = event.buf
-
-    if not vim.api.nvim_buf_is_valid(buf) then
-      return
+    if vim.api.nvim_buf_is_valid(buf) then
+      local file_name = vim.api.nvim_buf_get_name(buf)
+      self:prep_chat(buf, file_name)
+      self:prep_context(buf, file_name)
     end
-
-    local file_name = vim.api.nvim_buf_get_name(buf)
-
-    self:prep_chat(buf, file_name)
-    self:prep_context(buf, file_name)
   end, gid)
 end
 
 function ChatHandler:prep_chat(buf, file_name)
-  if not utils.is_chat(buf, file_name, self.options.chat_dir) then
-    return
-  end
-
-  if buf ~= vim.api.nvim_get_current_buf() then
+  if not utils.is_chat(buf, file_name, self.options.chat_dir) or buf ~= vim.api.nvim_get_current_buf() then
     return
   end
 
   chatutils.prep_md(buf)
 
+  self:setup_chat_prompt(buf)
+  self:setup_chat_commands(buf)
+  self:setup_chat_shortcuts(buf)
+
+  -- remember last opened chat file
+  self.state:set_last_chat(file_name)
+  self.state:refresh(self.available_providers, self.available_models)
+end
+
+function ChatHandler:setup_chat_prompt(buf)
   if self.options.chat_prompt_buf_type then
     vim.api.nvim_set_option_value("buftype", "prompt", { buf = buf })
     vim.fn.prompt_setprompt(buf, "")
@@ -125,8 +112,9 @@ function ChatHandler:prep_chat(buf, file_name)
       self:chat_respond({ args = "" })
     end)
   end
+end
 
-  -- setup chat specific commands
+function ChatHandler:setup_chat_commands(buf)
   local range_commands = {
     {
       command = "ChatRespond",
@@ -142,21 +130,26 @@ function ChatHandler:prep_chat(buf, file_name)
     },
   }
   for _, rc in ipairs(range_commands) do
-    local cmd = self.options.cmd_prefix .. rc.command .. "<cr>"
-    for _, mode in ipairs(rc.modes) do
-      if mode == "n" or mode == "i" then
-        utils.set_keymap({ buf }, mode, rc.shortcut, function()
-          vim.api.nvim_command(self.options.cmd_prefix .. rc.command)
-          -- go to normal mode
-          vim.api.nvim_command("stopinsert")
-          utils.feedkeys("<esc>", "xn")
-        end, rc.comment)
-      else
-        utils.set_keymap({ buf }, mode, rc.shortcut, ":<C-u>'<,'>" .. cmd, rc.comment)
-      end
+    self:setup_chat_command(buf, rc)
+  end
+end
+
+function ChatHandler:setup_chat_command(buf, command_info)
+  local cmd = self.options.cmd_prefix .. command_info.command .. "<cr>"
+  for _, mode in ipairs(command_info.modes) do
+    if mode == "n" or mode == "i" then
+      utils.set_keymap({ buf }, mode, command_info.shortcut, function()
+        vim.api.nvim_command(self.options.cmd_prefix .. command_info.command)
+        vim.api.nvim_command("stopinsert")
+        utils.feedkeys("<esc>", "xn")
+      end, command_info.comment)
+    else
+      utils.set_keymap({ buf }, mode, command_info.shortcut, ":<C-u>'<,'>" .. cmd, command_info.comment)
     end
   end
+end
 
+function ChatHandler:setup_chat_shortcuts(buf)
   local ds = self.options.chat_shortcut_delete
   utils.set_keymap({ buf }, ds.modes, ds.shortcut, function()
     self:chat_delete()
@@ -166,10 +159,6 @@ function ChatHandler:prep_chat(buf, file_name)
   utils.set_keymap({ buf }, ss.modes, ss.shortcut, function()
     self:stop()
   end, "Parrot Chat Stop")
-
-  -- remember last opened chat file
-  self.state:set_last_chat(file_name)
-  self.state:refresh(self.available_providers, self.available_models)
 end
 
 function ChatHandler:prep_context(buf, file_name)
@@ -1392,7 +1381,6 @@ end
 ---@param handler function # response handler
 ---@param on_exit function | nil # optional on_exit handler
 function ChatHandler:query(buf, provider, payload, handler, on_exit)
-  -- make sure handler is a function
   if type(handler) ~= "function" then
     logger.error(
       string.format("query() expects a handler function, but got %s:\n%s", type(handler), vim.inspect(handler))
@@ -1401,6 +1389,7 @@ function ChatHandler:query(buf, provider, payload, handler, on_exit)
   end
 
   if not provider:verify() then
+    logger.error("Provider verification failed")
     return
   end
 
@@ -1448,8 +1437,11 @@ function ChatHandler:query(buf, provider, payload, handler, on_exit)
     on_exit = function(response, exit_code)
       logger.debug("on_exit: " .. vim.inspect(response:result()))
       if exit_code ~= 0 then
-        logger.error("An error occured calling curl .. " .. table.concat(curl_params, " "))
-        on_exit(qid)
+        logger.error("An error occurred calling curl: " .. table.concat(curl_params, " "))
+        if on_exit then
+          on_exit(qid)
+        end
+        return
       end
       local result = response:result()
       result = utils.parse_raw_response(result)
@@ -1459,11 +1451,13 @@ function ChatHandler:query(buf, provider, payload, handler, on_exit)
         response.handle:close()
       end
 
-      on_exit(qid)
+      if on_exit then
+        on_exit(qid)
+      end
       local qt = self.queries:get(qid)
-      if qt.ns_id and qt.buf then
+      if qt and qt.ns_id and qt.buf then
         vim.schedule(function()
-          vim.api.nvim_buf_clear_namespace(qt.buf, qt.ns_id, 0, -1)
+          pcall(vim.api.nvim_buf_clear_namespace, qt.buf, qt.ns_id, 0, -1)
         end)
       end
       self.pool:remove(response.pid)
diff --git a/lua/parrot/chat_utils.lua b/lua/parrot/chat_utils.lua
index 22980f4..5a8e3cc 100644
--- a/lua/parrot/chat_utils.lua
+++ b/lua/parrot/chat_utils.lua
@@ -6,27 +6,17 @@ local M = {}
 ---@param params table | string # table with args or string args
 ---@return number # buf target
 M.resolve_buf_target = function(params)
-  local args = ""
-  if type(params) == "table" then
-    args = params.args or ""
-  else
-    args = params
-  end
-
-  if args == "popup" then
-    return ui.BufTarget.popup
-  elseif args == "split" then
-    return ui.BufTarget.split
-  elseif args == "vsplit" then
-    return ui.BufTarget.vsplit
-  elseif args == "tabnew" then
-    return ui.BufTarget.tabnew
-  else
-    return ui.BufTarget.current
-  end
+  local args = type(params) == "table" and (params.args or "") or params
+  local target_map = {
+    popup = ui.BufTarget.popup,
+    split = ui.BufTarget.split,
+    vsplit = ui.BufTarget.vsplit,
+    tabnew = ui.BufTarget.tabnew,
+  }
+  return target_map[args] or ui.BufTarget.current
 end
-
 -- response handler
+---@param queries table
 ---@param buf number | nil # buffer to insert response into
 ---@param win number | nil # window to insert response into
 ---@param line number | nil # line to insert response into
@@ -35,100 +25,103 @@ end
 ---@param cursor boolean # whether to move cursor to the end of the response
 M.create_handler = function(queries, buf, win, line, first_undojoin, prefix, cursor)
   buf = buf or vim.api.nvim_get_current_buf()
+  win = win or vim.api.nvim_get_current_win()
   prefix = prefix or ""
-  local first_line = line or vim.api.nvim_win_get_cursor(win)[1] - 1
+  local first_line = line or (vim.api.nvim_win_get_cursor(win)[1] - 1)
   local finished_lines = 0
   local skip_first_undojoin = not first_undojoin
 
   local hl_handler_group = "PrtHandlerStandout"
-  vim.cmd("highlight default link " .. hl_handler_group .. " CursorLine")
+  vim.api.nvim_set_hl(0, hl_handler_group, { link = "CursorLine" })
 
   local ns_id = vim.api.nvim_create_namespace("PrtHandler_" .. utils.uuid())
-
   local ex_id = vim.api.nvim_buf_set_extmark(buf, ns_id, first_line, 0, {
     strict = false,
     right_gravity = false,
   })
 
   local response = ""
-  return vim.schedule_wrap(function(qid, chunk)
-    local qt = queries:get(qid)
-    if not qt then
-      return
-    end
-    -- if buf is not valid, stop
+
+  local function update_buffer(qid, chunk)
     if not vim.api.nvim_buf_is_valid(buf) then
       return
     end
-    -- undojoin takes previous change into account, so skip it for the first chunk
-    if skip_first_undojoin then
-      skip_first_undojoin = false
-    else
-      utils.undojoin(buf)
-    end
-
-    if not qt.ns_id then
-      qt.ns_id = ns_id
-    end
-
-    if not qt.ex_id then
-      qt.ex_id = ex_id
-    end
-
-    first_line = vim.api.nvim_buf_get_extmark_by_id(buf, ns_id, ex_id, {})[1]
-
-    -- clean previous response
-    local line_count = #vim.split(response, "\n")
-    vim.api.nvim_buf_set_lines(buf, first_line + finished_lines, first_line + line_count, false, {})
-
-    -- append new response
-    response = response .. chunk
-    utils.undojoin(buf)
-
-    -- prepend prefix to each line
-    local lines = vim.split(response, "\n")
-    for i, l in ipairs(lines) do
-      lines[i] = prefix .. l
-    end
-
-    local unfinished_lines = {}
-    for i = finished_lines + 1, #lines do
-      table.insert(unfinished_lines, lines[i])
-    end
 
-    vim.api.nvim_buf_set_lines(buf, first_line + finished_lines, first_line + finished_lines, false, unfinished_lines)
-
-    local new_finished_lines = math.max(0, #lines - 1)
-    for i = finished_lines, new_finished_lines do
-      vim.api.nvim_buf_add_highlight(buf, qt.ns_id, hl_handler_group, first_line + i, 0, -1)
-    end
-    finished_lines = new_finished_lines
+    vim.api.nvim_buf_call(buf, function()
+      if not skip_first_undojoin then
+        vim.cmd("undojoin")
+      end
+      skip_first_undojoin = false
 
-    local end_line = first_line + #vim.split(response, "\n")
-    qt.first_line = first_line
-    qt.last_line = end_line - 1
+      first_line = vim.api.nvim_buf_get_extmark_by_id(buf, ns_id, ex_id, {})[1]
+
+      -- Clean previous response and append new chunk
+      local line_count = #vim.split(response, "\n")
+      vim.api.nvim_buf_set_lines(buf, first_line + finished_lines, first_line + line_count, false, {})
+      response = response .. chunk
+      vim.cmd("undojoin")
+
+      -- Prepend prefix to each line and update buffer
+      local lines = vim.tbl_map(function(l)
+        return prefix .. l
+      end, vim.split(response, "\n", { plain = true }))
+      local unfinished_lines = vim.list_slice(lines, finished_lines + 1)
+      vim.api.nvim_buf_set_lines(buf, first_line + finished_lines, first_line + finished_lines, false, unfinished_lines)
+
+      -- Update highlighting
+      local new_finished_lines = math.max(0, #lines - 1)
+      for i = finished_lines, new_finished_lines - 1 do
+        vim.api.nvim_buf_add_highlight(buf, ns_id, hl_handler_group, first_line + i, 0, -1)
+      end
+      finished_lines = new_finished_lines
+
+      -- Update query table
+      local end_line = first_line + #lines
+      if queries:get(qid) then
+        queries:get(qid).first_line = first_line
+        queries:get(qid).last_line = end_line - 1
+        queries:get(qid).ns_id = ns_id
+        queries:get(qid).ex_id = ex_id
+      end
+
+      -- Move cursor if needed
+      if cursor then
+        utils.cursor_to_line(end_line, buf, win)
+      end
+    end)
+  end
 
-    -- move cursor to the end of the response
-    if cursor then
-      utils.cursor_to_line(end_line, buf, win)
+  return vim.schedule_wrap(function(qid, chunk)
+    if not queries:get(qid) then
+      return
     end
+    update_buffer(qid, chunk)
   end)
 end
 
 ---@param buf number | nil
 M.prep_md = function(buf)
-  vim.api.nvim_set_option_value("swapfile", false, { buf = buf })
-  vim.api.nvim_set_option_value("filetype", "markdown", { buf = buf })
-
-  -- better text wrapping
-  vim.api.nvim_command("setlocal wrap linebreak")
-  -- auto save on TextChanged, InsertLeave
-  vim.api.nvim_command("autocmd TextChanged,InsertLeave <buffer=" .. buf .. "> silent! write")
-
-  -- register shortcuts local to this buffer
   buf = buf or vim.api.nvim_get_current_buf()
+  local buf_options = {
+    swapfile = false,
+    filetype = "markdown",
+  }
+  local win_options = {
+    wrap = true,
+    linebreak = true,
+  }
+  for option, value in pairs(buf_options) do
+    vim.api.nvim_buf_set_option(buf, option, value)
+  end
+  for option, value in pairs(win_options) do
+    vim.api.nvim_win_set_option(0, option, value)
+  end
+
+  vim.api.nvim_create_autocmd({ "TextChanged", "InsertLeave" }, {
+    buffer = buf,
+    command = "silent! write",
+  })
 
-  -- ensure normal mode
   vim.api.nvim_command("stopinsert")
   utils.feedkeys("<esc>", "xn")
 end
diff --git a/lua/parrot/config.lua b/lua/parrot/config.lua
index e992be1..56f958c 100644
--- a/lua/parrot/config.lua
+++ b/lua/parrot/config.lua
@@ -371,7 +371,7 @@ function M.setup(opts)
     Retry = "retry",
   }
 
-  M.chat_handler = ChatHandler:new(M.options, M.providers, M.available_providers, available_models, M.cmd)
+  M.chat_handler = ChatHandler:new(M.options, M.providers, M.available_providers, M.available_models, M.cmd)
   M.chat_handler:prepare_commands()
   M.add_default_commands(M.cmd, M.hooks, M.options)
   M.chat_handler:buf_handler()
@@ -427,7 +427,7 @@ M.add_default_commands = function(commands, hooks, options)
       end, {
         nargs = "?",
         range = true,
-        desc = "Parrot LLM plugin",
+        desc = "Parrot LLM plugin: " .. cmd,
         complete = function()
           if completions[cmd] then
             return completions[cmd]
diff --git a/lua/parrot/state.lua b/lua/parrot/state.lua
index 5d844be..f52bff6 100644
--- a/lua/parrot/state.lua
+++ b/lua/parrot/state.lua
@@ -42,13 +42,7 @@ end
 --- @param available_models table
 function State:load_models(provider, model_type, available_models)
   local state_model = self.file_state and self.file_state[provider] and self.file_state[provider][model_type]
-  local is_valid_model = false
-
-  if model_type == "chat_model" then
-    is_valid_model = utils.contains(available_models[provider], state_model)
-  elseif model_type == "command_model" then
-    is_valid_model = utils.contains(available_models[provider], state_model)
-  end
+  local is_valid_model = state_model and utils.contains(available_models[provider], state_model)
 
   if self._state[provider][model_type] == nil then
     if state_model and is_valid_model then
@@ -65,20 +59,21 @@ end
 function State:refresh(available_providers, available_models)
   self:init_file_state(available_providers)
   self:init_state(available_providers, available_models)
-  self._state.current_provider.chat = self._state.current_provider.chat
-    or self.file_state.current_provider.chat
-    or available_providers[1]
-  self._state.current_provider.command = self._state.current_provider.command
-    or self.file_state.current_provider.command
-    or available_providers[1]
-  self._state.last_chat = self._state.last_chat or self.file_state.last_chat or nil
 
-  if not utils.contains(available_providers, self._state.current_provider.chat) then
-    self._state.current_provider.chat = available_providers[1]
-  end
-  if not utils.contains(available_providers, self._state.current_provider.command) then
-    self._state.current_provider.command = available_providers[1]
+  local function set_current_provider(key)
+    self._state.current_provider[key] = self._state.current_provider[key]
+      or self.file_state.current_provider[key]
+      or available_providers[1]
+    if not utils.contains(available_providers, self._state.current_provider[key]) then
+      self._state.current_provider[key] = available_providers[1]
+    end
   end
+
+  set_current_provider("chat")
+  set_current_provider("command")
+
+  self._state.last_chat = self._state.last_chat or self.file_state.last_chat or nil
+
   self:save()
 end
 
@@ -124,11 +119,8 @@ end
 --- @param model_type string # Type of model ('chat' or 'command').
 --- @return table|nil
 function State:get_model(provider, model_type)
-  if model_type == "chat" then
-    return self._state[provider].chat_model or self.file_state[provider].chat_model
-  elseif model_type == "command" then
-    return self._state[provider].command_model or self.file_state[provider].command_model
-  end
+  local key = model_type .. "_model"
+  return self._state[provider][key] or self.file_state[provider][key]
 end
 
 --- Sets the last opened chat file path.
diff --git a/lua/parrot/utils.lua b/lua/parrot/utils.lua
index d78f2b6..dca17dd 100644
--- a/lua/parrot/utils.lua
+++ b/lua/parrot/utils.lua
@@ -3,58 +3,57 @@ local pft = require("plenary.filetype")
 local M = {}
 
 -- Trim leading whitespace and tabs from a string.
----@param str string # The input string to be trimmed.
----@return string # The trimmed string.
-M.trim = function(str)
-  return str:gsub("^[\t ]+", ""):gsub("\n[\t ]+", "\n")
+---@param str string The input string to be trimmed.
+---@return string The trimmed string.
+function M.trim(str)
+  return str:gsub("^%s+", ""):gsub("\n%s+", "\n")
 end
 
 -- Feed keys to Neovim.
----@param keys string # string of keystrokes
----@param mode string # string of vim mode ('n', 'i', 'c', etc.), default is 'n'
-M.feedkeys = function(keys, mode)
+---@param keys string String of keystrokes
+---@param mode string String of vim mode ('n', 'i', 'c', etc.), default is 'n'
+function M.feedkeys(keys, mode)
   mode = mode or "n"
-  keys = vim.api.nvim_replace_termcodes(keys, true, false, true)
+  keys = vim.api.nvim_replace_termcodes(keys, true, true, true)
   vim.api.nvim_feedkeys(keys, mode, true)
 end
 
 -- Set keymap for multiple buffers.
----@param buffers table # table of buffers
----@param mode table | string # mode(s) to set keymap for
----@param key string # shortcut key
----@param callback function | string # callback or string to set keymap
----@param desc string | nil # optional description for keymap
-M.set_keymap = function(buffers, mode, key, callback, desc)
+---@param buffers table Table of buffers
+---@param mode table|string Mode(s) to set keymap for
+---@param key string Shortcut key
+---@param callback function|string Callback or string to set keymap
+---@param desc string|nil Optional description for keymap
+function M.set_keymap(buffers, mode, key, callback, desc)
+  local opts = {
+    noremap = true,
+    silent = true,
+    nowait = true,
+    desc = desc,
+  }
   for _, buf in ipairs(buffers) do
-    vim.keymap.set(mode, key, callback, {
-      noremap = true,
-      silent = true,
-      nowait = true,
-      buffer = buf,
-      desc = desc,
-    })
+    opts.buffer = buf
+    vim.keymap.set(mode, key, callback, opts)
   end
 end
 
 -- Create an autocommand for specified events and buffers.
----@param events string | table # events to listen to
----@param buffers table | nil # buffers to listen to (nil for all buffers)
----@param callback function # callback to call
----@param gid number # augroup id
-M.autocmd = function(events, buffers, callback, gid)
+---@param events string|table Events to listen to
+---@param buffers table|nil Buffers to listen to (nil for all buffers)
+---@param callback function Callback to call
+---@param gid number Augroup id
+function M.autocmd(events, buffers, callback, gid)
+  local opts = {
+    group = gid,
+    callback = vim.schedule_wrap(callback),
+  }
   if buffers then
     for _, buf in ipairs(buffers) do
-      vim.api.nvim_create_autocmd(events, {
-        group = gid,
-        buffer = buf,
-        callback = vim.schedule_wrap(callback),
-      })
+      opts.buffer = buf
+      vim.api.nvim_create_autocmd(events, opts)
     end
   else
-    vim.api.nvim_create_autocmd(events, {
-      group = gid,
-      callback = vim.schedule_wrap(callback),
-    })
+    vim.api.nvim_create_autocmd(events, opts)
   end
 end
 

From 480eea94107d00929d6bcdb212ecabbf22446e5f Mon Sep 17 00:00:00 2001
From: Frank Roeder <fr.coding@proton.me>
Date: Tue, 3 Sep 2024 14:14:37 +0200
Subject: [PATCH 2/4] Add improvements

---
 lua/parrot/chat_utils.lua | 145 ++++++++++++++++++--------------------
 1 file changed, 69 insertions(+), 76 deletions(-)

diff --git a/lua/parrot/chat_utils.lua b/lua/parrot/chat_utils.lua
index 5a8e3cc..e7e3fce 100644
--- a/lua/parrot/chat_utils.lua
+++ b/lua/parrot/chat_utils.lua
@@ -3,9 +3,10 @@ local utils = require("parrot.utils")
 
 local M = {}
 
+-- Buffer target resolution
 ---@param params table | string # table with args or string args
 ---@return number # buf target
-M.resolve_buf_target = function(params)
+function M.resolve_buffer_target(params)
   local args = type(params) == "table" and (params.args or "") or params
   local target_map = {
     popup = ui.BufTarget.popup,
@@ -15,25 +16,27 @@ M.resolve_buf_target = function(params)
   }
   return target_map[args] or ui.BufTarget.current
 end
--- response handler
----@param queries table
+
+-- Response handler creation
+---@param queries table # queries object
 ---@param buf number | nil # buffer to insert response into
 ---@param win number | nil # window to insert response into
 ---@param line number | nil # line to insert response into
 ---@param first_undojoin boolean | nil # whether to skip first undojoin
 ---@param prefix string | nil # prefix to insert before each response line
 ---@param cursor boolean # whether to move cursor to the end of the response
-M.create_handler = function(queries, buf, win, line, first_undojoin, prefix, cursor)
+function M.create_handler(queries, buf, win, line, first_undojoin, prefix, cursor)
   buf = buf or vim.api.nvim_get_current_buf()
-  win = win or vim.api.nvim_get_current_win()
   prefix = prefix or ""
-  local first_line = line or (vim.api.nvim_win_get_cursor(win)[1] - 1)
+  local first_line = line or vim.api.nvim_win_get_cursor(win)[1] - 1
   local finished_lines = 0
   local skip_first_undojoin = not first_undojoin
 
+  -- Set up highlighting
   local hl_handler_group = "PrtHandlerStandout"
-  vim.api.nvim_set_hl(0, hl_handler_group, { link = "CursorLine" })
+  vim.cmd("highlight default link " .. hl_handler_group .. " CursorLine")
 
+  -- Create namespace and extmark
   local ns_id = vim.api.nvim_create_namespace("PrtHandler_" .. utils.uuid())
   local ex_id = vim.api.nvim_buf_set_extmark(buf, ns_id, first_line, 0, {
     strict = false,
@@ -42,86 +45,76 @@ M.create_handler = function(queries, buf, win, line, first_undojoin, prefix, cur
 
   local response = ""
 
-  local function update_buffer(qid, chunk)
-    if not vim.api.nvim_buf_is_valid(buf) then
+  return vim.schedule_wrap(function(qid, chunk)
+    local qt = queries:get(qid)
+    if not qt or not vim.api.nvim_buf_is_valid(buf) then
       return
     end
 
-    vim.api.nvim_buf_call(buf, function()
-      if not skip_first_undojoin then
-        vim.cmd("undojoin")
-      end
-      skip_first_undojoin = false
-
-      first_line = vim.api.nvim_buf_get_extmark_by_id(buf, ns_id, ex_id, {})[1]
-
-      -- Clean previous response and append new chunk
-      local line_count = #vim.split(response, "\n")
-      vim.api.nvim_buf_set_lines(buf, first_line + finished_lines, first_line + line_count, false, {})
-      response = response .. chunk
-      vim.cmd("undojoin")
-
-      -- Prepend prefix to each line and update buffer
-      local lines = vim.tbl_map(function(l)
-        return prefix .. l
-      end, vim.split(response, "\n", { plain = true }))
-      local unfinished_lines = vim.list_slice(lines, finished_lines + 1)
-      vim.api.nvim_buf_set_lines(buf, first_line + finished_lines, first_line + finished_lines, false, unfinished_lines)
-
-      -- Update highlighting
-      local new_finished_lines = math.max(0, #lines - 1)
-      for i = finished_lines, new_finished_lines - 1 do
-        vim.api.nvim_buf_add_highlight(buf, ns_id, hl_handler_group, first_line + i, 0, -1)
-      end
-      finished_lines = new_finished_lines
-
-      -- Update query table
-      local end_line = first_line + #lines
-      if queries:get(qid) then
-        queries:get(qid).first_line = first_line
-        queries:get(qid).last_line = end_line - 1
-        queries:get(qid).ns_id = ns_id
-        queries:get(qid).ex_id = ex_id
-      end
-
-      -- Move cursor if needed
-      if cursor then
-        utils.cursor_to_line(end_line, buf, win)
-      end
-    end)
-  end
+    -- Handle undojoin
+    if not skip_first_undojoin then
+      utils.undojoin(buf)
+    end
+    skip_first_undojoin = false
+
+    -- Set namespace and extmark IDs if not set
+    qt.ns_id = qt.ns_id or ns_id
+    qt.ex_id = qt.ex_id or ex_id
+
+    first_line = vim.api.nvim_buf_get_extmark_by_id(buf, ns_id, ex_id, {})[1]
+
+    -- Update response
+    response = response .. chunk
+    local lines = vim.split(response, "\n")
+    local prefixed_lines = vim.tbl_map(function(l)
+      return prefix .. l
+    end, lines)
+
+    -- Update buffer content
+    vim.api.nvim_buf_set_lines(buf, first_line + finished_lines, first_line + #vim.split(response, "\n"), false, {})
+    vim.api.nvim_buf_set_lines(
+      buf,
+      first_line + finished_lines,
+      first_line + finished_lines,
+      false,
+      vim.list_slice(prefixed_lines, finished_lines + 1)
+    )
+
+    -- Update highlighting
+    local new_finished_lines = math.max(0, #lines - 1)
+    for i = finished_lines, new_finished_lines do
+      vim.api.nvim_buf_add_highlight(buf, qt.ns_id, hl_handler_group, first_line + i, 0, -1)
+    end
+    finished_lines = new_finished_lines
 
-  return vim.schedule_wrap(function(qid, chunk)
-    if not queries:get(qid) then
-      return
+    -- Update query object
+    local end_line = first_line + #lines
+    qt.first_line = first_line
+    qt.last_line = end_line - 1
+
+    -- Move cursor if needed
+    if cursor then
+      utils.cursor_to_line(end_line, buf, win)
     end
-    update_buffer(qid, chunk)
   end)
 end
 
+-- Markdown buffer preparation
 ---@param buf number | nil
-M.prep_md = function(buf)
+function M.prepare_markdown_buffer(buf)
   buf = buf or vim.api.nvim_get_current_buf()
-  local buf_options = {
-    swapfile = false,
-    filetype = "markdown",
-  }
-  local win_options = {
-    wrap = true,
-    linebreak = true,
-  }
-  for option, value in pairs(buf_options) do
-    vim.api.nvim_buf_set_option(buf, option, value)
-  end
-  for option, value in pairs(win_options) do
-    vim.api.nvim_win_set_option(0, option, value)
-  end
-
-  vim.api.nvim_create_autocmd({ "TextChanged", "InsertLeave" }, {
-    buffer = buf,
-    command = "silent! write",
-  })
 
+  -- Set buffer options
+  vim.api.nvim_set_option_value("swapfile", false, { buf = buf })
+  vim.api.nvim_set_option_value("filetype", "markdown", { buf = buf })
+
+  -- Set text wrapping
+  vim.api.nvim_command("setlocal wrap linebreak")
+
+  -- Set up auto-save
+  vim.api.nvim_command("autocmd TextChanged,InsertLeave <buffer=" .. buf .. "> silent! write")
+
+  -- Ensure normal mode
   vim.api.nvim_command("stopinsert")
   utils.feedkeys("<esc>", "xn")
 end

From 044584c623d7eb7ad0d413bcd04fa38528832fb4 Mon Sep 17 00:00:00 2001
From: Frank Roeder <fr.coding@proton.me>
Date: Wed, 9 Oct 2024 10:09:51 +0200
Subject: [PATCH 3/4] Update chat_utils test

---
 tests/parrot/chat_utils_spec.lua | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/tests/parrot/chat_utils_spec.lua b/tests/parrot/chat_utils_spec.lua
index 469ad4d..8bb18f9 100644
--- a/tests/parrot/chat_utils_spec.lua
+++ b/tests/parrot/chat_utils_spec.lua
@@ -3,18 +3,19 @@ local ui = require("parrot.ui")
 local async = require("plenary.async")
 
 describe("chat_utils", function()
-  describe("resolve_buf_target", function()
+  describe("resolve_buffer_target", function()
     it("should resolve buffer target correctly", function()
-      assert.are.equal(ui.BufTarget.popup, chat_utils.resolve_buf_target("popup"))
-      assert.are.equal(ui.BufTarget.split, chat_utils.resolve_buf_target("split"))
-      assert.are.equal(ui.BufTarget.vsplit, chat_utils.resolve_buf_target("vsplit"))
-      assert.are.equal(ui.BufTarget.tabnew, chat_utils.resolve_buf_target("tabnew"))
-      assert.are.equal(ui.BufTarget.current, chat_utils.resolve_buf_target(""))
-      assert.are.equal(ui.BufTarget.current, chat_utils.resolve_buf_target({}))
+      assert.are.equal(ui.BufTarget.popup, chat_utils.resolve_buffer_target({ args = "popup" }))
+      assert.are.equal(ui.BufTarget.split, chat_utils.resolve_buffer_target({ args = "split" }))
+      assert.are.equal(ui.BufTarget.vsplit, chat_utils.resolve_buffer_target({ args = "vsplit" }))
+      assert.are.equal(ui.BufTarget.tabnew, chat_utils.resolve_buffer_target({ args = "tabnew" }))
+      assert.are.equal(ui.BufTarget.current, chat_utils.resolve_buffer_target({ args = "" }))
+      assert.are.equal(ui.BufTarget.current, chat_utils.resolve_buffer_target({}))
+      assert.are.equal(ui.BufTarget.current, chat_utils.resolve_buffer_target(""))
+      assert.are.equal(ui.BufTarget.popup, chat_utils.resolve_buffer_target("popup"))
     end)
   end)
-
-  describe("prep_md", function()
+  describe("prepare_markdown_buffer", function()
     it("should set buffer and window options correctly", function()
       async.run(function()
         local buf = vim.api.nvim_create_buf(false, true)
@@ -26,7 +27,7 @@ describe("chat_utils", function()
           col = 5,
         })
 
-        chat_utils.prep_md(buf)
+        chat_utils.prepare_markdown_buffer(buf)
 
         -- Check buffer option
         assert.is_false(vim.api.nvim_buf_get_option(buf, "swapfile"))

From 5e1fcb7d0386c3e812c4bc7d5a04404d0d2beb5c Mon Sep 17 00:00:00 2001
From: Frank Roeder <fr.coding@proton.me>
Date: Fri, 11 Oct 2024 10:55:42 +0200
Subject: [PATCH 4/4] WIP

---
 lua/parrot/buffer_manager.lua  | 137 ++++++++
 lua/parrot/chat_handler.lua    | 574 +++++++++++++++------------------
 lua/parrot/command_manager.lua | 240 ++++++++++++++
 lua/parrot/toggle_manager.lua  |  46 +++
 lua/parrot/utils.lua           |  28 +-
 5 files changed, 700 insertions(+), 325 deletions(-)
 create mode 100644 lua/parrot/buffer_manager.lua
 create mode 100644 lua/parrot/command_manager.lua
 create mode 100644 lua/parrot/toggle_manager.lua

diff --git a/lua/parrot/buffer_manager.lua b/lua/parrot/buffer_manager.lua
new file mode 100644
index 0000000..5070aab
--- /dev/null
+++ b/lua/parrot/buffer_manager.lua
@@ -0,0 +1,137 @@
+local utils = require("parrot.utils")
+local chatutils = require("parrot.chat_utils")
+local logger = require("parrot.logger")
+
+--- BufferManager Module
+-- Handles buffer-related operations such as preparing chat and context buffers,
+-- setting buffer options, and managing auto-save functionality.
+local BufferManager = {}
+BufferManager.__index = BufferManager
+
+--- Creates a new BufferManager instance.
+-- @param options table Configuration options.
+-- @param state State instance for managing persistent state.
+-- @return BufferManager
+function BufferManager:new(options, state)
+  local instance = setmetatable({
+    options = options,
+    state = state,
+    auto_save_groups = {}, -- Tracks autocmd groups per buffer
+  }, BufferManager)
+  return instance
+end
+
+--- Prepares a chat buffer.
+-- Sets up the buffer with necessary options, prepares markdown formatting,
+-- configures auto-save, and updates the plugin state.
+-- @param buf number Buffer number.
+-- @param file_name string Name of the chat file.
+function BufferManager:prepare_chat(buf, file_name)
+  if not utils.is_chat(buf, file_name, self.options.chat_dir) then
+    return
+  end
+
+  -- Prepare the buffer with markdown settings
+  chatutils.prepare_markdown_buffer(buf)
+
+  -- Set up auto-save for the buffer
+  self:setup_auto_save(buf)
+
+  -- Remember the last opened chat file in the state
+  self.state:set_last_chat(file_name)
+  self.state:refresh(self.options.available_providers, self.options.available_models)
+end
+
+--- Prepares a context buffer.
+-- Configures buffers that hold additional context information.
+-- @param buf number Buffer number.
+-- @param file_name string Name of the context file.
+function BufferManager:prepare_context(buf, file_name)
+  if not utils.ends_with(file_name, ".parrot.md") then
+    return
+  end
+
+  if buf ~= vim.api.nvim_get_current_buf() then
+    return
+  end
+
+  chatutils.prepare_markdown_buffer(buf)
+end
+
+--- Sets up auto-save for a buffer.
+-- Configures autocmds to automatically save the buffer on text changes and insert leave events.
+-- Prevents multiple autocmd registrations for the same buffer.
+-- @param buf number Buffer number.
+function BufferManager:setup_auto_save(buf)
+  -- Ensure that auto-save is not already set up for this buffer
+  if self.auto_save_groups[buf] then
+    return
+  end
+
+  local group_name = "PrtAutoSave_" .. buf
+  -- Ensure group_name is less than or equal to 30 characters
+  group_name = group_name:sub(1, 30)
+
+  local group = vim.api.nvim_create_augroup(group_name, { clear = true })
+
+  vim.api.nvim_create_autocmd({ "TextChanged", "InsertLeave" }, {
+    group = group, -- Use the group ID
+    buffer = buf,
+    callback = function()
+      local ok, err = pcall(function()
+        vim.api.nvim_buf_call(buf, function()
+          vim.cmd("silent! write")
+        end)
+      end)
+      if not ok then
+        logger.error(string.format("Auto-save failed for buffer %d: %s", buf, err))
+      end
+    end,
+  })
+
+  self.auto_save_groups[buf] = group
+end
+
+--- Cleans up auto-save autocmds for a buffer.
+-- Removes autocmd groups associated with a buffer to prevent memory leaks.
+-- @param buf number Buffer number.
+function BufferManager:cleanup_auto_save(buf)
+  local group = self.auto_save_groups[buf]
+  if group then
+    pcall(vim.api.nvim_del_augroup_by_id, group)
+    self.auto_save_groups[buf] = nil
+  end
+end
+
+--- Cleans up all resources.
+-- Removes all autocmd groups managed by BufferManager.
+function BufferManager:cleanup()
+  for buf, group in pairs(self.auto_save_groups) do
+    pcall(vim.api.nvim_del_augroup_by_id, group)
+  end
+  self.auto_save_groups = {}
+end
+
+--- Prepares a markdown buffer.
+-- Sets buffer options, text wrapping, and auto-save.
+-- @param buf number Buffer number.
+function BufferManager:prepare_markdown_buffer(buf)
+  buf = buf or vim.api.nvim_get_current_buf()
+
+  -- Set buffer options
+  vim.api.nvim_buf_set_option(buf, "swapfile", false)
+  vim.api.nvim_buf_set_option(buf, "filetype", "markdown")
+
+  -- Set text wrapping
+  vim.api.nvim_buf_set_option(buf, "wrap", true)
+  vim.api.nvim_buf_set_option(buf, "linebreak", true)
+
+  -- Set up auto-save for the buffer
+  self:setup_auto_save(buf)
+
+  -- Ensure normal mode
+  vim.api.nvim_command("stopinsert")
+  utils.feedkeys("<esc>", "xn")
+end
+
+return BufferManager
diff --git a/lua/parrot/chat_handler.lua b/lua/parrot/chat_handler.lua
index 2605f82..7126e3b 100644
--- a/lua/parrot/chat_handler.lua
+++ b/lua/parrot/chat_handler.lua
@@ -11,6 +11,9 @@ local Spinner = require("parrot.spinner")
 local Job = require("plenary.job")
 local pft = require("plenary.filetype")
 local ResponseHandler = require("parrot.response_handler")
+local BufferManager = require("parrot.buffer_manager")
+local CommandManager = require("parrot.command_manager")
+local ToggleManager = require("parrot.toggle_manager")
 
 local ChatHandler = {}
 ChatHandler.__index = ChatHandler
@@ -25,6 +28,8 @@ function ChatHandler:new(options, providers, available_providers, available_mode
     current_provider = { chat = nil, command = nil },
     pool = Pool:new(),
     queries = Queries:new(),
+    buffer_manager = BufferManager:new(options, state),
+		toggle_manager = ToggleManager:new(),
     commands = commands,
     state = state,
     _toggle = {},
@@ -45,6 +50,8 @@ function ChatHandler:new(options, providers, available_providers, available_mode
   }, self)
 end
 
+--- Retrieves status information about the current buffer.
+---@return table { is_chat = boolean, prov = table | nil, model = string }
 function ChatHandler:get_status_info()
   local buf = vim.api.nvim_get_current_buf()
   local file_name = vim.api.nvim_buf_get_name(buf)
@@ -53,6 +60,9 @@ function ChatHandler:get_status_info()
   return { is_chat = is_chat, prov = self.current_provider, model = model_obj.name }
 end
 
+--- Sets the current provider for chat or command.
+---@param selected_prov string Selected provider name.
+---@param is_chat boolean True for chat provider, false for command provider.
 function ChatHandler:set_provider(selected_prov, is_chat)
   local endpoint = self.providers[selected_prov].endpoint
   local api_key = self.providers[selected_prov].api_key
@@ -60,15 +70,18 @@ function ChatHandler:set_provider(selected_prov, is_chat)
   self.current_provider[is_chat and "chat" or "command"] = _prov
   self.state:set_provider(_prov.name, is_chat)
   self.state:refresh(self.available_providers, self.available_models)
-  self:prepare_commands()
+  self.command_manager:prepare_commands() -- Delegate to CommandManager
 end
 
+--- Retrieves the current provider for chat or command.
+---@param is_chat boolean True for chat provider, false for command provider.
+---@return table | nil Provider table or nil if not found.
 function ChatHandler:get_provider(is_chat)
   local current_prov = self.current_provider[is_chat and "chat" or "command"]
   if not current_prov then
     local prov = self.state:get_provider(is_chat)
     if not prov then
-      logger.error("No provider found for " .. (is_chat and "chat" or "command"))
+      logger.error(string.format("No provider found for %s", is_chat and "chat" or "command"))
       return nil
     end
     self:set_provider(prov, is_chat)
@@ -77,32 +90,23 @@ function ChatHandler:get_provider(is_chat)
   return current_prov
 end
 
+--- Handles buffer events by delegating to BufferManager.
 function ChatHandler:buf_handler()
   local gid = utils.create_augroup("PrtBufHandler", { clear = true })
+
   utils.autocmd({ "BufEnter" }, nil, function(event)
     local buf = event.buf
-    if vim.api.nvim_buf_is_valid(buf) then
-      local file_name = vim.api.nvim_buf_get_name(buf)
-      self:prep_chat(buf, file_name)
-      self:prep_context(buf, file_name)
-    end
-  end, gid)
-end
 
-function ChatHandler:prep_chat(buf, file_name)
-  if not utils.is_chat(buf, file_name, self.options.chat_dir) or buf ~= vim.api.nvim_get_current_buf() then
-    return
-  end
-
-  chatutils.prep_md(buf)
+    if not vim.api.nvim_buf_is_valid(buf) then
+      return
+    end
 
-  self:setup_chat_prompt(buf)
-  self:setup_chat_commands(buf)
-  self:setup_chat_shortcuts(buf)
+    local file_name = vim.api.nvim_buf_get_name(buf)
 
-  -- remember last opened chat file
-  self.state:set_last_chat(file_name)
-  self.state:refresh(self.available_providers, self.available_models)
+    -- Delegate to BufferManager
+    self.buffer_manager:prepare_chat(buf, file_name)
+    self.buffer_manager:prepare_context(buf, file_name)
+  end, gid)
 end
 
 function ChatHandler:setup_chat_prompt(buf)
@@ -162,68 +166,46 @@ function ChatHandler:setup_chat_shortcuts(buf)
   end, "Parrot Chat Stop")
 end
 
+--- Delegates context preparation to BufferManager.
+---@param buf number Buffer number.
+---@param file_name string Name of the context file.
 function ChatHandler:prep_context(buf, file_name)
-  if not utils.ends_with(file_name, ".parrot.md") then
-    return
-  end
-
-  if buf ~= vim.api.nvim_get_current_buf() then
-    return
-  end
-
-  chatutils.prep_md(buf)
+  self.buffer_manager:prepare_context(buf, file_name)
 end
 
----@param kind number # kind of toggle
----@return boolean # true if toggle was closed
+--- Toggles (closes) a specific toggle kind using ToggleManager.
+---@param kind number Kind of toggle.
+---@return boolean True if toggle was closed.
 function ChatHandler:toggle_close(kind)
-  if
-    self._toggle[kind]
-    and self._toggle[kind].win
-    and self._toggle[kind].buf
-    and self._toggle[kind].close
-    and vim.api.nvim_win_is_valid(self._toggle[kind].win)
-    and vim.api.nvim_buf_is_valid(self._toggle[kind].buf)
-    and vim.api.nvim_win_get_buf(self._toggle[kind].win) == self._toggle[kind].buf
-  then
-    if #vim.api.nvim_list_wins() == 1 then
-      logger.warning("Can't close the last window.")
-    else
-      self._toggle[kind].close()
-      self._toggle[kind] = nil
-    end
-    return true
-  end
-  self._toggle[kind] = nil
-  return false
+  return self.toggle_manager:close(kind)
 end
 
----@param kind number # kind of toggle
----@param toggle table # table containing `win`, `buf`, and `close` information
+--- Adds a toggle using ToggleManager.
+---@param kind number Kind of toggle.
+---@param toggle table Table containing `win`, `buf`, and `close` information.
 function ChatHandler:toggle_add(kind, toggle)
-  self._toggle[kind] = toggle
+  self.toggle_manager:add(kind, toggle)
 end
 
----@param kind string # string representation of the toggle kind
----@return number # numeric kind of the toggle
+--- Resolves a toggle kind string to its numeric representation using ToggleManager.
+---@param kind string String representation of the toggle kind.
+---@return number Numeric kind of the toggle.
 function ChatHandler:toggle_resolve(kind)
-  kind = kind:lower()
-  if kind == "chat" then
-    return self._toggle_kind.chat
-  elseif kind == "popup" then
-    return self._toggle_kind.popup
-  elseif kind == "context" then
-    return self._toggle_kind.context
-  end
-  logger.warning("Unknown toggle kind: " .. kind)
-  return self._toggle_kind.unknown
+  return self.toggle_manager:resolve(kind)
 end
 
----@return table # { name, system_prompt, provider }
+--- Retrieves the model information based on the model type.
+---@param model_type string "chat" or "command".
+---@return table { name, system_prompt, provider }
 function ChatHandler:get_model(model_type)
-  local prov = self:get_provider(model_type == "chat")
+  local is_chat = model_type == "chat"
+  local prov = self:get_provider(is_chat)
+  if not prov then
+    logger.error("Provider not available for model type: " .. model_type)
+    return {}
+  end
   local model = self.state:get_model(prov.name, model_type)
-  local system_prompt = self.options.system_prompt[model_type]
+  local system_prompt = self.options.system_prompt[model_type] or ""
   return {
     name = model,
     system_prompt = system_prompt,
@@ -231,55 +213,23 @@ function ChatHandler:get_model(model_type)
   }
 end
 
--- creates prompt commands for each target
+--- Prepares commands by delegating to CommandManager.
 function ChatHandler:prepare_commands()
-  for name, target in pairs(ui.Target) do
-    -- uppercase first letter
-    local command = name:gsub("^%l", string.upper)
-
-    local model_obj = self:get_model("command")
-    -- popup is like ephemeral one off chat
-    if target == ui.Target.popup then
-      model_obj = self:get_model("chat")
-    end
-
-    local cmd = function(params)
-      -- template is chosen dynamically based on mode in which the command is called
-      local template = self.options.template_command
-      if params.range == 2 then
-        template = self.options.template_selection
-        -- rewrite needs custom template
-        if target == ui.Target.rewrite then
-          template = self.options.template_rewrite
-        end
-        if target == ui.Target.append then
-          template = self.options.template_append
-        end
-        if target == ui.Target.prepend then
-          template = self.options.template_prepend
-        end
-      end
-      local cmd_prefix = utils.template_render_from_list(
-        self.options.command_prompt_prefix_template,
-        { ["{{llm}}"] = self:get_model("command").name }
-      )
-      self:prompt(params, target, model_obj, cmd_prefix, utils.trim(template), true)
-    end
-    self.commands[command] = command
-    self:addCommand(command, function(params)
-      cmd(params)
-    end)
+  if not self.command_manager then
+    self.command_manager = CommandManager:new(self.options, self)
   end
+  self.command_manager:prepare_commands()
 end
 
+--- Adds a command by delegating to CommandManager.
+-- @param command string Command name.
+-- @param cmd function Command callback.
 function ChatHandler:addCommand(command, cmd)
-  self[command] = function(self, params)
-    cmd(params)
-  end
+  self.command_manager:add_command(command, cmd)
 end
 
--- stop receiving responses for all processes and create a new pool
----@param signal number | nil # signal to send to the process
+--- Stops all ongoing processes by killing associated jobs.
+---@param signal number | nil Signal to send to the processes.
 function ChatHandler:stop(signal)
   if self.pool:is_empty() then
     return
@@ -294,9 +244,11 @@ function ChatHandler:stop(signal)
   self.pool = Pool:new()
 end
 
+--- Handles context-related actions.
+---@param params table Parameters for the context action.
 function ChatHandler:context(params)
   self:toggle_close(self._toggle_kind.popup)
-  -- if there is no selection, try to close context toggle
+  -- If there is no selection, try to close context toggle
   if params.range ~= 2 then
     if self:toggle_close(self._toggle_kind.context) then
       return
@@ -326,7 +278,7 @@ function ChatHandler:context(params)
   if params.args == "" then
     params.args = self.options.toggle_target
   end
-  local target = chatutils.resolve_buf_target(params)
+  local target = chatutils.resolve_buffer_target(params)
   buf = self:open_buf(file_name, target, self._toggle_kind.context, true)
 
   if params.range == 2 then
@@ -336,10 +288,16 @@ function ChatHandler:context(params)
   utils.feedkeys("G", "xn")
 end
 
+--- Opens a buffer based on the provided parameters.
+---@param file_name string Name of the file to open.
+---@param target number Buffer target.
+---@param kind number Kind of toggle.
+---@param toggle boolean Whether to toggle the buffer.
+---@return number Buffer number.
 function ChatHandler:open_buf(file_name, target, kind, toggle)
   target = target or ui.BufTarget.current
 
-  -- close previous popup if it exists
+  -- Close previous popup if it exists
   self:toggle_close(self._toggle_kind.popup)
 
   if toggle then
@@ -373,19 +331,19 @@ function ChatHandler:open_buf(file_name, target, kind, toggle)
     end
 
     if old_buf == nil then
-      -- read file into buffer and force write it
+      -- Read file into buffer and force write it
       vim.api.nvim_command("silent 0read " .. file_name)
       vim.api.nvim_command("silent file " .. file_name)
       vim.api.nvim_set_option_value("filetype", "markdown", { buf = buf })
     else
-      -- move cursor to the beginning of the file and scroll to the end
+      -- Move cursor to the beginning of the file and scroll to the end
       utils.feedkeys("ggG", "xn")
     end
 
-    -- delete whitespace lines at the end of the file
+    -- Delete whitespace lines at the end of the file
     local last_content_line = utils.last_content_line(buf)
     vim.api.nvim_buf_set_lines(buf, last_content_line, -1, false, {})
-    -- insert a new line at the end of the file
+    -- Insert a new line at the end of the file
     vim.api.nvim_buf_set_lines(buf, -1, -1, false, { "" })
     vim.api.nvim_command("silent write! " .. file_name)
   elseif target == ui.BufTarget.split then
@@ -395,7 +353,7 @@ function ChatHandler:open_buf(file_name, target, kind, toggle)
   elseif target == ui.BufTarget.tabnew then
     vim.api.nvim_command("tabnew " .. file_name)
   else
-    -- is it already open in a buffer?
+    -- Check if it's already open in a buffer
     for _, b in ipairs(vim.api.nvim_list_bufs()) do
       if vim.api.nvim_buf_get_name(b) == file_name then
         for _, w in ipairs(vim.api.nvim_list_wins()) do
@@ -407,7 +365,7 @@ function ChatHandler:open_buf(file_name, target, kind, toggle)
       end
     end
 
-    -- open in new buffer
+    -- Open in new buffer
     vim.api.nvim_command("edit " .. file_name)
   end
 
@@ -442,14 +400,19 @@ function ChatHandler:open_buf(file_name, target, kind, toggle)
   return buf
 end
 
+--- Creates a new chat file.
+---@param params table Parameters for creating a new chat.
+---@param toggle boolean Whether to toggle the chat buffer.
+---@param chat_prompt string Optional chat prompt.
+---@return number Buffer number.
 function ChatHandler:_new_chat(params, toggle, chat_prompt)
   self:toggle_close(self._toggle_kind.popup)
 
-  -- prepare filename
+  -- Prepare filename
   local time = os.date("%Y-%m-%d.%H-%M-%S")
   local stamp = tostring(math.floor(vim.uv.hrtime() / 1000000) % 1000)
   local cbuf = vim.api.nvim_get_current_buf()
-  -- make sure stamp is 3 digits
+  -- Make sure stamp is 3 digits
   while #stamp < 3 do
     stamp = "0" .. stamp
   end
@@ -471,14 +434,14 @@ function ChatHandler:_new_chat(params, toggle, chat_prompt)
     ["{{user}}"] = self.options.chat_user_prefix,
     ["{{optional}}"] = chat_prompt,
   })
-  -- escape underscores (for markdown)
+  -- Escape underscores (for markdown)
   template = template:gsub("_", "\\_")
-  -- strip leading and trailing newlines
+  -- Strip leading and trailing newlines
   template = template:gsub("^%s*(.-)%s*$", "%1") .. "\n"
 
-  -- create chat file
+  -- Create chat file
   vim.fn.writefile(vim.split(template, "\n"), filename)
-  local target = chatutils.resolve_buf_target(params)
+  local target = chatutils.resolve_buffer_target(params)
   local buf = self:open_buf(filename, target, self._toggle_kind.chat, toggle)
 
   if params.range == 2 then
@@ -488,9 +451,12 @@ function ChatHandler:_new_chat(params, toggle, chat_prompt)
   return buf
 end
 
----@return number # buffer number
+--- Creates a new chat.
+---@param params table Parameters for creating a new chat.
+---@param chat_prompt string Optional chat prompt.
+---@return number Buffer number.
 function ChatHandler:chat_new(params, chat_prompt)
-  -- if chat toggle is open, close it and start a new one
+  -- If chat toggle is open, close it and start a new one
   if self:toggle_close(self._toggle_kind.chat) then
     params.args = params.args or ""
     if params.args == "" then
@@ -505,23 +471,25 @@ function ChatHandler:chat_new(params, chat_prompt)
   return self:_new_chat(params, false, chat_prompt)
 end
 
+--- Toggles the chat buffer.
+---@param params table Parameters for toggling the chat.
 function ChatHandler:chat_toggle(params)
   self:toggle_close(self._toggle_kind.popup)
   if self:toggle_close(self._toggle_kind.chat) and params.range ~= 2 then
     return
   end
 
-  -- create new chat file otherwise
+  -- Create new chat file otherwise
   params.args = params.args or ""
   if params.args == "" then
     params.args = self.options.toggle_target
   end
 
-  -- if the range is 2, we want to create a new chat file with the selection
+  -- If the range is 2, we want to create a new chat file with the selection
   if params.range ~= 2 then
     local last_chat_file = self.state:get_last_chat()
     if last_chat_file and vim.fn.filereadable(last_chat_file) == 1 then
-      self:open_buf(last_chat_file, chatutils.resolve_buf_target(params), self._toggle_kind.chat, true)
+      self:open_buf(last_chat_file, chatutils.resolve_buffer_target(params), self._toggle_kind.chat, true)
       return
     end
   end
@@ -529,19 +497,21 @@ function ChatHandler:chat_toggle(params)
   self:_new_chat(params, true)
 end
 
+--- Pastes selected text into the last chat.
+---@param params table Parameters for pasting.
 function ChatHandler:chat_paste(params)
-  -- if there is no selection, do nothing
+  -- If there is no selection, do nothing
   if params.range ~= 2 then
     logger.warning("Please select some text to paste into the chat.")
     return
   end
 
-  -- get current buffer
+  -- Get current buffer
   local cbuf = vim.api.nvim_get_current_buf()
 
   local last_chat_file = self.state:get_last_chat()
   if last_chat_file and vim.fn.filereadable(last_chat_file) ~= 1 then
-    -- skip rest since new chat will handle snippet on it's own
+    -- Skip rest since new chat will handle snippet on its own
     self:chat_new(params)
     return
   end
@@ -550,7 +520,7 @@ function ChatHandler:chat_paste(params)
   if params.args == "" then
     params.args = self.options.toggle_target
   end
-  local target = chatutils.resolve_buf_target(params)
+  local target = chatutils.resolve_buffer_target(params)
   local buf = utils.get_buffer(last_chat_file)
   local win_found = false
   if buf then
@@ -569,24 +539,25 @@ function ChatHandler:chat_paste(params)
   utils.feedkeys("G", "xn")
 end
 
+--- Deletes the current chat file.
 function ChatHandler:chat_delete()
-  -- get buffer and file
+  -- Get buffer and file
   local buf = vim.api.nvim_get_current_buf()
   local file_name = vim.api.nvim_buf_get_name(buf)
 
-  -- check if file is in the chat dir
+  -- Check if file is in the chat dir
   if not utils.starts_with(file_name, self.options.chat_dir) then
     logger.warning("File " .. vim.inspect(file_name) .. " is not in chat dir")
     return
   end
 
-  -- delete without confirmation
+  -- Delete without confirmation
   if not self.options.chat_confirm_delete then
     futils.delete_file(file_name, self.options.chat_dir)
     return
   end
 
-  -- ask for confirmation
+  -- Ask for confirmation
   vim.ui.input({ prompt = "Delete " .. file_name .. "? [y/N] " }, function(input)
     if input and input:lower() == "y" then
       futils.delete_file(file_name, self.options.chat_dir)
@@ -594,6 +565,8 @@ function ChatHandler:chat_delete()
   end)
 end
 
+--- Handles chat responses by delegating to the querying system.
+---@param params table Parameters for responding.
 function ChatHandler:_chat_respond(params)
   local buf = vim.api.nvim_get_current_buf()
   local win = vim.api.nvim_get_current_win()
@@ -606,31 +579,31 @@ function ChatHandler:_chat_respond(params)
     return
   end
 
-  -- go to normal mode
+  -- Go to normal mode
   vim.cmd("stopinsert")
 
-  -- get all lines
+  -- Get all lines
   local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
 
-  -- check if file looks like a chat file
+  -- Check if file looks like a chat file
   local file_name = vim.api.nvim_buf_get_name(buf)
   if not utils.is_chat(buf, file_name, self.options.chat_dir) then
     logger.warning("File " .. vim.inspect(file_name) .. " does not look like a chat file")
     return
   end
 
-  -- headers are fields before first message ---
+  -- Headers are fields before first message ---
   local headers = {}
   local header_end = nil
-  local line_idx = 0
-  ---parse headers
+  local line_idx = 1 -- Lua tables are 1-based
+  -- Parse headers
   for _, line in ipairs(lines) do
-    -- first line starts with ---
+    -- First line starts with ---
     if line:sub(1, 3) == "---" then
       header_end = line_idx
       break
     end
-    -- parse header fields
+    -- Parse header fields
     local key, value = line:match("^[-#] (%w+): (.*)")
     if key ~= nil then
       headers[key] = value
@@ -644,12 +617,12 @@ function ChatHandler:_chat_respond(params)
     return
   end
 
-  -- message needs role and content
+  -- Messages need role and content
   local messages = {}
   local role = ""
   local content = ""
 
-  -- iterate over lines
+  -- Iterate over lines
   local start_index = header_end + 1
   local end_index = #lines
   if params.range == 2 then
@@ -658,7 +631,6 @@ function ChatHandler:_chat_respond(params)
   end
 
   if headers.system and headers.system:match("%S") then
-    ---@diagnostic disable-next-line: cast-local-type
     model_name = model_name .. " & custom system prompt"
   end
 
@@ -668,7 +640,6 @@ function ChatHandler:_chat_respond(params)
   local llm_prefix = self.options.llm_prefix
   local llm_suffix = "[{{llm}}]"
   local provider = query_prov.name
-  ---@diagnostic disable-next-line: cast-local-type
   llm_suffix = utils.template_render_from_list(llm_suffix, { ["{{llm}}"] = model_name .. " - " .. provider })
 
   for index = start_index, end_index do
@@ -685,10 +656,10 @@ function ChatHandler:_chat_respond(params)
       content = content .. "\n" .. line
     end
   end
-  -- insert last message not handled in loop
+  -- Insert last message not handled in loop
   table.insert(messages, { role = role, content = content })
 
-  -- replace first empty message with system prompt
+  -- Replace first empty message with system prompt
   content = ""
   if headers.system and headers.system:match("%S") then
     content = headers.system
@@ -696,12 +667,12 @@ function ChatHandler:_chat_respond(params)
     content = model_obj.system_prompt
   end
   if content:match("%S") then
-    -- make it multiline again if it contains escaped newlines
+    -- Make it multiline again if it contains escaped newlines
     content = content:gsub("\\n", "\n")
     messages[1] = { role = "system", content = content }
   end
 
-  -- write assistant prompt
+  -- Write assistant prompt
   local last_content_line = utils.last_content_line(buf)
   vim.api.nvim_buf_set_lines(buf, last_content_line, last_content_line, false, { "", llm_prefix .. llm_suffix, "" })
 
@@ -711,7 +682,7 @@ function ChatHandler:_chat_respond(params)
     spinner:start("calling API...")
   end
 
-  -- call the model and write response
+  -- Call the model and write response
   self:query(
     buf,
     query_prov,
@@ -728,7 +699,7 @@ function ChatHandler:_chat_respond(params)
         return
       end
 
-      -- write user prompt
+      -- Write user prompt
       last_content_line = utils.last_content_line(buf)
       utils.undojoin(buf)
       vim.api.nvim_buf_set_lines(
@@ -739,28 +710,28 @@ function ChatHandler:_chat_respond(params)
         { "", "", self.options.chat_user_prefix, "" }
       )
 
-      -- delete whitespace lines at the end of the file
+      -- Delete whitespace lines at the end of the file
       last_content_line = utils.last_content_line(buf)
       utils.undojoin(buf)
       vim.api.nvim_buf_set_lines(buf, last_content_line, -1, false, {})
-      -- insert a new line at the end of the file
+      -- Insert a new line at the end of the file
       utils.undojoin(buf)
       vim.api.nvim_buf_set_lines(buf, -1, -1, false, { "" })
 
-      -- if topic is ?, then generate it
+      -- If topic is ?, then generate it
       if headers.topic == "?" then
-        -- insert last model response
+        -- Insert last model response
         table.insert(messages, { role = "assistant", content = qt.response })
 
         local topic_prov = model_obj.provider
 
-        -- ask model to generate topic/title for the chat
+        -- Ask model to generate topic/title for the chat
         local topic_prompt = self.providers[topic_prov.name].topic_prompt
         if topic_prompt ~= "" then
           table.insert(messages, { role = "user", content = topic_prompt })
         end
 
-        -- prepare invisible buffer for the model to write to
+        -- Prepare invisible buffer for the model to write to
         local topic_buf = vim.api.nvim_create_buf(false, true)
         local topic_resp_handler = ResponseHandler:new(self.queries, topic_buf, nil, 0, false, "", false)
         local topic_handler = topic_resp_handler:create_handler()
@@ -781,7 +752,7 @@ function ChatHandler:_chat_respond(params)
             self.providers[topic_prov.name].topic.params
           ),
         }))
-        -- call the model
+        -- Call the model
         self:query(
           nil,
           topic_prov,
@@ -795,21 +766,21 @@ function ChatHandler:_chat_respond(params)
             if self.options.enable_spinner and topic_spinner then
               topic_spinner:stop()
             end
-            -- get topic from invisible buffer
+            -- Get topic from invisible buffer
             local topic = vim.api.nvim_buf_get_lines(topic_buf, 0, -1, false)[1]
-            -- close invisible buffer
+            -- Close invisible buffer
             vim.api.nvim_buf_delete(topic_buf, { force = true })
-            -- strip whitespace from ends of topic
+            -- Strip whitespace from ends of topic
             topic = topic:gsub("^%s*(.-)%s*$", "%1")
-            -- strip dot from end of topic
+            -- Strip dot from end of topic
             topic = topic:gsub("%.$", "")
 
-            -- if topic is empty do not replace it
+            -- If topic is empty do not replace it
             if topic == "" then
               return
             end
 
-            -- replace topic in current buffer
+            -- Replace topic in current buffer
             utils.undojoin(buf)
             vim.api.nvim_buf_set_lines(buf, 0, 1, false, { "# topic: " .. topic })
           end)
@@ -825,13 +796,15 @@ function ChatHandler:_chat_respond(params)
   )
 end
 
+--- Responds to chat inputs.
+---@param params table Parameters for responding.
 function ChatHandler:chat_respond(params)
   if params.args == "" then
     self:_chat_respond(params)
     return
   end
 
-  -- ensure args is a single positive number
+  -- Ensure args is a single positive number
   local n_requests = tonumber(params.args)
   if n_requests == nil or math.floor(n_requests) ~= n_requests or n_requests <= 0 then
     logger.warning("args for ChatRespond should be a single positive number, not: " .. params.args)
@@ -853,6 +826,7 @@ function ChatHandler:chat_respond(params)
   self:_chat_respond(params)
 end
 
+--- Opens the chat file finder using fzf-lua or native UI.
 function ChatHandler:chat_finder()
   local has_fzf, fzf_lua = pcall(require, "fzf-lua")
 
@@ -866,19 +840,24 @@ function ChatHandler:chat_finder()
       local filename = filename_from_selection(selected)
       return self:open_buf(self.options.chat_dir .. "/" .. filename, ui.BufTarget.popup, self._toggle_kind.chat, false)
     end
-    -- add custom action to delete chat files
-    actions["ctrl-d"] = {
-      fn = function(selected)
-        local filename = filename_from_selection(selected)
-        if vim.fn.confirm("Are you sure you want to delete " .. filename .. "?", "&Yes\n&No", 2) == 1 then
-          futils.delete_file(self.options.chat_dir .. "/" .. filename, self.options.chat_dir)
-          logger.info(filename .. " deleted")
+    -- Add custom action to delete chat files
+    actions["ctrl-d"] = function(selected)
+      local filename = filename_from_selection(selected)
+      if not filename then
+        logger.warning("No valid filename selected for deletion.")
+        return
+      end
+      if vim.fn.confirm(string.format("Are you sure you want to delete '%s'?", filename), "&Yes\n&No", 2) == 1 then
+        local success, err = futils.delete_file(self.options.chat_dir .. "/" .. filename, self.options.chat_dir)
+        if success then
+          logger.info(string.format("Deleted chat file: %s", filename))
+        else
+          logger.error(string.format("Failed to delete '%s': %s", filename, err))
         end
-      end,
-      -- TODO: Fix bug, currently not possible --
-      reload = false,
-    }
+      end
+    end
 
+    -- Set default action based on toggle_target
     if self.options.toggle_target == "popup" then
       actions["default"] = actions["ctrl-p"]
     elseif self.options.toggle_target == "split" then
@@ -901,6 +880,7 @@ function ChatHandler:chat_finder()
     })
     return
   else
+    local scan = require("plenary.scandir") -- Ensure plenary.scandir is available
     local chat_files = scan.scan_dir(self.options.chat_dir, { depth = 1, search_pattern = "%d+%.md$" })
     vim.ui.select(chat_files, {
       prompt = "Select your chat file:",
@@ -924,7 +904,7 @@ function ChatHandler:chat_finder()
       end
       self:open_buf(
         selected_chat,
-        chatutils.resolve_buf_target(self.options.toggle_target),
+        chatutils.resolve_buffer_target(self.options.toggle_target),
         self._toggle_kind.chat,
         false
       )
@@ -932,6 +912,9 @@ function ChatHandler:chat_finder()
   end
 end
 
+--- Switches the current provider.
+---@param selected_prov string Selected provider name.
+---@param is_chat boolean True for chat provider, false for command provider.
 function ChatHandler:switch_provider(selected_prov, is_chat)
   if selected_prov == nil then
     logger.warning("Empty provider selection")
@@ -943,11 +926,19 @@ function ChatHandler:switch_provider(selected_prov, is_chat)
     logger.info("Switched to provider: " .. selected_prov)
     return
   else
-    logger.error("Provider not found: " .. selected_prov)
+    logger.error(
+      string.format(
+        "Provider '%s' not found. Available providers: %s",
+        selected_prov,
+        table.concat(self.available_providers, ", ")
+      )
+    )
     return
   end
 end
 
+--- Handles provider selection via command or UI.
+---@param params table Parameters for provider selection.
 function ChatHandler:provider(params)
   local prov_arg = string.gsub(params.args, "^%s*(.-)%s*$", "%1")
   local has_fzf, fzf_lua = pcall(require, "fzf-lua")
@@ -960,9 +951,11 @@ function ChatHandler:provider(params)
     fzf_lua.fzf_exec(self.available_providers, {
       prompt = "Provider selection ❯",
       fzf_opts = self.options.fzf_lua_opts,
-      complete = function(selection)
-        self:switch_provider(selection[1], is_chat)
-      end,
+      actions = {
+        ["default"] = function(selected)
+          self:switch_provider(selected[1], is_chat)
+        end,
+      },
     })
   else
     vim.ui.select(self.available_providers, {
@@ -973,6 +966,10 @@ function ChatHandler:provider(params)
   end
 end
 
+--- Switches the model for chat or command.
+---@param is_chat boolean True for chat model, false for command model.
+---@param selected_model string Selected model name.
+---@param prov table Provider table.
 function ChatHandler:switch_model(is_chat, selected_model, prov)
   if selected_model == nil then
     logger.warning("Empty model selection")
@@ -986,9 +983,11 @@ function ChatHandler:switch_model(is_chat, selected_model, prov)
     logger.info("Command model: " .. selected_model)
   end
   self.state:refresh(self.available_providers, self.available_models)
-  self:prepare_commands()
+  self.command_manager:prepare_commands() -- Delegate to CommandManager
 end
 
+--- Handles model selection via command or UI.
+---@param params table Parameters for model selection.
 function ChatHandler:model(params)
   local buf = vim.api.nvim_get_current_buf()
   local file_name = vim.api.nvim_buf_get_name(buf)
@@ -1004,14 +1003,16 @@ function ChatHandler:model(params)
     fzf_lua.fzf_exec(prov:get_available_models(fetch_online), {
       prompt = "Model selection ❯",
       fzf_opts = self.options.fzf_lua_opts,
-      complete = function(selection)
-        if #selection == 0 then
-          logger.warning("No model selected")
-          return
-        end
-        local selected_model = selection[1]
-        self:switch_model(is_chat, selected_model, prov)
-      end,
+      actions = {
+        ["default"] = function(selected)
+          if #selected == 0 then
+            logger.warning("No model selected")
+            return
+          end
+          local selected_model = selected[1]
+          self:switch_model(is_chat, selected_model, prov)
+        end,
+      },
     })
   else
     vim.ui.select(prov:get_available_models(fetch_online), {
@@ -1022,6 +1023,8 @@ function ChatHandler:model(params)
   end
 end
 
+--- Retries the last command action.
+---@param params table Parameters for retrying.
 function ChatHandler:retry(params)
   if self.history.last_line1 == nil and self.history.last_line2 == nil then
     return logger.error("No history available to retry: " .. vim.inspect(self.history))
@@ -1040,18 +1043,25 @@ function ChatHandler:retry(params)
   elseif self.history.last_target == ui.Target.prepend then
     template = self.options.template_prepend
   else
-    logger.error("Invalid last target" .. self.history.last_target)
+    logger.error("Invalid last target: " .. tostring(self.history.last_target))
   end
   self:prompt(params, self.history.last_target, model_obj, nil, utils.trim(template), false)
 end
 
+--- Prompts the user and sends the request to the model.
+---@param params table Parameters for prompting.
+---@param target table | number Buffer target.
+---@param model_obj table Model information.
+---@param prompt string Optional prompt for user input.
+---@param template string Template for generating the user prompt.
+---@param reset_history boolean Whether to reset history.
 function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_history)
-  -- enew, new, vnew, tabnew should be resolved into table
+  -- If target is a function, call it to get the actual target
   if type(target) == "function" then
     target = target()
   end
 
-  logger.debug("ChatHandler:prompt - `reset_history`: " .. vim.inspect(reset_history))
+  logger.debug("ChatHandler:prompt - `reset_history`: " .. tostring(reset_history))
   if reset_history == nil or reset_history then
     self.history = {
       last_selection = nil,
@@ -1063,7 +1073,7 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
 
   target = target or ui.Target.enew()
 
-  -- get current buffer
+  -- Get current buffer and window
   local buf = vim.api.nvim_get_current_buf()
   local win = vim.api.nvim_get_current_win()
 
@@ -1072,13 +1082,13 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
     return
   end
 
-  -- defaults to normal mode
+  -- Defaults to normal mode
   local selection = nil
   local prefix = ""
   local start_line = vim.api.nvim_win_get_cursor(0)[1]
   local end_line = start_line
 
-  -- handle range
+  -- Handle range selection
   if params.range == 2 then
     start_line = params.line1
     end_line = params.line2
@@ -1086,13 +1096,13 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
 
     local min_indent = nil
     local use_tabs = false
-    -- measure minimal common indentation for lines with content
+    -- Measure minimal common indentation for lines with content
     for i, line in ipairs(lines) do
       lines[i] = line
-      -- skip whitespace only lines
+      -- Skip whitespace-only lines
       if not line:match("^%s*$") then
         local indent = line:match("^%s*")
-        -- contains tabs
+        -- Check if indentation uses tabs
         if indent:match("\t") then
           use_tabs = true
         end
@@ -1130,15 +1140,15 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
       logger.debug("LAST COMMAND in use " .. self.history.last_command)
       command = self.history.last_command
     end
-    -- dummy handler
+    -- Dummy handler
     local handler = function() end
-    -- default on_exit strips trailing backticks if response was markdown snippet
+    -- Default on_exit strips trailing backticks if response was markdown snippet
     local on_exit = function(qid)
       local qt = self.queries:get(qid)
       if not qt then
         return
       end
-      -- if buf is not valid, return
+      -- If buffer is not valid, return
       if not vim.api.nvim_buf_is_valid(buf) then
         return
       end
@@ -1146,9 +1156,9 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
       local flc, llc
       local fl = qt.first_line
       local ll = qt.last_line
-      -- remove empty lines from the start and end of the response
+      -- Remove empty lines from the start and end of the response
       while true do
-        -- get content of first_line and last_line
+        -- Get content of first_line and last_line
         flc = vim.api.nvim_buf_get_lines(buf, fl, fl + 1, false)[1]
         llc = vim.api.nvim_buf_get_lines(buf, ll, ll + 1, false)[1]
 
@@ -1159,12 +1169,12 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
         local flm = flc:match("%S")
         local llm = llc:match("%S")
 
-        -- break loop if both lines contain non-whitespace characters
+        -- Break loop if both lines contain non-whitespace characters
         if flm and llm then
           break
         end
 
-        -- break loop lines are equal
+        -- Break loop if lines are equal
         if fl >= ll then
           break
         end
@@ -1179,12 +1189,12 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
         ll = ll - 1
       end
 
-      -- if fl and ll starts with triple backticks, remove these lines
+      -- If fl and ll start with triple backticks, remove these lines
       if flc and llc and flc:match("^%s*```") and llc:match("^%s*```") then
-        -- remove first line with undojoin
+        -- Remove first line with undojoin
         utils.undojoin(buf)
         vim.api.nvim_buf_set_lines(buf, fl, fl + 1, false, {})
-        -- remove last line
+        -- Remove last line
         utils.undojoin(buf)
         vim.api.nvim_buf_set_lines(buf, ll - 1, ll, false, {})
         ll = ll - 2
@@ -1192,17 +1202,17 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
       qt.first_line = fl
       qt.last_line = ll
 
-      -- option to not select response automatically
+      -- Option to not select response automatically
       if not self.options.command_auto_select_response then
         return
       end
 
-      -- don't select popup response
+      -- Don't select popup response
       if target == ui.Target.popup then
         return
       end
 
-      -- default works for rewrite and enew
+      -- Default works for rewrite and enew
       local start = fl
       local finish = ll
 
@@ -1214,13 +1224,13 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
         finish = self._selection_last_line + ll - fl
       end
 
-      -- select from first_line to last_line
+      -- Select from first_line to last_line
       vim.api.nvim_win_set_cursor(0, { start + 1, 0 })
       vim.api.nvim_command("normal! V")
       vim.api.nvim_win_set_cursor(0, { finish + 1, 0 })
     end
 
-    -- prepare messages
+    -- Prepare messages
     local messages = {}
     local filetype = pft.detect(vim.api.nvim_buf_get_name(buf), {})
     local filename = vim.api.nvim_buf_get_name(buf)
@@ -1232,7 +1242,7 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
     if sys_prompt ~= "" then
       local repo_instructions = futils.find_repo_instructions()
       if repo_instructions ~= "" and sys_prompt ~= "" then
-        -- append the repository instructions from .parrot.md to the system prompt
+        -- Append the repository instructions from .parrot.md to the system prompt
         sys_prompt = sys_prompt .. "\n" .. repo_instructions
       end
       table.insert(messages, { role = "system", content = sys_prompt })
@@ -1245,7 +1255,7 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
     table.insert(messages, { role = "user", content = user_prompt })
     logger.debug("ChatHandler:prompt - `user_prompt`: " .. user_prompt)
 
-    -- cancel possible visual mode before calling the model
+    -- Cancel possible visual mode before calling the model
     utils.feedkeys("<esc>", "xn")
 
     local cursor = true
@@ -1253,89 +1263,10 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
       cursor = false
     end
 
-    -- mode specific logic
-    if target == ui.Target.rewrite then
-      -- delete selection
-      vim.api.nvim_buf_set_lines(buf, start_line - 1, end_line - 1, false, {})
-      -- prepare handler
-      handler = ResponseHandler:new(self.queries, buf, win, start_line - 1, true, prefix, cursor):create_handler()
-    elseif target == ui.Target.append then
-      -- move cursor to the end of the selection
-      vim.api.nvim_win_set_cursor(0, { end_line, 0 })
-      -- put newline after selection
-      vim.api.nvim_put({ "" }, "l", true, true)
-      -- prepare handler
-      handler = ResponseHandler:new(self.queries, buf, win, end_line, true, prefix, cursor):create_handler()
-    elseif target == ui.Target.prepend then
-      -- move cursor to the start of the selection
-      vim.api.nvim_win_set_cursor(0, { start_line, 0 })
-      -- put newline before selection
-      vim.api.nvim_put({ "" }, "l", false, true)
-      -- prepare handler
-      handler = ResponseHandler:new(self.queries, buf, win, start_line - 1, true, prefix, cursor):create_handler()
-    elseif target == ui.Target.popup then
-      self:toggle_close(self._toggle_kind.popup)
-      -- create a new buffer
-      local popup_close = nil
-      buf, win, popup_close, _ = ui.create_popup(
-        nil,
-        self._plugin_name .. " popup (close with <esc>/<C-c>)",
-        function(w, h)
-          local top = self.options.style_popup_margin_top or 2
-          local bottom = self.options.style_popup_margin_bottom or 8
-          local left = self.options.style_popup_margin_left or 1
-          local right = self.options.style_popup_margin_right or 1
-          local max_width = self.options.style_popup_max_width or 160
-          local ww = math.min(w - (left + right), max_width)
-          local wh = h - (top + bottom)
-          return ww, wh, top, (w - ww) / 2
-        end,
-        { on_leave = true, escape = true },
-        { border = self.options.style_popup_border or "single" }
-      )
-      -- set the created buffer as the current buffer
-      vim.api.nvim_set_current_buf(buf)
-      -- set the filetype to markdown
-      vim.api.nvim_set_option_value("filetype", "markdown", { buf = buf })
-      -- better text wrapping
-      vim.api.nvim_command("setlocal wrap linebreak")
-      -- prepare handler
-      handler = ResponseHandler:new(self.queries, buf, win, 0, false, "", false):create_handler()
-      self:toggle_add(self._toggle_kind.popup, { win = win, buf = buf, close = popup_close })
-    elseif type(target) == "table" then
-      if target.type == ui.Target.new().type then
-        vim.cmd("split")
-        win = vim.api.nvim_get_current_win()
-      elseif target.type == ui.Target.vnew().type then
-        vim.cmd("vsplit")
-        win = vim.api.nvim_get_current_win()
-      elseif target.type == ui.Target.tabnew().type then
-        vim.cmd("tabnew")
-        win = vim.api.nvim_get_current_win()
-      end
-
-      buf = vim.api.nvim_create_buf(true, true)
-      vim.api.nvim_set_current_buf(buf)
-
-      local group = utils.create_augroup("PrtScratchSave" .. utils.uuid(), { clear = true })
-      vim.api.nvim_create_autocmd({ "BufWritePre" }, {
-        buffer = buf,
-        group = group,
-        callback = function(ctx)
-          vim.api.nvim_set_option_value("buftype", "", { buf = ctx.buf })
-          vim.api.nvim_buf_set_name(ctx.buf, ctx.file)
-          vim.api.nvim_command("w!")
-          vim.api.nvim_del_augroup_by_id(ctx.group)
-        end,
-      })
-
-      local ft = target.filetype or filetype
-      vim.api.nvim_set_option_value("filetype", ft, { buf = buf })
-
-      handler = ResponseHandler:new(self.queries, buf, win, 0, false, "", cursor):create_handler()
-    end
+    -- Mode-specific logic delegated to CommandManager
+    self.command_manager:handle_command(target, buf, win, start_line, end_line, prefix, cursor, messages)
 
-    -- call the model and write the response
+    -- Call the model and write the response
     prov:set_model(model_obj.name)
 
     local spinner = nil
@@ -1365,7 +1296,7 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
       return
     end
 
-    -- if prompt is not provided, run the command directly
+    -- If prompt is not provided, run the command directly
     if not prompt or prompt == "" then
       callback(nil)
       return
@@ -1381,17 +1312,17 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h
         callback(input)
       end)
     else
-      logger.error("Invalid user input ui option: " .. self.options.user_input_ui)
+      logger.error("Invalid user input UI option: " .. self.options.user_input_ui)
     end
   end)
 end
 
--- call the API
----@param buf number | nil # buffer number
----@param provider table
----@param payload table # payload for api
----@param handler function # response handler
----@param on_exit function | nil # optional on_exit handler
+--- Sends a query to the provider's API.
+---@param buf number | nil Buffer number.
+---@param provider table Provider information.
+---@param payload table Payload for the API.
+---@param handler function Response handler function.
+---@param on_exit function | nil Optional on_exit handler.
 function ChatHandler:query(buf, provider, payload, handler, on_exit)
   if type(handler) ~= "function" then
     logger.error(
@@ -1491,8 +1422,25 @@ function ChatHandler:query(buf, provider, payload, handler, on_exit)
         end
       end
     end,
+    on_stderr = function(_, data)
+      if data and #data > 0 then
+        logger.error("Curl stderr: " .. table.concat(data, "\n"))
+      end
+    end,
   })
-  job:start()
+
+  -- Handle job start failure
+  local ok, err = pcall(function()
+    job:start()
+  end)
+  if not ok then
+    logger.error(string.format("Failed to start curl job: %s", err))
+    if on_exit then
+      on_exit(qid)
+    end
+    return
+  end
+
   self.pool:add(job, buf)
 end
 
diff --git a/lua/parrot/command_manager.lua b/lua/parrot/command_manager.lua
new file mode 100644
index 0000000..5a434f3
--- /dev/null
+++ b/lua/parrot/command_manager.lua
@@ -0,0 +1,240 @@
+local utils = require("parrot.utils")
+local logger = require("parrot.logger")
+local ui = require("parrot.ui")
+
+--- CommandManager Module
+-- Handles the registration and setup of commands and keybindings for the parrot.nvim plugin.
+local CommandManager = {}
+CommandManager.__index = CommandManager
+
+--- Creates a new CommandManager instance.
+-- @param options table Configuration options.
+-- @param chat_handler ChatHandler instance to delegate actions.
+-- @return CommandManager
+function CommandManager:new(options, chat_handler)
+  local instance = setmetatable({
+    options = options,
+    chat_handler = chat_handler,
+    commands = {}, -- Stores command definitions
+  }, CommandManager)
+  return instance
+end
+
+--- Prepares and registers all commands by iterating over ui.Target.
+function CommandManager:prepare_commands()
+  for name, target in pairs(ui.Target) do
+    -- Uppercase first letter to form the command name
+    local command = name:gsub("^%l", string.upper)
+
+    -- Determine the appropriate model based on the target
+    local model_obj = self.chat_handler:get_model("command")
+    if target == ui.Target.popup then
+      model_obj = self.chat_handler:get_model("chat")
+    end
+
+    -- Define the command callback function
+    local cmd = function(params)
+      -- Choose the template based on the range and target
+      local template = self.options.template_command
+      if params.range == 2 then
+        template = self.options.template_selection
+        -- Select custom templates for specific targets
+        if target == ui.Target.rewrite then
+          template = self.options.template_rewrite
+        elseif target == ui.Target.append then
+          template = self.options.template_append
+        elseif target == ui.Target.prepend then
+          template = self.options.template_prepend
+        end
+      end
+
+      -- Render the command prefix using templates
+      local cmd_prefix = utils.template_render_from_list(
+        self.options.command_prompt_prefix_template,
+        { ["{{llm}}"] = model_obj.name }
+      )
+
+      -- Delegate the prompt handling to ChatHandler
+      self.chat_handler:prompt(params, target, model_obj, cmd_prefix, utils.trim(template), true)
+    end
+
+    -- Store the command in the commands table
+    self.commands[command] = command
+
+    -- Register the command with Neovim and set up keybindings
+    self:add_command(command, function(params)
+      cmd(params)
+    end)
+  end
+end
+
+--- Adds and registers a single command.
+-- @param command string Command name.
+-- @param cmd function Command callback function.
+function CommandManager:add_command(command, cmd)
+  if not command or not cmd then
+    logger.error("CommandManager:add_command requires both command name and callback function.")
+    return
+  end
+
+  -- Register the user command with Neovim
+  vim.api.nvim_create_user_command(command, function(params)
+    cmd(params)
+  end, { nargs = "*", range = true })
+
+  logger.debug(string.format("Registered command: %s", command))
+
+  -- Extract the base name without "Chat" prefix if present for keybinding configuration
+  local base_command = command:gsub("^Chat", "")
+  local shortcut_key = self.options["chat_shortcut_" .. base_command:lower()]
+
+  if shortcut_key and shortcut_key.shortcut then
+    self:setup_keybinding(command, shortcut_key)
+  end
+end
+
+--- Sets up keybindings for a specific command.
+-- @param command string Command name.
+-- @param keybind table Keybinding definition containing modes, shortcut, and description.
+function CommandManager:setup_keybinding(command, keybind)
+  if not keybind or not keybind.modes or not keybind.shortcut then
+    logger.error(string.format("Invalid keybind configuration for command: %s", command))
+    return
+  end
+
+  for _, mode in ipairs(keybind.modes) do
+    -- Define a keymap callback to execute the command
+    local callback = function()
+      vim.api.nvim_command(command .. "<CR>")
+      -- Exit insert mode if in insert mode
+      if mode == "i" then
+        vim.api.nvim_command("stopinsert")
+      end
+    end
+
+    -- Set the keymap with buffer-local scope if specified
+    local opts = { noremap = true, silent = true, desc = keybind.desc or ("Execute " .. command) }
+    if keybind.buffer then
+      vim.api.nvim_buf_set_keymap(keybind.buffer, mode, keybind.shortcut, "", { callback = callback, noremap = opts.noremap, silent = opts.silent, desc = opts.desc })
+    else
+      vim.api.nvim_set_keymap(mode, keybind.shortcut, "", { callback = callback, noremap = opts.noremap, silent = opts.silent, desc = opts.desc })
+    end
+
+    logger.debug(string.format("Set keybinding for command '%s' in mode '%s' with shortcut '%s'", command, mode, keybind.shortcut))
+  end
+end
+
+--- Cleans up all registered commands and keybindings.
+function CommandManager:cleanup()
+  for command_name, _ in pairs(self.commands) do
+    pcall(vim.api.nvim_del_user_command, command_name)
+    logger.info(string.format("Unregistered command: %s", command_name))
+    -- Note: Neovim does not provide a direct way to remove keybindings programmatically.
+    -- To remove keybindings, you would need to track them and manually unset them if necessary.
+  end
+end
+--- Handles the execution of a command.
+---@param params table Parameters passed from the command invocation.
+---@param target ui.Target The target UI component.
+---@param model_obj Model The language model to use.
+---@param cmd_prefix string The command prompt prefix.
+---@param template string The template to use for the prompt.
+function CommandManager:handle_command(params, target, model_obj, cmd_prefix, template)
+  -- Mode specific logic
+  local handler
+  local buf = vim.api.nvim_get_current_buf()
+  local win = vim.api.nvim_get_current_win()
+  local start_line = params.line1
+  local end_line = params.line2
+  local cursor = vim.api.nvim_win_get_cursor(0)
+  local prefix = cmd_prefix
+
+  if target == ui.Target.rewrite then
+    -- Delete selection
+    vim.api.nvim_buf_set_lines(buf, start_line - 1, end_line, false, {})
+    -- Prepare handler
+    handler = ResponseHandler:new(self.chat_handler.queries, buf, win, start_line - 1, true, prefix, cursor):create_handler()
+  elseif target == ui.Target.append then
+    -- Move cursor to the end of the selection
+    vim.api.nvim_win_set_cursor(0, { end_line, 0 })
+    -- Put newline after selection
+    vim.api.nvim_put({ "" }, "l", true, true)
+    -- Prepare handler
+    handler = ResponseHandler:new(self.chat_handler.queries, buf, win, end_line, true, prefix, cursor):create_handler()
+  elseif target == ui.Target.prepend then
+    -- Move cursor to the start of the selection
+    vim.api.nvim_win_set_cursor(0, { start_line, 0 })
+    -- Put newline before selection
+    vim.api.nvim_put({ "" }, "l", false, true)
+    -- Prepare handler
+    handler = ResponseHandler:new(self.chat_handler.queries, buf, win, start_line - 1, true, prefix, cursor):create_handler()
+  elseif target == ui.Target.popup then
+    self.chat_handler:toggle_close(self.chat_handler._toggle_kind.popup)
+    -- Create a new buffer
+    local popup_close = nil
+    buf, win, popup_close, _ = ui.create_popup(
+      nil,
+      self.chat_handler._plugin_name .. " popup (close with <esc>/<C-c>)",
+      function(w, h)
+        local top = self.options.style_popup_margin_top or 2
+        local bottom = self.options.style_popup_margin_bottom or 8
+        local left = self.options.style_popup_margin_left or 1
+        local right = self.options.style_popup_margin_right or 1
+        local max_width = self.options.style_popup_max_width or 160
+        local ww = math.min(w - (left + right), max_width)
+        local wh = h - (top + bottom)
+        return ww, wh, top, (w - ww) / 2
+      end,
+      { on_leave = true, escape = true },
+      { border = self.options.style_popup_border or "single" }
+    )
+    -- Set the created buffer as the current buffer
+    vim.api.nvim_set_current_buf(buf)
+    -- Set the filetype to markdown
+    vim.api.nvim_set_option_value("filetype", "markdown", { buf = buf })
+    -- Better text wrapping
+    vim.api.nvim_command("setlocal wrap linebreak")
+    -- Prepare handler
+    handler = ResponseHandler:new(self.chat_handler.queries, buf, win, 0, false, "", false):create_handler()
+    self.chat_handler:toggle_add(self.chat_handler._toggle_kind.popup, { win = win, buf = buf, close = popup_close })
+  elseif type(target) == "table" then
+    if target.type == ui.Target.new().type then
+      vim.cmd("split")
+      win = vim.api.nvim_get_current_win()
+    elseif target.type == ui.Target.vnew().type then
+      vim.cmd("vsplit")
+      win = vim.api.nvim_get_current_win()
+    elseif target.type == ui.Target.tabnew().type then
+      vim.cmd("tabnew")
+      win = vim.api.nvim_get_current_win()
+    end
+
+    buf = vim.api.nvim_create_buf(true, true)
+    vim.api.nvim_set_current_buf(buf)
+
+    local group = utils.create_augroup("PrtScratchSave" .. utils.uuid(), { clear = true })
+    vim.api.nvim_create_autocmd({ "BufWritePre" }, {
+      buffer = buf,
+      group = group,
+      callback = function(ctx)
+        vim.api.nvim_set_option_value("buftype", "", { buf = ctx.buf })
+        vim.api.nvim_buf_set_name(ctx.buf, ctx.file)
+        vim.api.nvim_command("w!")
+        vim.api.nvim_del_augroup_by_id(ctx.group)
+      end,
+    })
+
+    local ft = target.filetype or vim.api.nvim_buf_get_option(0, 'filetype')
+    vim.api.nvim_set_option_value("filetype", ft, { buf = buf })
+
+    handler = ResponseHandler:new(self.chat_handler.queries, buf, win, 0, false, "", cursor):create_handler()
+  else
+    -- Default handler
+    handler = ResponseHandler:new(self.chat_handler.queries, buf, win, 0, true, prefix, cursor):create_handler()
+  end
+
+  -- Delegate the prompt handling to ChatHandler with the prepared handler
+  self.chat_handler:prompt(params, target, model_obj, cmd_prefix, utils.trim(template), true, handler)
+end
+
+return CommandManager
diff --git a/lua/parrot/toggle_manager.lua b/lua/parrot/toggle_manager.lua
new file mode 100644
index 0000000..e16e01e
--- /dev/null
+++ b/lua/parrot/toggle_manager.lua
@@ -0,0 +1,46 @@
+local logger = require("parrot.logger")
+
+local TOGGLE_KIND = {
+  UNKNOWN = 0,
+  CHAT = 1,
+  POPUP = 2,
+  CONTEXT = 3,
+}
+
+local ToggleManager = {}
+ToggleManager.__index = ToggleManager
+
+function ToggleManager:new()
+  return setmetatable({ toggles = {} }, self)
+end
+
+function ToggleManager:close(kind)
+  local toggle = self.toggles[kind]
+  if toggle and vim.api.nvim_win_is_valid(toggle.win) and vim.api.nvim_buf_is_valid(toggle.buf) then
+    if #vim.api.nvim_list_wins() > 1 then
+      toggle.close()
+      self.toggles[kind] = nil
+      return true
+    else
+      logger.warning("Can't close the last window.")
+    end
+  end
+  self.toggles[kind] = nil
+  return false
+end
+
+function ToggleManager:add(kind, toggle)
+  self.toggles[kind] = toggle
+end
+
+function ToggleManager:resolve(kind_str)
+  kind_str = kind_str:lower()
+  local kind_map = {
+    chat = TOGGLE_KIND.CHAT,
+    popup = TOGGLE_KIND.POPUP,
+    context = TOGGLE_KIND.CONTEXT,
+  }
+  return kind_map[kind_str] or TOGGLE_KIND.UNKNOWN
+end
+
+return ToggleManager
diff --git a/lua/parrot/utils.lua b/lua/parrot/utils.lua
index dca17dd..3cf4753 100644
--- a/lua/parrot/utils.lua
+++ b/lua/parrot/utils.lua
@@ -25,14 +25,14 @@ end
 ---@param callback function|string Callback or string to set keymap
 ---@param desc string|nil Optional description for keymap
 function M.set_keymap(buffers, mode, key, callback, desc)
-  local opts = {
-    noremap = true,
-    silent = true,
-    nowait = true,
-    desc = desc,
-  }
   for _, buf in ipairs(buffers) do
-    opts.buffer = buf
+    local opts = {
+      noremap = true,
+      silent = true,
+      nowait = true,
+      buffer = buf,
+      desc = desc,
+    }
     vim.keymap.set(mode, key, callback, opts)
   end
 end
@@ -43,16 +43,20 @@ end
 ---@param callback function Callback to call
 ---@param gid number Augroup id
 function M.autocmd(events, buffers, callback, gid)
-  local opts = {
-    group = gid,
-    callback = vim.schedule_wrap(callback),
-  }
   if buffers then
     for _, buf in ipairs(buffers) do
-      opts.buffer = buf
+      local opts = {
+        group = gid,
+        buffer = buf,
+        callback = vim.schedule_wrap(callback),
+      }
       vim.api.nvim_create_autocmd(events, opts)
     end
   else
+    local opts = {
+      group = gid,
+      callback = vim.schedule_wrap(callback),
+    }
     vim.api.nvim_create_autocmd(events, opts)
   end
 end