configs.lua

local api = vim.api

local queries = require "nvim-treesitter.query"
local ts = require "nvim-treesitter.compat"
local parsers = require "nvim-treesitter.parsers"
local utils = require "nvim-treesitter.utils"
local caching = require "nvim-treesitter.caching"

local M = {}

---@class TSConfig
---@field modules {[string]:TSModule}
---@field sync_install boolean
---@field ensure_installed string[]|string
---@field ignore_install string[]
---@field auto_install boolean
---@field parser_install_dir string|nil

---@type TSConfig
local config = {
  modules = {},
  sync_install = false,
  ensure_installed = {},
  auto_install = false,
  ignore_install = {},
  parser_install_dir = nil,
}

-- List of modules that need to be setup on initialization.
---@type TSModule[][]
local queued_modules_defs = {}
-- Whether we've initialized the plugin yet.
local is_initialized = false

---@class TSModule
---@field module_path string
---@field enable boolean|string[]|function(string): boolean
---@field disable boolean|string[]|function(string): boolean
---@field keymaps table<string, string>
---@field is_supported function(string): boolean
---@field attach function(string)
---@field detach function(string)
---@field enabled_buffers table<integer, boolean>
---@field additional_vim_regex_highlighting boolean|string[]

---@type {[string]: TSModule}
local builtin_modules = {
  highlight = {
    module_path = "nvim-treesitter.highlight",
    -- @deprecated: use `highlight.set_custom_captures` instead
    custom_captures = {},
    enable = false,
    is_supported = function(lang)
      return queries.has_highlights(lang)
    end,
    additional_vim_regex_highlighting = false,
  },
  incremental_selection = {
    module_path = "nvim-treesitter.incremental_selection",
    enable = false,
    keymaps = {
      init_selection = "gnn", -- set to `false` to disable one of the mappings
      node_incremental = "grn",
      scope_incremental = "grc",
      node_decremental = "grm",
    },
    is_supported = function()
      return true
    end,
  },
  indent = {
    module_path = "nvim-treesitter.indent",
    enable = false,
    is_supported = queries.has_indents,
  },
}

local attached_buffers_by_module = caching.create_buffer_cache()

---Resolves a module by requiring the `module_path` or using the module definition.
---@param mod_name string
---@return TSModule|nil
local function resolve_module(mod_name)
  local config_mod = M.get_module(mod_name)

  if not config_mod then
    return
  end

  if type(config_mod.attach) == "function" and type(config_mod.detach) == "function" then
    return config_mod
  elseif type(config_mod.module_path) == "string" then
    return require(config_mod.module_path)
  end
end

---Enables and attaches the module to a buffer for lang.
---@param mod string path to module
---@param bufnr integer|nil buffer number, defaults to current buffer
---@param lang string|nil language, defaults to current language
local function enable_module(mod, bufnr, lang)
  local module = M.get_module(mod)
  if not module then
    return
  end

  bufnr = bufnr or api.nvim_get_current_buf()
  lang = lang or parsers.get_buf_lang(bufnr)

  if not module.enable then
    if module.enabled_buffers then
      module.enabled_buffers[bufnr] = true
    else
      module.enabled_buffers = { [bufnr] = true }
    end
  end

  M.attach_module(mod, bufnr, lang)
end

---Enables autocomands for the module.
---After the module is loaded `loaded` will be set to true for the module.
---@param mod string path to module
local function enable_mod_conf_autocmd(mod)
  local config_mod = M.get_module(mod)
  if not config_mod or config_mod.loaded then
    return
  end

  api.nvim_create_autocmd("FileType", {
    group = api.nvim_create_augroup("NvimTreesitter-" .. mod, {}),
    callback = function(args)
      require("nvim-treesitter.configs").reattach_module(mod, args.buf)
    end,
    desc = "Reattach module",
  })

  config_mod.loaded = true
end

---Enables the module globally and for all current buffers.
---After enabled, `enable` will be set to true for the module.
---@param mod string path to module
local function enable_all(mod)
  local config_mod = M.get_module(mod)
  if not config_mod then
    return
  end

  enable_mod_conf_autocmd(mod)
  config_mod.enable = true
  config_mod.enabled_buffers = nil

  for _, bufnr in pairs(api.nvim_list_bufs()) do
    enable_module(mod, bufnr)
  end
end

