-- Since Neovim 0.12.0, the bundled Markdown Tree-sitter parser is enabled by
-- default. Fences get no highlighting when their language parser is missing.
-- `g:markdown_fenced_languages` can highlight such fences with legacy syntax
-- files, but eagerly loads every listed file, making it unsuitable for
-- Neovim's ~780 syntaxes. This file loads a syntax file only when a matching
-- fence is found and no Tree-sitter parser is available.

-- You can define optional fence-to-syntax aliases with:
--
--   vim.g.markdown_fenced_syntax_aliases = {
--     js = 'javascript',
--     shell = 'sh',
--   }

-- Do not let syntax/markdown.vim include any fenced languages eagerly.
vim.g.markdown_fenced_languages = {}

-- The Markdown Tree-sitter capture covers the entire fenced block. Any style
-- on it (especially a foreground color) takes precedence over the legacy
-- syntax fallback underneath it. Keep the block capture itself unstyled, and
-- reapply this after every colorscheme change.
local function clear_raw_block_highlights()
  for _, capture in ipairs({
    '@markup.raw.block.markdown',
    '@markdown.raw.block.markdown', -- Compatibility with older/custom queries.
  }) do
    vim.api.nvim_set_hl(0, capture, {})
  end
end

clear_raw_block_highlights()
vim.api.nvim_create_autocmd('ColorScheme', {
  callback = clear_raw_block_highlights,
})

local fence_query = vim.treesitter.query.parse(
  'markdown',
  '(fenced_code_block (info_string (language) @language))'
)

---@class MarkdownFallbackState
---@field syntaxes table<string, string> syntax name -> syntax cluster
---@field fences table<string, boolean>
---@field pending table<string, boolean>

---@type table<integer, MarkdownFallbackState>
local states = {}

local function normalize_language(language)
  return language:gsub('%s+', ''):lower():gsub('%-', '_')
end

-- Match the language resolution used by LanguageTree injections. Calling add()
-- here is still lazy: it is only attempted after that language occurs in a
-- fence. It also handles parsers loaded from an explicit path more accurately
-- than merely looking for parser/{lang}.* on 'runtimepath'.
local function has_parser(language)
  local alias = normalize_language(language)
  if not alias:match('^[%w_]+$') then
    return false
  end

  local function can_add(lang)
    local ok, loaded = pcall(vim.treesitter.language.add, lang)
    return ok and loaded == true
  end

  if can_add(alias) then
    return true
  end

  local lang = vim.treesitter.language.get_lang(alias)
  return lang ~= nil and lang ~= alias and can_add(lang)
end

local function syntax_identifier(name)
  return name:gsub('[^%w_]', function(char)
    return string.format('_%02x', string.byte(char))
  end)
end

local function syntax_file(language)
  local aliases = vim.g.markdown_fenced_syntax_aliases
  aliases = type(aliases) == 'table' and aliases or {}
  local syntax = aliases[language] or language
  if not syntax:match('^[%w_.+-]+$') then
    return nil
  end

  if #vim.api.nvim_get_runtime_file('syntax/' .. syntax .. '.vim', false) == 0 then
    return nil
  end

  return syntax
end

