dotfiles

local function resolve_lang(lang)
  if lang ~= 'bash' then
    return lang
  end
  local shell = vim.o.shell
  if shell and (shell:match '^nu$' or shell:match '/nu$') then
    return 'nu'
  end
  return 'bash'
end

local function apply(buf, ns, text, lang, row, col_offset)
  lang = resolve_lang(lang)
  if not lang or lang == '' then
    return
  end
  if vim.api.nvim_buf_line_count(buf) < 1 then
    return
  end
  vim.api.nvim_buf_clear_namespace(buf, ns, row, row + 1)
  local parser_ok, parser = pcall(vim.treesitter.get_string_parser, text, lang)
  if not parser_ok or not parser then
    return
  end
  local parse_ok, trees = pcall(function()
    return parser:parse()
  end)
  if not parse_ok or not trees or #trees == 0 then
    return
  end
  local tree = trees[1]
  local query_ok, query = pcall(vim.treesitter.query.get, lang, 'highlights')
  if not query_ok or not query then
    return
  end
  for id, node in query:iter_captures(tree:root(), text) do
    local start_row, start_col, end_row, end_col = node:range()
    if start_row == end_row and start_col < end_col then
      local hl_group = '@' .. query.captures[id]
      vim.api.nvim_buf_set_extmark(buf, ns, row, col_offset + start_col, {
        end_col = col_offset + end_col,
        hl_group = hl_group,
        priority = 100,
      })
    end
  end
end

return {
  apply = apply,
}