---Disables and detaches the module for a buffer.
---@param mod string path to module
---@param bufnr integer buffer number, defaults to current buffer
local function disable_module(mod, bufnr)
  local module = M.get_module(mod)
  if not module then
    return
  end

  bufnr = bufnr or api.nvim_get_current_buf()
  if module.enabled_buffers then
    module.enabled_buffers[bufnr] = false
  end
  M.detach_module(mod, bufnr)
end

---Disables autocomands for the module.
---After the module is unloaded `loaded` will be set to false for the module.
---@param mod string path to module
local function disable_mod_conf_autocmd(mod)
  local config_mod = M.get_module(mod)
  if not config_mod or not config_mod.loaded then
    return
  end
  api.nvim_clear_autocmds { event = "FileType", group = "NvimTreesitter-" .. mod }
  config_mod.loaded = false
end

---Disables the module globally and for all current buffers.
---After disabled, `enable` will be set to false for the module.
---@param mod string path to module
local function disable_all(mod)
  local config_mod = M.get_module(mod)
  if not config_mod then
    return
  end

  config_mod.enabled_buffers = nil
  disable_mod_conf_autocmd(mod)
  config_mod.enable = false

  for _, bufnr in pairs(api.nvim_list_bufs()) do
    disable_module(mod, bufnr)
  end
end

---Toggles a module for a buffer
---@param mod string path to module
---@param bufnr integer buffer number, defaults to current buffer
---@param lang string language, defaults to current language
local function toggle_module(mod, bufnr, lang)
  bufnr = bufnr or api.nvim_get_current_buf()
  lang = lang or parsers.get_buf_lang(bufnr)

  if attached_buffers_by_module.has(mod, bufnr) then
    disable_module(mod, bufnr)
  else
    enable_module(mod, bufnr, lang)
  end
end

-- Toggles the module globally and for all current buffers.
-- @param mod path to module
local function toggle_all(mod)
  local config_mod = M.get_module(mod)
  if not config_mod then
    return
  end

  if config_mod.enable then
    disable_all(mod)
  else
    enable_all(mod)
  end
end

---Recurses through all modules including submodules
---@param accumulator function called for each module
---@param root {[string]: TSModule}|nil root configuration table to start at
---@param path string|nil prefix path
local function recurse_modules(accumulator, root, path)
  root = root or config.modules

  for name, module in pairs(root) do
    local new_path = path and (path .. "." .. name) or name

    if M.is_module(module) then
      accumulator(name, module, new_path, root)
    elseif type(module) == "table" then
      recurse_modules(accumulator, module, new_path)
    end
  end
end

