Dynamically enable/disable some LSP features

---@class _MethodState
---@field is_enabled boolean
---@field enable function
---@field disable function

---@class _BufferState
---@field document_highlight? _MethodState
---@field inlay_hint? _MethodState
---@field codelens? _MethodState

---@type table<integer, _BufferState>
local STATE_PER_ATTACHED_BUFFER = {}

-- help functions to aide avoid large levels of indentation
local H = {
  inlay_hint = {},
  codelens = {},
  document_highlight = {},
}

local timer = vim.uv.new_timer()
timer:start(
  tonumber(vim.g.lsp_refresh_time) or 1000,
  tonumber(vim.g.lsp_refresh_time) or 1000,
  vim.schedule_wrap(function()
    local function handler(buf, state, kind)
      if not vim.api.nvim_buf_is_valid(buf) or not (state and state[kind]) then
        return
      end

      if H.should_be_enabled(buf, kind) then
        state[kind].enable()
      else
        state[kind].disable()
      end
    end

    for buf, state in pairs(STATE_PER_ATTACHED_BUFFER) do
      for _, kind in ipairs({ 'document_highlight', 'inlay_hint', 'codelens' }) do
        handler(buf, state, kind)
      end
    end
  end)
)

vim.api.nvim_create_augroup('lsp_attach', { clear = true })
vim.api.nvim_create_augroup('lsp_deatach', { clear = true })
vim.api.nvim_create_augroup('lsp_document_highlight_clear', { clear = true })
vim.api.nvim_create_augroup('lsp_document_highlight', { clear = false })
vim.api.nvim_create_augroup('lsp_inlay_hint', { clear = false })
vim.api.nvim_create_augroup('lsp_codelens', { clear = false })

vim.api.nvim_create_autocmd('LspAttach', {
  group = 'lsp_attach',
  callback = function(ev)
    local client = vim.lsp.get_client_by_id(ev.data.client_id)

    for kind, method in pairs({
      document_highlight = 'textDocument/documentHighlight',
      inlay_hint = 'textDocument/inlayHint',
      codelens = 'textDocument/codeLens',
    }) do
      if client:supports_method(method) then
        STATE_PER_ATTACHED_BUFFER[ev.buf] = STATE_PER_ATTACHED_BUFFER[ev.buf] or {}
        STATE_PER_ATTACHED_BUFFER[ev.buf][kind] = STATE_PER_ATTACHED_BUFFER[ev.buf][kind] or {}

        local state = STATE_PER_ATTACHED_BUFFER[ev.buf][kind]

        function state.enable()
          H[kind].enable(ev.buf, state)
        end

        function state.disable()
          H[kind].disable(ev.buf, state)
        end

        if H.should_be_enabled(ev.buf, kind) then
          state.enable()
        else
          state.is_enabled = false
        end
      end
    end
  end,
})

vim.api.nvim_create_autocmd('LspDetach', {
  group = 'lsp_deatach',
  callback = function(ev)
    for kind, method in pairs({
      document_highlight = 'textDocument/documentHighlight',
      inlay_hint = 'textDocument/inlayHint',
      codelens = 'textDocument/codeLens',
    }) do
      if
        not vim.iter(vim.lsp.get_clients({ bufnr = ev.buf })):any(function(client)
          return client.id ~= ev.data.client_id and client:supports_method(method)
        end)
      then
        STATE_PER_ATTACHED_BUFFER[ev.buf][kind] = nil
      end
    end

    if vim.tbl_isempty(STATE_PER_ATTACHED_BUFFER[ev.buf]) then
      STATE_PER_ATTACHED_BUFFER[ev.buf] = nil
    end
  end,
})

function H.should_be_enabled(buf, kind)
  local index = 'lsp_' .. kind .. '_enable'
  local b_var = vim.b[buf][index]
  local g_var = vim.g[index]

  return (b_var == nil) and (g_var == true) or (b_var == true)
end

function H.document_highlight.enable(buf, state)
  if not state.is_enabled then
    state.is_enabled = true

    vim.api.nvim_create_autocmd({ 'CursorHold', 'CursorHoldI' }, {
      buffer = buf,
      group = 'lsp_document_highlight',
      callback = vim.lsp.buf.document_highlight,
    })

    vim.api.nvim_create_autocmd({ 'CursorMoved', 'CursorMovedI', 'LspDetach' }, {
      buffer = buf,
      group = 'lsp_document_highlight',
      callback = vim.lsp.buf.clear_references,
    })

    vim.api.nvim_create_autocmd('LspDetach', {
      group = 'lsp_document_highlight_clear',
      callback = function(ev)
        -- remove the above autocmds if the buffer does not have any
        -- attached client that supports document highlight
        if
          not vim.iter(vim.lsp.get_clients({ bufnr = buf })):any(function(client)
            return client.id ~= ev.data.client_id
              and client:supports_method('textDocument/documentHighlight')
          end)
        then
          vim.api.nvim_clear_autocmds({
            buffer = buf,
            group = 'lsp_document_highlight',
          })
        end
      end,
    })
  end
end

function H.document_highlight.disable(buf, state)
  if state.is_enabled then
    state.is_enabled = false

    vim.api.nvim_buf_call(buf, function()
      vim.lsp.buf.clear_references()
    end)

    vim.api.nvim_clear_autocmds({
      buffer = buf,
      group = 'lsp_document_highlight',
    })

    vim.api.nvim_clear_autocmds({
      buffer = buf,
      group = 'lsp_document_highlight_clear',
    })
  end
end

function H.inlay_hint.enable(buf, state)
  if not state.is_enabled then
    state.is_enabled = true

    vim.lsp.inlay_hint.enable(true, { bufnr = buf })

    vim.api.nvim_create_autocmd('InsertEnter', {
      group = 'lsp_inlay_hint',
      buffer = buf,
      callback = function()
        vim.lsp.inlay_hint.enable(false, { bufnr = buf })
      end,
    })

    vim.api.nvim_create_autocmd('InsertLeave', {
      group = 'lsp_inlay_hint',
      buffer = buf,
      callback = function()
        vim.lsp.inlay_hint.enable(true, { bufnr = buf })
      end,
    })
  end
end

function H.inlay_hint.disable(buf, state)
  if state.is_enabled then
    state.is_enabled = false
    vim.lsp.inlay_hint.enable(false, { bufnr = buf })
    vim.api.nvim_clear_autocmds({
      group = 'lsp_inlay_hint',
      buffer = buf,
    })
  end
end

function H.codelens.enable(buf, state)
  if not state.is_enabled then
    state.is_enabled = true

    vim.lsp.codelens.refresh({ bufnr = buf })

    vim.api.nvim_create_autocmd({ 'BufEnter', 'CursorHold', 'InsertLeave' }, {
      group = 'lsp_codelens',
      buffer = buf,
      callback = function()
        vim.lsp.codelens.refresh({ bufnr = buf })
      end,
    })

    vim.api.nvim_create_autocmd('InsertEnter', {
      group = 'lsp_codelens',
      buffer = buf,
      callback = function()
        vim.lsp.codelens.clear(nil, buf)
      end,
    })
  end
end

function H.codelens.disable(buf, state)
  if state.is_enabled then
    state.is_enabled = false
    vim.lsp.codelens.clear(nil, buf)

    vim.api.nvim_clear_autocmds({
      group = 'lsp_codelens',
      buffer = buf,
    })
  end
end