--
-- TODO highlighting (treesitter-aware extmarks)
--
-- The `dotfiles.Todo` highlight group and the storage namespace are owned by
-- 40-colors.lua, because the same per-window namespaces also drive the
-- focused/unfocused cursor line. That file is numbered 40- so it is sourced
-- before this one and `vim.g.dotfiles.todo_ns` / `.todo_hl` exist below. Here
-- we only place the extmarks.

-- Returns an iterator over the ascendants of `node`, starting at the root node
local function node_ascendants(node)
	local root = node:tree():root()
	local current = root
	return function()
		if current:equal(node) then
			return nil
		end
		current = assert(current:child_with_descendant(node))
		return current
	end
end

-- Returns true if `node` or any of its parent has the type `type`
local function node_within_type(node, type)
	local next_ascendant = node_ascendants(node)
	local next = next_ascendant()
	while next do
		if next:type() == type then
			return true
		end
		next = next_ascendant()
	end
	return false
end

-- Returns an iterator over the positions (`{ col1, col2, row }`) of matches of pattern in lines
local function pmatches(lines, pattern)
	local col1, col2, row = nil, nil, 1
	return function()
		while row <= #lines do
			col1, col2 = string.find(lines[row], pattern, (col2 or 0) + 1)
			if col1 then
				return { col1 = col1, col2 = col2, row = row }
			else
				col1, col2, row = 0, 0, row + 1
			end
		end
		return nil
	end
end

-- Set extmarks are the positions (`{ col1, col2, row }`) given by `next_pos()`
-- `opts`: {
--     ns: highlight namespace
--     predicate: is passed row and col, if true then set the mark
--     hl_group: highlight group
--  }
local function set_extmarks(buf, next_pos, opts)
	for pos in next_pos do
		local marks = vim.api.nvim_buf_get_extmarks(
			buf,
			opts.ns,
			{ pos.row - 1, pos.col1 - 1 },
			{ pos.row - 1, pos.col2 - 1 },
			{}
		)
		if #marks == 0 and opts.predicate(buf, pos.row - 1, pos.col1 - 1) then
			vim.api.nvim_buf_set_extmark(
				buf,
				opts.ns,
				pos.row - 1,
				pos.col1 - 1,
				{ end_col = pos.col2, hl_group = opts.hl_group }
			)
		end
	end
end

-- TODO a better way to do this:
-- For code, only check in "comment" node ranges
-- For non-code, do it per type. For Markdown, only check "paragraph" nodes
local function hl_todo_predicate(buf, row, col)
	if vim.bo[buf].filetype == "markdown" or vim.treesitter.get_parser(buf) == nil then
		return true
	end
	local node = assert(vim.treesitter.get_node({ bufnr = buf, pos = { row, col } }))
	return node_within_type(node, "comment")
end

local patterns = { todo = "TODO" }
local extmarks_todo_opts = {
	ns = vim.g.dotfiles.todo_ns,
	hl_group = vim.g.dotfiles.todo_hl,
	predicate = hl_todo_predicate,
}

local function scan(buf, first, last)
	local lines = vim.api.nvim_buf_get_lines(buf, first, last, false)
	local iter = pmatches(lines, patterns.todo)
	set_extmarks(buf, function()
		local pos = iter()
		if pos then
			pos.row = pos.row + first
		end
		return pos
	end, extmarks_todo_opts)
end

vim.api.nvim_create_autocmd("FileType", {
	desc = "Initialize and watch TODO extmarks",
	group = vim.g.dotfiles.augroup,
	callback = function(args)
		local buf = args.buf
		scan(buf, 0, -1)
		if vim.b[buf].dotfiles_todo_attached then
			return
		end
		vim.b[buf].dotfiles_todo_attached = true
		vim.api.nvim_buf_attach(buf, false, {
			on_lines = function(_, b, _, first, _, last_new)
				vim.schedule(function()
					if vim.api.nvim_buf_is_valid(b) then
						scan(b, first, last_new)
					end
				end)
			end,
		})
	end,
})