-- lx-bundle:true

local core = require("core")
local command = require("core.command")

local arithmetic = {}
local preview_limit = 10
local expression_environment = {
	assert = assert,
	error = error,
	ipairs = ipairs,
	math = math,
	next = next,
	pairs = pairs,
	select = select,
	string = string,
	table = table,
	tonumber = tonumber,
	tostring = tostring,
	type = type,
	utf8 = utf8,
	_VERSION = _VERSION,
}
expression_environment._G = expression_environment

--- Compile one Lua expression with arithmetic selection variables as locals
-- @param expression string Lua expression
-- @return function|nil evaluator
-- @return string|nil err
local function compile_expression(expression)
	-- Pattern breakdown: '^' and '$' require the entire expression to contain only '%s' whitespace
	if type(expression) ~= "string" or expression:match("^%s*$") then
		return nil, "Expression is empty"
	end

	local source = table.concat({
		"return function(i, n, x, s, row, col, total)",
		"	return (" .. expression .. ")",
		"end",
	}, "\n")
	local chunk, compile_error = load(source, "Arithmetic", "t", expression_environment)

	if not chunk then
		return nil, tostring(compile_error)
	end

	local ok, evaluator = pcall(chunk)

	if not ok then
		return nil, tostring(evaluator)
	end

	return evaluator
end

--- Evaluate one expression independently for every target
-- Available locals include selection indexes, text, document position, and total target count
-- Standard Lua globals such as 'math', 'string', and 'table' remain available
-- @param expression string Lua expression
-- @param targets table selected targets in document order
-- @param total number complete target count
-- @return table|nil results string results
-- @return string|nil err
function arithmetic.evaluate(expression, targets, total)
	local evaluator, compile_error = compile_expression(expression)

	if not evaluator then
		return nil, compile_error
	end

	local results = {}

	for index = 1, #targets do
		local target = targets[index]
		local selected_text = target.text
		local ok, result = pcall(
			evaluator,
			index - 1,
			index,
			tonumber(selected_text),
			selected_text,
			target.line1,
			target.col1,
			total
		)

		if not ok then
			return nil, string.format("Selection %d: %s", index, tostring(result))
		end

		results[index] = result == nil and "" or tostring(result)
	end

	return results
end

--- Return the exclusive end column of one logical line
-- @param doc table active text document
-- @param line number line index
-- @return number end_col
local function get_line_end_col(doc, line)
	local text = doc.lines[line]

	return text:sub(-1) == "\n" and #text or #text + 1
end

--- Split active selections into ordered per-line arithmetic targets
-- Newline characters remain outside the targets so expressions cannot remove line boundaries
-- @param doc table active text document
-- @param limit number|nil maximum returned targets
-- @return table targets
-- @return number target_count
-- @return table groups target ranges belonging to each original selection
-- @return number primary_selection_index
local function get_line_targets(doc, limit)
	local ordered_selections, primary_selection_index = doc:get_selections(true)
	local original_selections = doc:get_selections()
	local targets = {}
	local groups = {}
	local target_count = 0

	for selection_index = 1, #ordered_selections do
		local selection = ordered_selections[selection_index]
		local original_selection = original_selections[selection_index]
		local line1, col1, line2, col2 = table.unpack(selection)

		if line2 > line1 and col2 == 1 then
			line2 = line2 - 1
			col2 = get_line_end_col(doc, line2)
		end

		local first_target_index = target_count + 1
		local selection_target_count = line2 - line1 + 1
		target_count = target_count + selection_target_count
		-- Lx stores the live caret in the first selection endpoint before sorting
		groups[selection_index] = {
			line1 = line1,
			col1 = col1,
			line2 = line2,
			col2 = col2,
			first_target_index = first_target_index,
			last_target_index = target_count,
			caret_at_start = original_selection[1] < original_selection[3]
				or original_selection[1] == original_selection[3]
				and original_selection[2] < original_selection[4],
			id = selection.id,
		}

		for line = line1, line2 do
			if limit and #targets >= limit then
				break
			end

			local target_col1 = line == line1 and col1 or 1
			local target_col2 = line == line2 and col2 or get_line_end_col(doc, line)
			targets[#targets + 1] = {
				line1 = line,
				col1 = target_col1,
				line2 = line,
				col2 = target_col2,
				text = doc:get_text(line, target_col1, line, target_col2),
			}
		end
	end

	return targets, target_count, groups, primary_selection_index
