--
-- 50-extmarks.lua
--
-- * Creates a global highlight namespace for extended marks.
-- * Keeps track of the current global highlight namespace in `vim.g.hl_ns`.
-- * Creates an autocommand that initializes and updates extended marks when 'background' or the
--   color scheme change.
--   For now, only "TODO" strings are marked but the system can be easily extended by adding items
--   in `marks`.
--   Doesn't support multi-line patterns.
--
-- Global variables:
--   `vim.g.hl_ns`: global highlight namespace
--
-- User commands:
--   `ExtmarkToggle`: toggle extended marks highlight namespace
--

vim.g.hl_ns = 0

local extmarks_hl_ns = vim.api.nvim_create_namespace("dotfiles_extmarks")

local hl_group_todo = "dotfiles.Todo"
vim.api.nvim_set_hl(extmarks_hl_ns, hl_group_todo, { link = "Todo" })

-- Wrapper around `nvim_set_hl_ns()`; sets the global highlight namespace to `hl_ns` and save it to
-- `vim.g.hl_ns`. Useful because the global highlight namespace is not saved anywhere by default.
local function set_global_hl_ns(hl_ns)
	vim.g.hl_ns = hl_ns
	vim.api.nvim_set_hl_ns(hl_ns)
end

-- Toggles extended marks highlight namespace
local function toggle_marks()
	local new_hl_ns = vim.g.hl_ns == 0 and extmarks_hl_ns or 0
	set_global_hl_ns(new_hl_ns)
end

-- Parses the treesitter tree synchronously, including injections, in the 0-based, end-exclusive
-- line range `first`:`last` of buffer `buf`.
--- @return boolean parser_exists
local function parse_tree(buf, first, last)
	local parser = vim.treesitter.get_parser(buf)
	if parser then
		parser:parse({ first, last })
	end
	return parser ~= nil
end

-- Returns a predicate telling whether a "TODO" string should be marked at a given position.
-- In markdown buffers, positions inside code are invalid, otherwise positions outside comments are
-- invalid. All positions are valid if there is no parser for the buffer filetype.
-- Only the line range `first`:`last` is parsed; the predicate must not be called outside it.
--- @param buf number
--- @return fun(row: number, col: number): boolean predicate taking a 0-based position
local function todo_filter(buf, first, last)
	local has_parser = parse_tree(buf, first, last)
	return function(row, col)
		assert(first <= row and row < last, "Position outside the parsed line range")
		if not has_parser then
			return true
		end
		local block = vim.treesitter.get_node({
			bufnr = buf,
			pos = { row, col },
		})
		assert(block, "Expected node to be truthy, is the tree parsed?")
		if vim.bo[buf].filetype == "markdown" then
			if string.match(block:type(), "code") then
				return false
			end
			local inline = vim.treesitter.get_node({
				bufnr = buf,
				ignore_injections = false,
				pos = { row, col },
			})
			assert(inline, "Expected node to be truthy, is the tree parsed?")
			return inline:type() ~= "code_span"
		else
			return string.match(block:type(), "comment") ~= nil
		end
	end
end

-- Returns an iterator over the positions of valid "TODO" strings in buffer `buf` in the line range
-- `first`:`last`. Each item is a table with 0-based fields: row `y`, start column `x1` (inclusive)
-- and end column `x2` (exclusive). See `todo_filter()` for what makes a position valid.
local function find_todo(buf, first, last)
	local lines = vim.api.nvim_buf_get_lines(buf, first, last, false)
	local should_mark = todo_filter(buf, first, first + #lines)
	local x1, x2, y = nil, nil, 1
	return function()
		while y <= #lines do
			x1, x2 = string.find(lines[y], "TODO", (x2 or 0) + 1)
			if x1 and x2 then
				local row = y + first - 1
				if should_mark(row, x1 - 1) then
					return { y = row, x1 = x1 - 1, x2 = x2 }
				end
			else
				x1, x2, y = 0, 0, y + 1
			end
		end
		return nil
	end
end

-- List of extended marks to apply with `mark_all()`
--
-- Each item is a table with two fields: `hl_group`, the name of the highlight group to use
-- for the marks, and `find`, which must return an iterator over the positions to mark (see
-- `find_todo()` for the format).
local marks = {
	{ hl_group = hl_group_todo, find = find_todo },
}

-- Updates extended marks in buffer `buf` in range `first`:`last`.
-- For each item in `marks`, sets extended marks on positions returned by `item.find` to the
-- highlight group `item.hl_group`.
local function mark_all(buf, first, last)
	vim.api.nvim_buf_clear_namespace(buf, extmarks_hl_ns, first, last)
	for _, marks_opts in ipairs(marks) do
		for p in marks_opts.find(buf, first, last) do
			vim.api.nvim_buf_set_extmark(
				buf,
				extmarks_hl_ns,
				p.y,
				p.x1,
				{ end_col = p.x2, hl_group = marks_opts.hl_group }
			)
		end
	end
end

-- Create extended marks in buffer `ev.buf` and updates them when the buffer changes.
local function watch_extmarks(ev)
	mark_all(ev.buf, 0, -1)
	if vim.b[ev.buf].dotfiles_todo_attached then
		return
	end
	vim.b[ev.buf].dotfiles_todo_attached = true
	vim.api.nvim_buf_attach(ev.buf, false, {
		on_lines = function(_, b, _, first, _, last_new)
			vim.schedule(function()
				if vim.api.nvim_buf_is_valid(b) then
					mark_all(b, first, last_new)
				end
			end)
		end,
	})
end

----------------------------------------------------------------------------------------------------

vim.api.nvim_create_autocmd("FileType", {
	desc = "Set extended marks watcher",
	group = vim.g.dotfiles.augroup,
	callback = watch_extmarks,
})

vim.api.nvim_create_user_command("ExtmarkToggle", toggle_marks, { desc = "Toggle extended marks" })

set_global_hl_ns(extmarks_hl_ns)