-- Shows current configuration of all nvim-treesitter modules
---@param process_function function used as the `process` parameter
---       for vim.inspect (https://github.com/kikito/inspect.lua#optionsprocess)
local function config_info(process_function)
  process_function = process_function
    or function(item, path)
      if path[#path] == vim.inspect.METATABLE then
        return
      end
      if path[#path] == "is_supported" then
        return
      end
      return item
    end
  print(vim.inspect(config, { process = process_function }))
end

---@param query_group string
---@param lang string
function M.edit_query_file(query_group, lang)
  lang = lang or parsers.get_buf_lang()
  local files = ts.get_query_files(lang, query_group, true)
  if #files == 0 then
    utils.notify "No query file found! Creating a new one!"
    M.edit_query_file_user_after(query_group, lang)
  elseif #files == 1 then
    vim.cmd(":edit " .. files[1])
  else
    vim.ui.select(files, { prompt = "Select a file:" }, function(file)
      if file then
        vim.cmd(":edit " .. file)
      end
    end)
  end
end

---@param query_group string
---@param lang string
function M.edit_query_file_user_after(query_group, lang)
  lang = lang or parsers.get_buf_lang()
  local folder = utils.join_path(vim.fn.stdpath "config", "after", "queries", lang)
  local file = utils.join_path(folder, query_group .. ".scm")
  if vim.fn.isdirectory(folder) ~= 1 then
    vim.ui.select({ "Yes", "No" }, { prompt = '"' .. folder .. '" does not exist. Create it?' }, function(choice)
      if choice == "Yes" then
        vim.fn.mkdir(folder, "p", "0755")
        vim.cmd(":edit " .. file)
      end
    end)
  else
    vim.cmd(":edit " .. file)
  end
end

M.commands = {
  TSBufEnable = {
    run = enable_module,
    args = {
      "-nargs=1",
      "-complete=custom,nvim_treesitter#available_modules",
    },
  },
  TSBufDisable = {
    run = disable_module,
    args = {
      "-nargs=1",
      "-complete=custom,nvim_treesitter#available_modules",
    },
  },
  TSBufToggle = {
    run = toggle_module,
    args = {
      "-nargs=1",
      "-complete=custom,nvim_treesitter#available_modules",
    },
  },
  TSEnable = {
    run = enable_all,
    args = {
      "-nargs=+",
      "-complete=custom,nvim_treesitter#available_modules",
    },
  },
  TSDisable = {
    run = disable_all,
    args = {
      "-nargs=+",
      "-complete=custom,nvim_treesitter#available_modules",
    },
  },
  TSToggle = {
    run = toggle_all,
    args = {
      "-nargs=+",
      "-complete=custom,nvim_treesitter#available_modules",
    },
  },
  TSConfigInfo = {
    run = config_info,
    args = {
      "-nargs=0",
    },
  },
  TSEditQuery = {
    run = M.edit_query_file,
    args = {
      "-nargs=+",
      "-complete=custom,nvim_treesitter#available_query_groups",
    },
  },
  TSEditQueryUserAfter = {
    run = M.edit_query_file_user_after,
    args = {
      "-nargs=+",
      "-complete=custom,nvim_treesitter#available_query_groups",
    },
  },
}

---@param mod string module
---@param lang string the language of the buffer
---@param bufnr integer the buffer
function M.is_enabled(mod, lang, bufnr)
  if not parsers.has_parser(lang) then
    return false
  end

  local module_config = M.get_module(mod)
  if not module_config then
    return false
  end

  local buffer_enabled = module_config.enabled_buffers and module_config.enabled_buffers[bufnr]
  local config_enabled = module_config.enable or buffer_enabled
  if not config_enabled or not module_config.is_supported(lang) then
    return false
  end

  local disable = module_config.disable
  if type(disable) == "function" then
    if disable(lang, bufnr) then
      return false
    end
  elseif type(disable) == "table" then
    -- Otherwise it's a list of languages
    for _, parser in pairs(disable) do
      if lang == parser then
        return false
      end
    end
  end

  return true
end

---Setup call for users to override module configurations.
---@param user_data TSConfig module overrides
function M.setup(user_data)
  config.modules = vim.tbl_deep_extend("force", config.modules, user_data)
  config.ignore_install = user_data.ignore_install or {}
  config.parser_install_dir = user_data.parser_install_dir or nil
  if config.parser_install_dir then
    config.parser_install_dir = vim.fn.expand(config.parser_install_dir, ":p")
  end

  config.auto_install = user_data.auto_install or false
  if config.auto_install then
    require("nvim-treesitter.install").setup_auto_install()
  end

  local ensure_installed = user_data.ensure_installed or {}
  if #ensure_installed > 0 then
    if user_data.sync_install then
      require("nvim-treesitter.install").ensure_installed_sync(ensure_installed)
    else
      require("nvim-treesitter.install").ensure_installed(ensure_installed)
    end
  end

  config.modules.ensure_installed = nil
  config.ensure_installed = ensure_installed

  recurse_modules(function(_, _, new_path)
    local data = utils.get_at_path(config.modules, new_path)
    if data.enable then
      enable_all(new_path)
    end
  end, config.modules)
end

-- Defines a table of modules that can be attached/detached to buffers
-- based on language support. A module consist of the following properties:
---* @enable Whether the modules is enabled. Can be true or false.
---* @disable A list of languages to disable the module for. Only relevant if enable is true.
---* @keymaps A list of user mappings for a given module if relevant.
---* @is_supported A function which, given a ft, will return true if the ft works on the module.
---* @module_path A string path to a module file using `require`. The exported module must contain
---               an `attach` and `detach` function. This path is not required if `attach` and `detach`
---               functions are provided directly on the module definition.
---* @attach An attach function that is called for each buffer that the module is enabled for. This is required
---          if a `module_path` is not specified.
---* @detach A detach function that is called for each buffer that the module is enabled for. This is required
---          if a `module_path` is not specified.
--
-- Modules are not setup until `init` is invoked by the plugin. This allows modules to be defined in any order
-- and can be loaded lazily.
--
---* @example
---require"nvim-treesitter".define_modules {
---  my_cool_module = {
---    attach = function()
---      do_some_cool_setup()
---    end,
---    detach = function()
---      do_some_cool_teardown()
---    end
---  }
---}
---@param mod_defs TSModule[]
function M.define_modules(mod_defs)
  if not is_initialized then
    table.insert(queued_modules_defs, mod_defs)
    return
  end

  recurse_modules(function(key, mod, _, group)
    group[key] = vim.tbl_extend("keep", mod, {
      enable = false,
      disable = {},
      is_supported = function()
        return true
      end,
    })
  end, mod_defs)

  config.modules = vim.tbl_deep_extend("keep", config.modules, mod_defs)

  for _, mod in ipairs(M.available_modules(mod_defs)) do
    local module_config = M.get_module(mod)
    if module_config and module_config.enable then
      enable_mod_conf_autocmd(mod)
    end
  end
end

---Attaches a module to a buffer
---@param mod_name string the module name
---@param bufnr integer the buffer
---@param lang string the language of the buffer
function M.attach_module(mod_name, bufnr, lang)
  bufnr = bufnr or api.nvim_get_current_buf()
  lang = lang or parsers.get_buf_lang(bufnr)
  local resolved_mod = resolve_module(mod_name)

  if resolved_mod and not attached_buffers_by_module.has(mod_name, bufnr) and M.is_enabled(mod_name, lang, bufnr) then
    attached_buffers_by_module.set(mod_name, bufnr, true)
    resolved_mod.attach(bufnr, lang)
  end
end

-- Detaches a module to a buffer
---@param mod_name string the module name
---@param bufnr integer the buffer
function M.detach_module(mod_name, bufnr)
  local resolved_mod = resolve_module(mod_name)
  bufnr = bufnr or api.nvim_get_current_buf()

  if resolved_mod and attached_buffers_by_module.has(mod_name, bufnr) then
    attached_buffers_by_module.remove(mod_name, bufnr)
    resolved_mod.detach(bufnr)
  end
end

-- Same as attach_module, but if the module is already attached, detach it first.
---@param mod_name string the module name
---@param bufnr integer the buffer
---@param lang string the language of the buffer
function M.reattach_module(mod_name, bufnr, lang)
  M.detach_module(mod_name, bufnr)
  M.attach_module(mod_name, bufnr, lang)
end

-- Gets available modules
---@param root {[string]:TSModule}|nil table to find modules
---@return string[] modules list of module paths
function M.available_modules(root)
  local modules = {}

  recurse_modules(function(_, _, path)
    table.insert(modules, path)
  end, root)

  return modules
end

---Gets a module config by path
---@param mod_path string path to the module
---@return TSModule|nil: the module or nil
function M.get_module(mod_path)
  local mod = utils.get_at_path(config.modules, mod_path)

  return M.is_module(mod) and mod or nil
end

-- Determines whether the provided table is a module.
-- A module should contain an attach and detach function.
---@param mod table|nil the module table
---@return boolean
function M.is_module(mod)
  return type(mod) == "table"
    and ((type(mod.attach) == "function" and type(mod.detach) == "function") or type(mod.module_path) == "string")
end

-- Initializes built-in modules and any queued modules
-- registered by plugins or the user.
function M.init()
  is_initialized = true
  M.define_modules(builtin_modules)

  for _, mod_def in ipairs(queued_modules_defs) do
    M.define_modules(mod_def)
  end
end

-- If parser_install_dir is not nil is used or created.
-- If parser_install_dir is nil try the package dir of the nvim-treesitter
-- plugin first, followed by the "site" dir from "runtimepath". "site" dir will
-- be created if it doesn't exist. Using only the package dir won't work when
-- the plugin is installed with Nix, since the "/nix/store" is read-only.
---@param folder_name string|nil
---@return string|nil, string|nil
function M.get_parser_install_dir(folder_name)
  folder_name = folder_name or "parser"

  local install_dir = config.parser_install_dir or utils.get_package_path()
  local parser_dir = utils.join_path(install_dir, folder_name)

  return utils.create_or_reuse_writable_dir(
    parser_dir,
    utils.join_space("Could not create parser dir '", parser_dir, "': "),
    utils.join_space(
      "Parser dir '",
      parser_dir,
      "' should be read/write (see README on how to configure an alternative install location)"
    )
  )
end

function M.get_parser_info_dir()
  return M.get_parser_install_dir "parser-info"
end

function M.get_ignored_parser_installs()
  return config.ignore_install or {}
end

function M.get_ensure_installed_parsers()
  if type(config.ensure_installed) == "string" then
    return { config.ensure_installed }
  end
  return config.ensure_installed or {}
end

return M