end

--- Return the end position of inserted text from one result position
-- @param line number result start line
-- @param col number result start column
-- @param text string inserted text
-- @return number end_line
-- @return number end_col
local function get_result_end_position(line, col, text)
	local line_break_count = 0
	local trailing_length = 0

	for index = 1, #text do
		if text:byte(index) == 10 then
			line_break_count = line_break_count + 1
			trailing_length = 0
		else
			trailing_length = trailing_length + 1
		end
	end

	if line_break_count == 0 then
		return line, col + trailing_length
	end

	return line + line_break_count, trailing_length + 1
end

--- Apply one arithmetic expression to every active selection or caret
-- @param doc table active text document
-- @param expression string Lua expression
-- @return boolean changed
-- @return string|nil err
function arithmetic.apply(doc, expression)
	local targets, _, groups, primary_selection_index = get_line_targets(doc)
	local results, evaluation_error = arithmetic.evaluate(expression, targets, #targets)

	if not results then
		return false, evaluation_error
	end

	local replacements = {}
	local after_selections = {}
	local preceding

	for group_index = 1, #groups do
		local group = groups[group_index]
		local group_results = {}

		for target_index = group.first_target_index, group.last_target_index do
			group_results[#group_results + 1] = results[target_index]
		end

		local replacement_text = table.concat(group_results, "\n")
		local result_line1 = group.line1
		local result_col1 = group.col1

		if preceding then
			if group.line1 == preceding.line2 then
				result_line1 = preceding.result_line2
				result_col1 = preceding.result_col2 + group.col1 - preceding.col2
			else
				result_line1 = preceding.result_line2 + group.line1 - preceding.line2
			end
		end

		local result_line2, result_col2 = get_result_end_position(
			result_line1,
			result_col1,
			replacement_text
		)

		replacements[group_index] = {
			line1 = group.line1,
			col1 = group.col1,
			line2 = group.line2,
			col2 = group.col2,
			text = replacement_text,
		}

		local caret_line = group.caret_at_start and result_line1 or result_line2
		local caret_col = group.caret_at_start and result_col1 or result_col2
		after_selections[group_index] = {
			caret_line,
			caret_col,
			caret_line,
			caret_col,
			id = group.id,
		}
		preceding = {
			line2 = group.line2,
			col2 = group.col2,
			result_line2 = result_line2,
			result_col2 = result_col2,
		}
	end

	return doc:replace_ranges(replacements, after_selections, primary_selection_index)
end

--- Insert sequential numbers at the active selections or selected lines
-- @param doc table active text document
-- @param start any first sequence value
-- @param step any increment between sequence values
-- @return boolean changed
-- @return string|nil err
local function apply_numbering(doc, start, step)
	local starting_number = tonumber(start)
	local step_number = tonumber(step or 1)

	if not starting_number or starting_number ~= starting_number
		or starting_number == math.huge or starting_number == -math.huge
	then
		return false, "Starting number must be numeric"
	end
	if not step_number or step_number ~= step_number
		or step_number == math.huge or step_number == -math.huge
	then
		return false, "Step must be numeric"
	end

	step_number = math.max(1, step_number)

	return arithmetic.apply(doc, string.format(
		"(%s) + i * (%s)",
		tostring(starting_number),
		tostring(step_number)
	))
end

--- Keep preview results readable inside one command suggestion row
-- @param result string evaluated result
-- @return string preview_text
local function format_preview_result(result)
	-- Pattern breakdown: '\r', '\n', and '\t' each match one literal whitespace control
	result = result:gsub("\r", "\\r"):gsub("\n", "\\n"):gsub("\t", "\\t")

	if result == "" then
		return '""'
	end

	return result
end

--- Build live result preview rows for the first selections
-- @param doc table active text document
-- @param expression string Lua expression
-- @return table suggestions command prompt suggestions
local function get_preview_suggestions(doc, expression)
	-- Pattern breakdown: '^' and '$' require the entire expression to contain only '%s' whitespace
	if expression:match("^%s*$") then
		return {}
	end

	local targets, target_count = get_line_targets(doc, preview_limit)
	local results, evaluation_error = arithmetic.evaluate(expression, targets, target_count)

	if not results then
		return {
			{
				text = "Error: " .. evaluation_error,
				complete_text = expression,
			},
		}
	end

	local suggestions = {}

	for index = 1, #results do
		suggestions[index] = {
			text = format_preview_result(results[index]),
			complete_text = expression,
		}
	end

	if target_count > preview_limit then
		suggestions[#suggestions + 1] = {
			text = "etc.",
			complete_text = expression,
		}
	end

	-- Command suggestions draw their first item nearest the prompt, so reverse previews for top-down reading
	local display_suggestions = {}

	for index = #suggestions, 1, -1 do
		display_suggestions[#display_suggestions + 1] = suggestions[index]
	end

	return display_suggestions
end

--- Open the arithmetic prompt and keep failed expressions available for correction
-- @param doc table active text document
-- @param initial_expression string prefilled expression
-- @param initial_error string|nil submitted evaluation error
-- @param select_expression boolean select the prefilled expression
local function open_prompt(doc, initial_expression, initial_error, select_expression)
	local suggestions_ready = false
	local retry_error = initial_error
	local opened = core.command_view:enter({
		caption = "Lua Expression",
		submit = function(expression)
			local _, evaluation_error = arithmetic.apply(doc, expression)

			if evaluation_error then
				open_prompt(doc, expression, evaluation_error, false)
			end
		end,
		suggest = function(expression)
			if not suggestions_ready then
				return {}
			end

			if retry_error and expression == initial_expression then
				return {
					{
						text = "Error: " .. retry_error,
						complete_text = expression,
					},
				}
			end

			retry_error = nil

			return get_preview_suggestions(doc, expression)
		end,
	})

	if not opened then
		return
	end

	suggestions_ready = true
	core.command_view:set_text(initial_expression, select_expression)
	core.command_view:update_suggestions()
	core.command_view.last_change_id = core.command_view.doc:get_change_id()
end

--- Open a prompt for the increment used by sequential numbering
-- @param doc table active text document
-- @param starting_number number first sequence value
-- @param initial_value string prefilled step value
-- @param initial_error string|nil submitted value error
local function open_start_step_prompt(doc, starting_number, initial_value, initial_error)
	local opened = core.command_view:enter({
		caption = "Step",
		submit = function(value)
			local _, numbering_error = apply_numbering(doc, starting_number, value)

			if numbering_error then
				open_start_step_prompt(doc, starting_number, value, numbering_error)
			end
		end,
		suggest = function(value)
			if initial_error and value == initial_value then
				return {
					{
						text = "Error: " .. initial_error,
						complete_text = value,
					},
				}
			end

			return {}
		end,
	})

	if not opened then
		return
	end

	core.command_view:set_text(initial_value, true)
	core.command_view:update_suggestions()
end

--- Open a prompt for the first value used by sequential numbering
-- @param doc table active text document
-- @param initial_value string prefilled starting value
-- @param initial_error string|nil submitted value error
local function open_starting_number_prompt(doc, initial_value, initial_error)
	local opened = core.command_view:enter({
		caption = "Starting number",
		submit = function(value)
			local starting_number = tonumber(value)

			if not starting_number or starting_number ~= starting_number
				or starting_number == math.huge or starting_number == -math.huge
			then
				open_starting_number_prompt(doc, value, "Starting number must be numeric")

				return
			end

			open_start_step_prompt(doc, starting_number, "1", nil)
		end,
		suggest = function(value)
			if initial_error and value == initial_value then
				return {
					{
						text = "Error: " .. initial_error,
						complete_text = value,
					},
				}
			end

			return {}
		end,
	})

	if not opened then
		return
	end

	core.command_view:set_text(initial_value)
	core.command_view:update_suggestions()
end

command.register({
	["arithmetic:apply"] = {
		requires = { "doc", "edit" },
		perform = function()
			local doc = core.active_view.doc
			local _, target_count = get_line_targets(doc, 1)
			local default_expression = target_count > 1 and "i + 1" or ""
			open_prompt(doc, default_expression, nil, default_expression ~= "")
		end,
	},
	["arithmetic:number"] = {
		requires = { "doc", "edit" },
		perform = function(_, args)
			local start = type(args) == "table" and args.start or 1
			local step = type(args) == "table" and args.step or 1
			local changed, numbering_error = apply_numbering(core.active_view.doc, start, step)

			if numbering_error then
				core.error("Arithmetic numbering failed: %s", numbering_error)
			end

			return changed
		end,
	},
	["arithmetic:number-start-step"] = {
		requires = { "doc", "edit" },
		perform = function()
			open_starting_number_prompt(core.active_view.doc, "", nil)
		end,
	},
})

return arithmetic