local function include_fallback(buf, language)
  if not vim.api.nvim_buf_is_valid(buf) or vim.bo[buf].filetype ~= 'markdown' then
    return
  end

  language = language:lower()
  local state = states[buf]
  if not state or state.fences[language] then
    return
  end
  state.pending[language] = nil

  if has_parser(language) then
    state.fences[language] = true
    return
  end

  local syntax = syntax_file(language)
  if not syntax then
    -- Do not mark this as handled: a syntax file added to 'runtimepath' later
    -- can then be picked up by a subsequent parse.
    return
  end

  vim.api.nvim_buf_call(buf, function()
    local cluster = state.syntaxes[syntax]

    if not cluster then
      -- Include each syntax file once, while allowing multiple fence aliases
      -- (for example, bash and shell) to use the resulting cluster.
      cluster = 'markdownLazy_' .. syntax_identifier(syntax)
      local current_syntax = vim.b.current_syntax
      vim.b.current_syntax = nil

      local ok, err = pcall(vim.cmd, string.format(
        'syntax include @%s %s',
        cluster,
        vim.fn.fnameescape('syntax/' .. syntax .. '.vim')
      ))

      vim.b.current_syntax = current_syntax
      if not ok then
        vim.notify(
          string.format('Could not load Markdown fallback syntax %q: %s', syntax, err),
          vim.log.levels.WARN
        )
        return
      end

      state.syntaxes[syntax] = cluster
    end

    -- \V...\m quotes punctuation in an info string before returning to the
    -- normal magic mode needed by the remainder of the Vim regex.
    local quoted_language = '\\V' .. language .. '\\m'
    local region = 'markdownLazyFence_' .. syntax_identifier(language)

    for _, delimiter in ipairs({ '`', '\\~' }) do
      vim.cmd(string.format(
        [[syntax region %s matchgroup=markdownCodeDelimiter start="^\s*\z(%s\{3,\}\)\s*\%%({.\{-}\.\)\=%s}\=\S\@!.*$" end="^\s*\z1\ze\s*$" keepend contains=@%s]],
        region,
        delimiter,
        quoted_language,
        cluster
      ))
    end

    state.fences[language] = true
  end)
end

local function inspect_tree(buf, tree)
  if not tree or not vim.api.nvim_buf_is_valid(buf) then
    return
  end

  local state = states[buf]
  if not state then
    return
  end

  local found = {}
  for id, node in fence_query:iter_captures(tree:root(), buf) do
    if fence_query.captures[id] == 'language' then
      local language = vim.treesitter.get_node_text(node, buf)
      language = language and language:match('^%s*(.-)%s*$') or nil
      if language and language ~= '' then
        found[language:lower()] = true
      end
    end
  end

  -- Syntax commands should not run from inside Tree-sitter's parse callback.
  for language in pairs(found) do
    if not state.fences[language] and not state.pending[language] then
      state.pending[language] = true
      vim.schedule(function()
        include_fallback(buf, language)
      end)
    end
  end
end

local function setup(buf)
  if states[buf] or not vim.api.nvim_buf_is_valid(buf) then
    return
  end

  states[buf] = { syntaxes = {}, fences = {}, pending = {} }

  -- vim.treesitter.start(), called by the 0.12 Markdown ftplugin, disables
  -- legacy syntax. Re-enable it so it can fill gaps left by injections.
  vim.bo[buf].syntax = 'ON'

  local ok, parser = pcall(vim.treesitter.get_parser, buf, 'markdown')
  if not ok then
    states[buf] = nil
    return
  end

  parser:register_cbs({
    on_changedtree = function(_, tree)
      inspect_tree(buf, tree)
    end,
    on_detach = function()
      states[buf] = nil
    end,
  })

  -- on_changedtree handles later edits and normally handles the initial parse.
  -- This explicit parse also covers the case where registration happened after
  -- the document's initial parse.
  parser:parse(false, function(err, trees)
    if err or not trees then
      return
    end
    for _, tree in pairs(trees) do
      inspect_tree(buf, tree)
    end
  end)
end

vim.api.nvim_create_autocmd('FileType', {
  pattern = 'markdown',
  callback = function(event)
    -- Run after runtime/ftplugin/markdown.lua has called treesitter.start().
    vim.schedule(function()
      setup(event.buf)
    end)
  end,
})

-- Also support sourcing this file while Markdown buffers are already open.
for _, buf in ipairs(vim.api.nvim_list_bufs()) do
  if vim.api.nvim_buf_is_loaded(buf) and vim.bo[buf].filetype == 'markdown' then
    vim.schedule(function()
      setup(buf)
    end)
  end
end