init.lua

local async = require("plenary.async")
local completion = require("codecompanion.providers.completion")
local config = require("codecompanion.config")
local ts = require("codecompanion.utils.treesitter")
local ui_utils = require("codecompanion.utils.ui")
local utils = require("codecompanion.utils")

local api = vim.api

local M = {}

-- CHAT MAPPINGS --------------------------------------------------------------
local _cached_options = {}
M.options = {
  callback = function()
    local float_opts = {
      filetype = "codecompanion",
      lock = true,
      style = "minimal",
      title = "Options",
      window = config.display.chat.window,
    }

    if next(_cached_options) ~= nil then
      return ui_utils.create_float(_cached_options, float_opts)
    end

    local lines = {}
    local indent = " "

    local function max(col, tbl)
      local max_length = 0
      for key, val in pairs(tbl) do
        if val.hide then
          goto continue
        end

        local get_length = (col == "key") and key or val[col]

        local length = #get_length
        if length > max_length then
          max_length = length
        end

        ::continue::
      end
      return max_length
    end

    local function pad(str, max_length, offset)
      return str .. string.rep(" ", max_length - #str + (offset or 0))
    end

    --- Cleans and truncates a string to a maximum width.
    ---@param desc string? The description to clean
    ---@param max_width number? The maximum width to truncate the description to
    ---@return string The cleaned and truncated description
    local function clean_and_truncate(desc, max_width)
      if not desc then
        return ""
      end
      desc = vim.trim(tostring(desc):gsub("\n", " "))
      if max_width and #desc > max_width then
        return desc:sub(1, max_width - 3) .. "..."
      end
      return desc
    end

    local function sorted_pairs(tbl, comp)
      local keys = {}
      for k in pairs(tbl) do
        table.insert(keys, k)
      end
      table.sort(keys, comp)
      local i = 0
      return function()
        i = i + 1
        local key = keys[i]
        if key ~= nil then
          return key, tbl[key]
        end
      end
    end

    -- Filter out private keymaps
    local keymaps = {}
    for k, v in pairs(config.interactions.chat.keymaps) do
      if k:sub(1, 1) ~= "_" then
        keymaps[k] = v
      end
    end
    -- Workout the column spacing
    local keymaps_max = max("description", keymaps)

    local vars = {}
    vim.iter(config.interactions.chat.variables):each(function(key, val)
      if not val.hide_in_help_window then
        vars[key] = val
      end
    end)
    local vars_max = max("key", vars)

    local tools = {}
    -- Add tools
    vim
      .iter(config.interactions.chat.tools)
      :filter(function(name)
        return name ~= "opts" and name ~= "groups"
      end)
      :each(function(tool)
        local tool_conf = config.interactions.chat.tools[tool]
        if not tool_conf.hide_in_help_window then
          tools[tool] = {
            description = tool_conf.description,
          }
        end
      end)

    -- Add groups
    vim.iter(config.interactions.chat.tools.groups):each(function(tool)
      local group_conf = config.interactions.chat.tools.groups[tool]
      if not group_conf.hide_in_help_window then
        tools[tool] = {
          description = group_conf.description,
        }
      end
    end)

    local tools_max = max("key", tools)

    local max_length = math.max(keymaps_max, vars_max, tools_max)

    -- Keymaps
    table.insert(lines, "### Keymaps")

    local function compare_keymaps(a, b)
      return (keymaps[a].description or "") < (keymaps[b].description or "")
    end

    for _, map in sorted_pairs(keymaps, compare_keymaps) do
      if type(map.condition) == "function" and not map.condition() then
        goto continue
      end
      if not map.hide then
        local modes = {
          n = "Normal",
          i = "Insert",
        }

        local output = {}
        for mode, key in pairs(map.modes) do
          if type(key) == "table" then
            local keys = {}
            for _, v in ipairs(key) do
              table.insert(keys, "`" .. v .. "`")
            end
            key = table.concat(key, "|")
            table.insert(output, "`" .. key .. "` in " .. modes[mode] .. " mode")
          else
            table.insert(output, "`" .. key .. "` in " .. modes[mode] .. " mode")
          end
        end
        local output_str = table.concat(output, " and ")

        table.insert(lines, indent .. pad("_" .. map.description .. "_", max_length, 4) .. " " .. output_str)
      end
      ::continue::
    end

    -- Variables
    table.insert(lines, "")
    table.insert(lines, "### Variables")

    for key, val in sorted_pairs(vars) do
      local desc = clean_and_truncate(val.description)
      table.insert(lines, indent .. pad("#{" .. key .. "}", max_length, 4) .. " " .. desc)
    end

    -- Tools
    table.insert(lines, "")
    table.insert(lines, "### Tools")

    for key, val in sorted_pairs(tools) do
      if key ~= "opts" then
        local desc = clean_and_truncate(val.description)
        table.insert(lines, indent .. pad("@{" .. key .. "}", max_length, 4) .. " " .. desc)
      end
    end

    _cached_options = lines
    ui_utils.create_float(lines, float_opts)
  end,
}

-- Native completion
M.completion = {
  callback = function(chat)
    local function complete_items(callback)
      async.run(function()
        local slash_cmds = completion.slash_commands()
        local tools = completion.tools()
        local vars = completion.variables()

        local items = {}

        if type(slash_cmds[1]) == "table" then
          vim.list_extend(items, slash_cmds)
        end
        if type(tools[1]) == "table" then
          vim.list_extend(items, tools)
        end
        if type(vars[1]) == "table" then
          vim.list_extend(items, vars)
        end

        -- Process each item to match the completion format
        for _, item in ipairs(items) do
          if item.label then
            -- Add bracket wrapping for variables and tools like cmp/blink do
            if item.type == "variable" then
              item.word = string.format("#{%s}", item.label:sub(2))
            elseif item.type == "tool" then
              item.word = string.format("@{%s}", item.label:sub(2))
            else
              item.word = item.label
            end

            item.abbr = item.label:sub(2)
            item.menu = item.description or item.detail
            item.icase = 1
            item.dup = 0
            item.empty = 0
            item.user_data = {
              command = item.label:sub(2),
              label = item.label,
              type = item.type,
              config = item.config,
              from_prompt_library = item.from_prompt_library,
            }
          end
        end

        vim.schedule(function()
          callback(items)
        end)
      end)
    end

    local function trigger_complete()
      local line = vim.api.nvim_get_current_line()
      local cursor = vim.api.nvim_win_get_cursor(0)
      local col = cursor[2]
      if col == 0 or #line == 0 then
        return
      end

      local before_cursor = line:sub(1, col)
      local find_current_word = string.find(before_cursor, "%s[^%s]*$")
      local start = find_current_word or 0
      local prefix = line:sub(start + 1, col)
      if not prefix then
        return
      end

      complete_items(function(items)
        vim.fn.complete(
          start + 1,
          vim.tbl_filter(function(item)
            return vim.startswith(item.label:lower(), prefix:lower())
          end, items)
        )
      end)
    end

    trigger_complete()
  end,
}

M.send = {
  callback = function(chat)
    chat:submit()
  end,
}

M.regenerate = {
  callback = function(chat)
    chat:regenerate()
  end,
}

M.close = {
  callback = function(chat)
    chat:close()

    local chats = require("codecompanion").buf_get_chat()
    if vim.tbl_count(chats) == 0 then
      return
    end

    local window_opts = chat.ui.window_opts or { default = true }
    chats[1].chat.ui:open({ window_opts = window_opts })
  end,
}

M.stop = {
  callback = function(chat)
    if chat.current_request or chat.tool_orchestrator then
      chat:stop()
    end
  end,
}

M.clear = {
  callback = function(chat)
    chat:clear()
  end,
}

M.codeblock = {
  desc = "Insert a codeblock",
  callback = function(chat)
    local bufnr = api.nvim_get_current_buf()
    local cursor_pos = api.nvim_win_get_cursor(0)
    local line = cursor_pos[1]

    local ft = chat.buffer_context.filetype or ""

    local codeblock = {
      "````" .. ft,
      "",
      "````",
    }

    api.nvim_buf_set_lines(bufnr, line - 1, line, false, codeblock)
    api.nvim_win_set_cursor(0, { line + 1, vim.fn.indent(line) })
  end,
}

---@param node TSNode to yank text from
local function yank_node(node)
  local start_row, start_col, end_row, end_col = node:range()
  local cursor_position = vim.fn.getcurpos()

  -- Create marks for the node range
  local ok, _ = pcall(function()
    vim.api.nvim_buf_set_mark(0, "[", start_row + 1, start_col, {})
    vim.api.nvim_buf_set_mark(0, "]", end_row + 1, end_col - 1, {})
  end)

  if not ok then
    return utils.notify("Failed to copy code block", vim.log.levels.WARN)
  end

  -- Yank using marks
  vim.cmd(string.format('normal! `["%sy`]', config.interactions.chat.opts.register))

  -- Restore position after delay
  vim.defer_fn(function()
    vim.fn.setpos(".", cursor_position)
  end, config.interactions.chat.opts.yank_jump_delay_ms)
end

M.yank_code = {
  desc = "Yank focused or the last codeblock",
  callback = function(chat)
    local node = chat:get_codeblock()
    if node ~= nil then
      yank_node(node)
    end
  end,
}

M.buffer_sync_all = {
  desc = "Sync the buffer to share all of its content",
  callback = function(chat)
    local current_line = vim.api.nvim_win_get_cursor(0)[1]
    local line = vim.api.nvim_buf_get_lines(chat.bufnr, current_line - 1, current_line, true)[1]

    if not vim.startswith(line, "> - ") then
      return
    end

    local icon = config.display.chat.icons.buffer_sync_all
    local id = line:gsub("^> %- ", "")

    if not chat.context:can_be_synced__all(id) then
      return utils.notify("This context type cannot be sync'd", vim.log.levels.WARN)
    end

    local filename = id
    local state = "unsynced"
    if line:find(icon) then
      state = "synced"
      filename = filename:gsub(icon, "")
      id = filename
    end

    -- Update the UI
    local new_line = (state == "synced") and string.format("> - %s", filename)
      or string.format("> - %s%s", icon, filename)
    api.nvim_buf_set_lines(chat.bufnr, current_line - 1, current_line, true, { new_line })

    -- Update the context items on the chat buffer
    for _, item in ipairs(chat.context_items) do
      if item.id == id then
        item.opts.synced = not item.opts.synced
        break
      end
    end
  end,
}

M.buffer_sync_diff = {
  desc = "Sync the buffer to share it's diffs",
  callback = function(chat)
    local current_line = vim.api.nvim_win_get_cursor(0)[1]
    local line = vim.api.nvim_buf_get_lines(chat.bufnr, current_line - 1, current_line, true)[1]

    if not vim.startswith(line, "> - ") then
      return
    end

    local id = line:gsub("^> %- ", "")
    if not chat.context:can_be_synced__diff(id) then
      return utils.notify("This context type cannot be sync'd", vim.log.levels.WARN)
    end

    -- Find the context and toggle diff state
    local icons = config.display.chat.icons
    for _, item in ipairs(chat.context_items) do
      local clean_id = id:gsub(icons.buffer_sync_all, ""):gsub(icons.buffer_sync_diff, "")
      if item.id == clean_id then
        if not item.opts then
          item.opts = {}
        end
        item.opts.sync_diff = not item.opts.sync_diff

        -- Update the UI for just this line
        local new_line
        if item.opts.sync_diff then
          -- Check if buffer is still valid before syncing
          if vim.api.nvim_buf_is_valid(item.bufnr) and vim.api.nvim_buf_is_loaded(item.bufnr) then
            chat.buffer_diffs:sync(item.bufnr)
            new_line = string.format("> - %s%s", icons.buffer_sync_diff, clean_id)
          else
            -- Buffer is invalid, can't sync with it
            item.opts.sync_diff = false
            new_line = string.format("> - %s", clean_id)
            utils.notify("Cannot sync - Invalid or unloaded buffer " .. item.id, vim.log.levels.WARN)
          end
        else
          chat.buffer_diffs:unsync(item.bufnr)
          new_line = string.format("> - %s", clean_id)
        end

        -- Update only the current line
        vim.api.nvim_buf_set_lines(chat.bufnr, current_line - 1, current_line, true, { new_line })
        break
      end
    end
  end,
}

---@param chat CodeCompanion.Chat
---@param direction number
local function move_buffer(chat, direction)
  local bufs = _G.codecompanion_buffers
  local len = #bufs
  local next_buf = vim
    .iter(bufs)
    :enumerate()
    :filter(function(_, v)
      return v == chat.bufnr
    end)
    :map(function(i, _)
      return direction > 0 and bufs[(i % len) + 1] or bufs[((i - 2 + len) % len) + 1]
    end)
    :next()

  local codecompanion = require("codecompanion")

  local prev_ui = codecompanion.buf_get_chat(chat.bufnr).ui
  prev_ui:hide()
  local window_opts = prev_ui.window_opts or { default = true }
  codecompanion.buf_get_chat(next_buf).ui:open({ window_opts = window_opts })
end

M.next_chat = {
  desc = "Move to the next chat",
  callback = function(chat)
    if vim.tbl_count(_G.codecompanion_buffers) == 1 then
      return
    end
    move_buffer(chat, 1)
  end,
}

M.previous_chat = {
  desc = "Move to the previous chat",
  callback = function(chat)
    if vim.tbl_count(_G.codecompanion_buffers) == 1 then
      return
    end
    move_buffer(chat, -1)
  end,
}

M.next_header = {
  desc = "Go to the next message",
  callback = function()
    ts.goto_heading("next", 1)
  end,
}

M.previous_header = {
  desc = "Go to the previous message",
  callback = function()
    ts.goto_heading("prev", 1)
  end,
}

M.change_adapter = {
  desc = "Change the adapter",
  callback = function(chat)
    return require("codecompanion.interactions.chat.keymaps.change_adapter").callback(chat)
  end,
}

M.fold_code = {
  callback = function(chat)
    chat.ui:fold_code()
  end,
}

M.debug = {
  desc = "Show debug information for the current chat",
  callback = function(chat)
    local settings, messages = chat:debug()
    if not settings and not messages then
      return
    end

    return require("codecompanion.interactions.chat.debug")
      .new({
        chat = chat,
        settings = settings,
      })
      :render()
  end,
}

M.toggle_system_prompt = {
  desc = "Toggle the system prompt",
  callback = function(chat)
    chat:toggle_system_prompt()
  end,
}

M.clear_rules = {
  desc = "Clear rules",
  callback = function(chat)
    chat:remove_tagged_message("rules")
    chat:refresh_context()
    return utils.notify("Cleared the rules", vim.log.levels.INFO)
  end,
}

M.clear_approvals = {
  desc = "Clear approvals in the current buffer",
  callback = function(chat)
    local approvals = require("codecompanion.interactions.chat.tools.approvals")
    approvals:reset(chat.bufnr)
    return utils.notify("Cleared the approvals", vim.log.levels.INFO)
  end,
}

M.yolo_mode = {
  desc = "Toggle YOLO mode",
  callback = function(chat)
    local approvals = require("codecompanion.interactions.chat.tools.approvals")
    local status = approvals:toggle_yolo_mode(chat.bufnr)
    if status then
      return utils.notify("YOLO mode enabled!", vim.log.levels.INFO)
    end
    return utils.notify("YOLO mode disabled!", vim.log.levels.INFO)
  end,
}

M.goto_file_under_cursor = {
  desc = "Open the file under cursor in a new tab.",
  ---@param chat CodeCompanion.Chat
  callback = function(chat)
    local file_name
    if vim.fn.mode() == "n" then
      file_name = vim.fn.expand("<cfile>")
    elseif string.lower(vim.fn.mode()):find("^.?v%a?") then
      -- one of the visual selection modes
      local start_pos = vim.fn.getpos("v")
      local end_pos = vim.fn.getpos(".")
      if start_pos[1] > end_pos[1] or (start_pos[1] == end_pos[1] and start_pos[2] > end_pos[2]) then
        start_pos, end_pos = end_pos, start_pos
      end
      local lines =
        vim.api.nvim_buf_get_text(chat.bufnr, start_pos[2] - 1, start_pos[3] - 1, end_pos[2] - 1, end_pos[3], {})
      if lines then
        file_name = table.concat(lines)
      end
    end
    if type(file_name) == "string" then
      file_name = vim.fs.normalize(file_name)
    else
      return
    end

    local stat = vim.uv.fs_stat(file_name)
    if stat == nil or stat.type ~= "file" then
      return
    end
    local action = nil
    local user_action = config.interactions.chat.opts.goto_file_action
    if type(user_action) == "string" then
      action = function(fname)
        vim.cmd(user_action .. " " .. fname)
      end
    elseif type(user_action) == "function" then
      action = user_action
    else
      error(string.format("%s is not a valid jump action!", vim.inspect(user_action)))
    end
    action(file_name)
  end,
}

M.copilot_stats = {
  desc = "Show Copilot usage statistics",
  callback = function(chat)
    if not chat.adapter.show_copilot_stats then
      return utils.notify("Stats are only available when using the Copilot adapter", vim.log.levels.WARN)
    end
    chat.adapter.show_copilot_stats()
  end,
}

M.super_diff = {
  desc = "Show super diff buffer",
  callback = function(chat)
    require("codecompanion.interactions.chat.super_diff").show_super_diff(chat)
  end,
}

return M