Hide and show comments in code.

conceal_comment.lua

---@class ConcealComment
local M = {}

---@alias BufferHandle number
---@alias LineNumber number
---@alias ColumnNumber number
---@alias ExtmarkId number
---@alias NamespaceId number

---@class CommentNode
---@field start_row LineNumber 0-based line number
---@field start_col ColumnNumber 0-based column number
---@field end_row LineNumber 0-based line number
---@field end_col ColumnNumber 0-based column number
---@field text string The comment text content

---@class ConcealedLine
---@field row LineNumber 0-based line number
---@field extmark_id ExtmarkId The extmark ID for this concealed line
---@field original_text string The original line content

---@class NavigationDirection
---@field down number
---@field up number
local direction = {
  down = 1,
  up = -1,
}

---@class ConcealCommentConfig
---@field auto_enable boolean Automatically enable concealing for all supported languages
---@field smart_navigation boolean Enable smart navigation that skips concealed lines
---@field conceal_level number The conceallevel to set when concealing (0-3)
---@field refresh_on_change boolean Refresh concealing when buffer content changes
---@field debug boolean Enable debug logging

---@type ConcealCommentConfig
local default_config = {
  auto_enable = false,
  smart_navigation = true,
  conceal_level = 2,
  refresh_on_change = true,
  debug = false,
}

---@type ConcealCommentConfig
local config = vim.deepcopy(default_config)

---@type string The universal treesitter query for comments
local comment_query = "(comment) @comment"

---@type table<BufferHandle, ConcealedLine[]> Track concealed lines per buffer
local concealed_buffers = {}

---@type NamespaceId The namespace for concealing extmarks
local namespace_id = vim.api.nvim_create_namespace("conceal-comment")

---@type number Augroup ID for autocommands
local augroup_id = vim.api.nvim_create_augroup("ConcealComment", { clear = true })

---Debug logging utility
---@param message string
---@param level? number vim.log.levels
local function debug_log(message, level)
  if not config.debug then
    return
  end
  level = level or vim.log.levels.DEBUG
  vim.notify("[ConcealComment] " .. message, level)
end

---Validate buffer handle
---@param bufnr BufferHandle
---@return boolean is_valid
---@return string? error_message
local function validate_buffer(bufnr)
  if not bufnr or type(bufnr) ~= "number" then
    return false, "Invalid buffer handle"
  end

  if not vim.api.nvim_buf_is_valid(bufnr) then
    return false, "Buffer is not valid"
  end

  return true, nil
end

---Get treesitter parser for buffer
---@param bufnr BufferHandle
---@return TSParser? parser
---@return string? error_message
local function get_treesitter_parser(bufnr)
  local is_valid, error_msg = validate_buffer(bufnr)
  if not is_valid then
    return nil, error_msg
  end

  local ok, parser = pcall(vim.treesitter.get_parser, bufnr)
  if not ok or not parser then
    return nil, "Failed to get treesitter parser"
  end

  return parser, nil
end

