-- lx-bundle:true

--- TeX-like math support owned only by the optional Markdown preview
-- @module plugins.mdpreview.math

local layout = require("plugins.mdpreview.math.layout")
local parser = require("plugins.mdpreview.math.parser")

local markdown_math = {}

local NUMBERED_ENVIRONMENTS = {
	equation = "single",
	align = "rows",
	alignat = "rows",
	gather = "rows",
	multline = "single",
}

--- Find the numbered row node parsed from one complete environment
-- @param expression table parsed expression
-- @param environment_name string expected environment name
-- @return table|nil rows matching row node
local function find_numbered_rows(expression, environment_name)
	for child_index = 1, #(expression.children or {}) do
		local child = expression.children[child_index]

		if child.kind == "rows" and child.environment == environment_name then
			return child
		end
	end

	return nil
end

--- Number one complete environment and attach its labels to the matching row
-- @param current table collected environment source and location
-- @param number number previous equation number
-- @param references table label lookup to update
-- @param equation_tags table source-line tag lookup to update
-- @return number number last assigned equation number
local function number_environment(current, number, references, equation_tags)
	local expression = parser.parse(table.concat(current.lines, "\n"))
	local rows = find_numbered_rows(expression, current.name)
	local mode = NUMBERED_ENVIRONMENTS[current.name]
	local row_count = mode == "rows" and rows and #rows.rows or 1
	local tags = {}

	for row_index = 1, math.max(1, row_count) do
		-- Suppressed rows neither show nor consume an equation number
		if not (rows and rows.suppressed_tags[row_index]) then
			number = number + 1
			tags[row_index] = tostring(number)
		end
	end
	for row_index, tag in pairs(rows and rows.tags or {}) do
		tags[row_index] = tag
	end

	for row_index, labels in pairs(rows and rows.labels or {}) do
		for label_index = 1, #labels do
			references[labels[label_index]] = tags[row_index] or tags[#tags]
		end
	end
	equation_tags[current.start_line] = tags

	return number
end

--- Collect equation labels before lazy preview layout begins
-- @param lines table complete Markdown document lines
-- @param strip_line_ending function line cleanup helper
-- @param read_fence function|nil Markdown fence reader
-- @return table references label to number or explicit tag
-- @return table equation_tags source-line visible tag lists
function markdown_math.scan_references(lines, strip_line_ending, read_fence)
	local references = {}
	local equation_tags = {}
	local number = 0
	local current
	local active_fence

	for line_index = 1, #lines do
		local line = strip_line_ending(lines[line_index])
		local fence = read_fence and read_fence(line)

		if active_fence then
			if fence
				and fence:sub(1, 1) == active_fence:sub(1, 1)
				and #fence >= #active_fence
			then
				active_fence = nil
			end
		elseif fence then
			active_fence = fence
		else
			if current then
				current.lines[#current.lines + 1] = line
			else
				for name in pairs(NUMBERED_ENVIRONMENTS) do
					local marker = "\\begin{" .. name .. "}"
					local start = line:find(marker, 1, true)

					if start then
						current = {
							name = name,
							start_line = line_index,
							lines = { line:sub(start) },
						}
						break
					end
				end
			end

			if current and line:find("\\end{" .. current.name .. "}", 1, true) then
				number = number_environment(current, number, references, equation_tags)
				current = nil
			end
		end
	end

	return references, equation_tags
end

--- Find display-dollar blocks while ignoring every fenced code region
-- @param lines table document lines
-- @param strip_line_ending function line cleanup helper
-- @param read_fence function Markdown fence reader
-- @return table contexts source-line math lookup
function markdown_math.scan_blocks(lines, strip_line_ending, read_fence)
	local contexts = {}
	local active_fence
	local active_block

	for line_index = 1, #lines do
		local line = strip_line_ending(lines[line_index])
		local fence = read_fence(line)

		if active_fence then
			local remainder = fence and line:gsub("^%s*" .. fence, "", 1) or line

			if fence
				and fence:sub(1, 1) == active_fence:sub(1, 1)
				and #fence >= #active_fence
				and remainder:match("^%s*$")
			then
				active_fence = nil
			end
		elseif fence then
			active_fence = fence
		elseif line:match("^%s*%$%$%s*$") then
			-- Pattern breakdown: surrounding whitespace is allowed around an otherwise bare '$$'
			if active_block then
				contexts[line_index] = { delimiter = true, block = active_block }
				active_block.last_line = line_index
				active_block = nil
			else
				active_block = { lines = {}, first_line = line_index }
				contexts[line_index] = { delimiter = true, block = active_block }
			end
		elseif active_block then
			active_block.lines[#active_block.lines + 1] = line
			contexts[line_index] = {
				block = active_block,
				first = #active_block.lines == 1,
			}
		end
	end

	if active_block then
		-- An unmatched opening delimiter remains ordinary Markdown text
		for line_index = active_block.first_line, #lines do
			contexts[line_index] = nil
		end
	end

	return contexts
end

--- Read an inline dollar expression using conservative Pandoc-style boundaries
-- @param text string complete Markdown inline source
-- @param position number opening-dollar byte position
-- @return string|nil source expression without delimiters
-- @return number|nil next_position first byte after closing delimiter
function markdown_math.read_inline(text, position)
	if text:sub(position, position) ~= "$"
		or text:sub(position + 1, position + 1) == "$"
		or text:sub(position + 1, position + 1):match("%s")
	then
		return nil
	end

	local search_position = position + 1

	while search_position <= #text do
		local close_position = text:find("$", search_position, true)

		if not close_position then
			return nil
		end

		local before = text:sub(close_position - 1, close_position - 1)
		local after = text:sub(close_position + 1, close_position + 1)
		local escaped = close_position > 1 and text:sub(close_position - 1, close_position - 1) == "\\"

		if not escaped and not before:match("%s") and not after:match("%d") then
			return text:sub(position + 1, close_position - 1), close_position + 1
		end

		search_position = close_position + 1
	end

	return nil
end

--- Parse and measure one math expression
-- @param source string TeX-like source
-- @param font_provider function font resolver
-- @param options table|nil layout measurements
-- @return table box measured drawing box
function markdown_math.build(source, font_provider, options)
	return layout.build(
		parser.parse(
			source,
			options and options.references,
			options and options.equation_tags
		),
		font_provider,
		options
	)
end

markdown_math.draw = layout.draw
markdown_math.get_link_at = layout.get_link_at
markdown_math.parse = parser.parse

return markdown_math