---Get comment nodes from buffer using treesitter
---@param bufnr BufferHandle
---@return CommentNode[] nodes
---@return string? error_message
local function get_comment_nodes(bufnr)
  debug_log("Getting comment nodes for buffer " .. bufnr)

  local parser, parser_error = get_treesitter_parser(bufnr)
  if not parser then
    return {}, parser_error
  end

  local trees = parser:parse()
  if not trees or #trees == 0 then
    return {}, "No syntax trees available"
  end

  local root = trees[1]:root()
  if not root then
    return {}, "No root node available"
  end

  -- Try to parse the universal comment query
  local filetype = vim.api.nvim_get_option_value("filetype", { buf = bufnr })
  local ok, query = pcall(vim.treesitter.query.parse, filetype, comment_query)
  if not ok or not query then
    return {}, "Failed to parse comment query"
  end

  local nodes = {}
  for _, node in query:iter_captures(root, bufnr) do
    if node and node.range then
      local start_row, start_col, end_row, end_col = node:range()
      local text_lines = vim.api.nvim_buf_get_text(bufnr, start_row, start_col, end_row, end_col, {})
      local text = table.concat(text_lines, "\n")

      table.insert(nodes, {
        start_row = start_row,
        start_col = start_col,
        end_row = end_row,
        end_col = end_col,
        text = text,
      })
    end
  end

  debug_log("Found " .. #nodes .. " comment nodes")
  return nodes, nil
end

---Create concealing extmarks for comment nodes
---@param bufnr BufferHandle
---@param nodes CommentNode[]
---@return ConcealedLine[]
local function create_concealing_extmarks(bufnr, nodes)
  local concealed_lines = {}

  for _, node in ipairs(nodes) do
    for row = node.start_row, node.end_row do
      local line_text = vim.api.nvim_buf_get_lines(bufnr, row, row + 1, false)[1] or ""

      local ok, extmark_id = pcall(vim.api.nvim_buf_set_extmark, bufnr, namespace_id, row, 0, {
        end_row = row,
        end_col = #line_text,
        conceal_lines = "",
        priority = 1000,
      })

      if ok then
        table.insert(concealed_lines, {
          row = row,
          extmark_id = extmark_id,
          original_text = line_text,
        })
      else
        debug_log("Failed to create extmark for row " .. row, vim.log.levels.WARN)
      end
    end
  end

  return concealed_lines
end

---Apply comment concealing to buffer
---@param bufnr BufferHandle
---@return boolean success
---@return string? error_message
local function apply_concealing(bufnr)
  debug_log("Applying concealing to buffer " .. bufnr)

  local is_valid, error_msg = validate_buffer(bufnr)
  if not is_valid then
    return false, error_msg
  end

  -- Clear existing concealing
  vim.api.nvim_buf_clear_namespace(bufnr, namespace_id, 0, -1)
  concealed_buffers[bufnr] = nil

  local nodes, nodes_error = get_comment_nodes(bufnr)
  if nodes_error then
    return false, nodes_error
  end

  if #nodes == 0 then
    debug_log("No comment nodes found")
    return true, nil
  end

  -- Set conceallevel
  local current_win = vim.api.nvim_get_current_win()
  local buf_win = vim.fn.bufwinid(bufnr)

  if buf_win ~= -1 then
    vim.api.nvim_set_option_value("conceallevel", config.conceal_level, { win = buf_win })
  else
    vim.api.nvim_set_option_value("conceallevel", config.conceal_level, { win = current_win })
  end

  local concealed_lines = create_concealing_extmarks(bufnr, nodes)
  concealed_buffers[bufnr] = concealed_lines

  debug_log("Successfully concealed " .. #concealed_lines .. " lines")
  return true, nil
end

---Remove concealing from buffer
---@param bufnr BufferHandle
---@return boolean success
---@return string? error_message
local function remove_concealing(bufnr)
  debug_log("Removing concealing from buffer " .. bufnr)

  local is_valid, error_msg = validate_buffer(bufnr)
  if not is_valid then
    return false, error_msg
  end

  vim.api.nvim_buf_clear_namespace(bufnr, namespace_id, 0, -1)
  concealed_buffers[bufnr] = nil

  -- Reset conceallevel
  local current_win = vim.api.nvim_get_current_win()
  local buf_win = vim.fn.bufwinid(bufnr)

  if buf_win ~= -1 then
    vim.api.nvim_set_option_value("conceallevel", 0, { win = buf_win })
  else
    vim.api.nvim_set_option_value("conceallevel", 0, { win = current_win })
  end

  debug_log("Successfully removed concealing")
  return true, nil
end

---Check if a line is concealed
---@param bufnr BufferHandle
---@param line_nr LineNumber 1-based line number
---@return boolean is_concealed
local function is_line_concealed(bufnr, line_nr)
  local concealed_lines = concealed_buffers[bufnr]
  if not concealed_lines then
    return false
  end

  for _, concealed_line in ipairs(concealed_lines) do
    if concealed_line.row + 1 == line_nr then -- Convert 0-based to 1-based
      return true
    end
  end

  return false
end

---Find next visible (non-concealed) line
---@param bufnr BufferHandle
---@param current_line LineNumber 1-based line number
---@param direction number direction.down or direction.up
---@return LineNumber next_line
local function find_next_visible_line(bufnr, current_line, direction)
  local total_lines = vim.api.nvim_buf_line_count(bufnr)
  local next_line = current_line + direction

  while next_line >= 1 and next_line <= total_lines do
    if not is_line_concealed(bufnr, next_line) then
      return next_line
    end
    next_line = next_line + direction
  end

  -- Return boundary if no visible line found
  return direction > 0 and total_lines or 1
end

---Smart navigation that skips concealed lines
---@param direction number direction.down or direction.up
---@param count number Number of moves to make
local function smart_navigate(direction, count)
  local bufnr = vim.api.nvim_get_current_buf()

  -- Use normal navigation if concealing is not active
  if not concealed_buffers[bufnr] then
    local key = direction > 0 and "j" or "k"
    local cmd = count > 1 and (count .. key) or key
    vim.cmd("normal! " .. cmd)
    return
  end

  local current_line = vim.fn.line(".")
  local target_line = current_line

  for _ = 1, count do
    local next_line = find_next_visible_line(bufnr, target_line, direction)
    if next_line == target_line then
      break
    end
    target_line = next_line
  end

  if target_line ~= current_line then
    vim.api.nvim_win_set_cursor(0, { target_line, vim.fn.col(".") - 1 })
  end
end

---Setup smart navigation keymaps
local function setup_navigation_keymaps()
  local keymap_opts = { desc = "Smart comment navigation" }

  local function move_down()
    smart_navigate(direction.down, vim.v.count1)
  end

  local function move_up()
    smart_navigate(direction.up, vim.v.count1)
  end

  -- Normal mode mappings
  vim.keymap.set("n", "j", move_down, keymap_opts)
  vim.keymap.set("n", "k", move_up, keymap_opts)
  vim.keymap.set("n", "<Down>", move_down, keymap_opts)
  vim.keymap.set("n", "<Up>", move_up, keymap_opts)

  -- Visual mode mappings
  vim.keymap.set("v", "j", move_down, keymap_opts)
  vim.keymap.set("v", "k", move_up, keymap_opts)
  vim.keymap.set("v", "<Down>", move_down, keymap_opts)
  vim.keymap.set("v", "<Up>", move_up, keymap_opts)
end

---Setup autocommands for automatic concealing
local function setup_autocommands()
  if not config.auto_enable then
    return
  end

  vim.api.nvim_create_autocmd("FileType", {
    group = augroup_id,
    pattern = "*",
    callback = function(args)
      vim.schedule(function()
        if vim.treesitter.get_parser(args.buf, nil, { error = false }) then
          M.enable(args.buf)
        end
      end)
    end,
    desc = "Auto-enable comment concealing for all supported languages",
  })

  if config.refresh_on_change then
    vim.api.nvim_create_autocmd({ "TextChanged", "TextChangedI" }, {
      group = augroup_id,
      callback = function(args)
        if concealed_buffers[args.buf] then
          vim.schedule(function()
            M.enable(args.buf)
          end)
        end
      end,
      desc = "Refresh comment concealing on text changes",
    })
  end

  -- Cleanup on buffer delete
  vim.api.nvim_create_autocmd("BufDelete", {
    group = augroup_id,
    callback = function(args)
      concealed_buffers[args.buf] = nil
    end,
    desc = "Cleanup concealing data on buffer delete",
  })
end

---Setup user commands
local function setup_user_commands()
  vim.api.nvim_create_user_command("ConcealComments", function()
    local success, error_msg = M.enable()
    if not success then
      vim.notify("Failed to conceal comments: " .. (error_msg or "Unknown error"), vim.log.levels.ERROR)
    end
  end, { desc = "Conceal comments in current buffer" })

  vim.api.nvim_create_user_command("RevealComments", function()
    local success, error_msg = M.disable()
    if not success then
      vim.notify("Failed to reveal comments: " .. (error_msg or "Unknown error"), vim.log.levels.ERROR)
    end
  end, { desc = "Reveal comments in current buffer" })

  vim.api.nvim_create_user_command("ToggleConcealComments", function()
    local success, error_msg = M.toggle()
    if not success then
      vim.notify("Failed to toggle comment concealing: " .. (error_msg or "Unknown error"), vim.log.levels.ERROR)
    end
  end, { desc = "Toggle comment concealing in current buffer" })

  vim.api.nvim_create_user_command("ConcealCommentsStatus", function()
    local bufnr = vim.api.nvim_get_current_buf()
    local is_enabled = M.is_enabled(bufnr)
    local status = is_enabled and "enabled" or "disabled"
    local count = is_enabled and #concealed_buffers[bufnr] or 0
    vim.notify(string.format("Comment concealing is %s (%d lines concealed)", status, count), vim.log.levels.INFO)
  end, { desc = "Show comment concealing status" })
end

---Public API

---Setup the comment concealing module
---@param opts? ConcealCommentConfig User configuration options
function M.setup(opts)
  -- Merge user configuration with defaults
  config = vim.tbl_deep_extend("force", default_config, opts or {})

  debug_log("Setting up ConcealComment with config: " .. vim.inspect(config))

  -- Setup components
  if config.smart_navigation then
    setup_navigation_keymaps()
  end

  setup_autocommands()
  setup_user_commands()

  debug_log("ConcealComment setup completed successfully")
end

---Enable comment concealing for a buffer
---@param bufnr? BufferHandle Buffer handle (defaults to current buffer)
---@return boolean success
---@return string? error_message
function M.enable(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()

  local success, error_msg = apply_concealing(bufnr)
  if success then
    vim.notify("Comments concealed", vim.log.levels.INFO)
  end

  return success, error_msg
end

---Disable comment concealing for a buffer
---@param bufnr? BufferHandle Buffer handle (defaults to current buffer)
---@return boolean success
---@return string? error_message
function M.disable(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()

  local success, error_msg = remove_concealing(bufnr)
  if success then
    vim.notify("Comments revealed", vim.log.levels.INFO)
  end

  return success, error_msg
end

---Toggle comment concealing for a buffer
---@param bufnr? BufferHandle Buffer handle (defaults to current buffer)
---@return boolean success
---@return string? error_message
function M.toggle(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()

  if M.is_enabled(bufnr) then
    return M.disable(bufnr)
  else
    return M.enable(bufnr)
  end
end

---Check if comment concealing is enabled for a buffer
---@param bufnr? BufferHandle Buffer handle (defaults to current buffer)
---@return boolean is_enabled
function M.is_enabled(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()
  return concealed_buffers[bufnr] ~= nil
end

---Get concealing statistics for a buffer
---@param bufnr? BufferHandle Buffer handle (defaults to current buffer)
---@return table stats Statistics object with concealed line count and other info
function M.get_stats(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()

  local concealed_lines = concealed_buffers[bufnr] or {}
  local total_lines = vim.api.nvim_buf_line_count(bufnr)

  return {
    buffer = bufnr,
    total_lines = total_lines,
    concealed_lines = #concealed_lines,
    concealed_percentage = total_lines > 0 and (#concealed_lines / total_lines * 100) or 0,
    is_enabled = M.is_enabled(bufnr),
    is_supported = vim.treesitter.get_parser(bufnr, nil, { error = false }) ~= nil,
  }
end

---Get current configuration
---@return ConcealCommentConfig config Current configuration
function M.get_config()
  return vim.deepcopy(config)
end

---Update configuration
---@param new_config ConcealCommentConfig New configuration options
function M.update_config(new_config)
  config = vim.tbl_deep_extend("force", config, new_config)
  debug_log("Configuration updated: " .. vim.inspect(new_config))
end

return M