-- lx-bundle:true

-- This file is part of Lx, created by Liqube <www.liqube.com>
-- Copyright Liqube Audio. All rights reserved.
-- Lx Interim License. See LICENSE for details.

local core = require("core")
local command = require("core.command")
local hook = require("core.hook")
local common = require("core.common")
local text_line = require("core.text.line")
local system = require("system")
local native_text = require("text") -- Split large line transformations without a Lua substring loop
local utf8 = require("utf8")

--- Build the same normalized caseless key used by editor text search
-- @param text string Text to fold
-- @return string key Normalized caseless text
local function caseless_key(text)
	local normalized = utf8.normalize_nfc(text)
	local folded = utf8.fold(normalized)
	---@cast folded string

	return utf8.normalize_nfc(folded)
end

--- Apply one text transform to the whole document or each selected line block
-- @param doc table text document
-- @param transform_fn function text transform
-- @param opts table|nil transform options
-- @return boolean changed whether text changed
-- @return string|nil transform_error validation failure
local function apply_lines_transform(doc, transform_fn, opts)
	if not doc or not transform_fn then
		return false
	end

	opts = opts or {}

	if opts.whole_document or not doc:has_any_selection() then
		local transform_error
		local changed = doc:replace(function(text)
			local result, result_error = transform_fn(text)
			transform_error = result_error

			return result or text
		end)

		return changed, transform_error
	end

	local ranges = text_line.get_target_ranges(doc)
	local replacements = {}

	-- Build every result before editing so one invalid block cannot leave a partial change
	for index = 1, #ranges do
		local range = ranges[index]
		local line2, col2 = text_line.get_block_end(doc, range.line2)
		local text = doc:get_text(range.line1, 1, line2, col2)
		local result, transform_error = transform_fn(text)

		if not result then
			return false, transform_error
		end

		local line1 = range.line1
		local col1 = 1

		-- Removing the final selected block also removes the LF before it
		if result == "" and range.line2 == #doc.lines and line1 > 1 then
			line1 = line1 - 1
			col1 = #doc.lines[line1]
		end

		replacements[index] = {
			line1 = line1,
			col1 = col1,
			line2 = line2,
			col2 = col2,
			text = result,
		}
	end

	return doc:replace_ranges(replacements), nil
end

--- Split text into logical lines while preserving every empty line
-- @param text string normalized document text without its final newline
-- @return table lines logical lines without line endings
local function split_lines(text)
	return native_text.split_lines(text, { keep_eol = false })
end

-- Joins logical lines back to text
local function join_lines(lines)
	return table.concat(lines, "\n")
end

--- Apply an array-of-lines transform while preserving the final newline
-- @param doc table text document
-- @param list_fn function line-list transform
local function apply_line_list_transform(doc, list_fn)
	return apply_lines_transform(doc, function(text)
		local has_final_newline = text:sub(-1) == "\n"
		local content = has_final_newline and text:sub(1, -2) or text
		local lines = split_lines(content)
		local out, transform_error = list_fn(lines)

		if not out then
			return nil, transform_error
		end

		local result = join_lines(out)

		if has_final_newline and #out > 0 then
			result = result .. "\n"
		end

		return result
	end)
end

-- Extracts line content without trailing LF and returns content with line ending
local function split_line_content_and_eol(line_text)
	if line_text:sub(-1) == "\n" then
		return line_text:sub(1, -2), "\n"
	end
	return line_text, ""
end

-- Starts one grouped undo edit batch
local function begin_grouped_edit(doc)
	doc:clear_redo()
	return system.get_time()
end

-- Returns end column for logical line content, excluding LF when present
local function get_line_content_end_col(doc, line)
	local text = doc.lines[line]
	if line == #doc.lines and text:sub(-1) ~= "\n" then
		return #text + 1
	end
	return #text
end

-- Applies one per-line replace when text changed, preserving original line ending
local function replace_line_content(doc, line, transform_fn, edit_time)
	local old_content = select(1, split_line_content_and_eol(doc.lines[line]))
	local new_content = transform_fn(old_content)
	if old_content == new_content then
		return edit_time
	end
	if not edit_time then
		edit_time = begin_grouped_edit(doc)
	end
	doc:raw_replace(
		line,
		1,
		line,
		get_line_content_end_col(doc, line),
		new_content,
		doc.undo_stack,
		edit_time,
		old_content
	)
	return edit_time
end

-- Deletes one full line while keeping document valid when it has a single line
local function remove_full_line(doc, line, edit_time)
	if not edit_time then
		edit_time = begin_grouped_edit(doc)
	end
	if #doc.lines == 1 then
		local old_line = doc.lines[1]
		if old_line ~= "" then
			doc:raw_replace(1, 1, 1, math.huge, "", doc.undo_stack, edit_time, old_line)
		end
		return edit_time
	end
	if line < #doc.lines then
		doc:raw_remove(line, 1, line + 1, 1, doc.undo_stack, edit_time)
	else
		doc:raw_remove(line, 1, line, get_line_content_end_col(doc, line), doc.undo_stack, edit_time)
	end
	return edit_time
end

--- Trim trailing whitespace on every targeted line
-- @param doc table text document
-- @param opts table|nil trim options
local function trim_trailing_whitespaces(doc, opts)
	if not doc then
		return
	end

	opts = opts or {}
	local ranges = text_line.get_target_ranges(doc, opts.whole_document)
	local edit_time = nil

	for index = 1, #ranges do
		local range = ranges[index]

		for line = range.line1, range.line2 do
			if line ~= opts.skip_line then
				edit_time = replace_line_content(doc, line, common.rtrim, edit_time)
			end
		end
	end
end

--- Trim leading whitespace on every targeted line
-- @param doc table text document
local function trim_leading_whitespaces(doc)
	if not doc then
		return
	end

	local ranges = text_line.get_target_ranges(doc)
	local edit_time = nil

	for index = 1, #ranges do
		local range = ranges[index]

		for line = range.line1, range.line2 do
			edit_time = replace_line_content(doc, line, common.ltrim, edit_time)
		end
	end
end

--- Trim leading and trailing whitespace on every targeted line
-- @param doc table text document
local function trim_whitespaces(doc)
	if not doc then
		return
	end

	local ranges = text_line.get_target_ranges(doc)
	local edit_time = nil

	for index = 1, #ranges do
		local range = ranges[index]

		for line = range.line1, range.line2 do
			edit_time = replace_line_content(doc, line, common.trim, edit_time)
		end
	end
end

--- Return a comparator for one line sort mode
-- @param case_sensitive boolean compare original text when true
-- @return function comparator
local function build_sort_comparator(case_sensitive)
	if case_sensitive then
		return function(a, b)
			return a < b
		end
	end
	return function(a, b)
		local a_folded = caseless_key(a)
		local b_folded = caseless_key(b)
		if a_folded == b_folded then
			return a < b
		end
		return a_folded < b_folded
	end
end

-- Sorts lines using case-sensitive or case-insensitive mode
local function sort_lines(doc, case_sensitive)
	apply_line_list_transform(doc, function(lines)
		table.sort(lines, build_sort_comparator(case_sensitive))
		return lines
	end)
end

--- Sort lines then remove adjacent duplicates
-- @param doc table text document
-- @param case_sensitive boolean compare original text when true
local function unique_lines(doc, case_sensitive)
	apply_line_list_transform(doc, function(lines)
		table.sort(lines, build_sort_comparator(case_sensitive))
		local out = {}
		local prev_key = nil
		for i = 1, #lines do
			local line = lines[i]
			local key = case_sensitive and line or caseless_key(line)
			if i == 1 or key ~= prev_key then
				out[#out + 1] = line
				prev_key = key
			end
		end
		return out
	end)
end

-- Reverses the line order
local function reverse_lines(doc)
	apply_line_list_transform(doc, function(lines)
		local i, j = 1, #lines
		while i < j do
			lines[i], lines[j] = lines[j], lines[i]
			i = i + 1
			j = j - 1
		end
		return lines
	end)
end

-- Randomizes line order using in-place Fisher-Yates
local function shuffle_lines(doc)
	apply_line_list_transform(doc, function(lines)
		for i = #lines, 2, -1 do
			local j = math.random(i)
			lines[i], lines[j] = lines[j], lines[i]
		end
		return lines
	end)
end

--- Check that a number is neither infinite nor NaN
-- @param value number value to check
-- @return boolean finite
local function is_finite_number(value)
	return value == value and value ~= math.huge and value ~= -math.huge
end

--- Read and validate one numeric command value
-- @param value any submitted value
-- @param label string user-facing value name
-- @param minimum number lowest accepted value
-- @param maximum number|nil highest accepted value
-- @param integer boolean require a whole number
-- @return number|nil parsed_value
-- @return string|nil value_error
local function parse_number(value, label, minimum, maximum, integer)
	local parsed = tonumber(value)

	if not parsed or not is_finite_number(parsed) then
		return nil, label .. " must be numeric"
	end
	if integer and parsed % 1 ~= 0 then
		return nil, label .. " must be a whole number"
	end
	if parsed < minimum or (maximum and parsed > maximum) then
		local range = maximum and minimum .. ".." .. maximum or minimum .. " or greater"

		return nil, label .. " must be " .. range
	end

	return parsed
end

--- Reject misspelled or unsupported named command arguments
-- @param args table submitted arguments
-- @param allowed table allowed argument names
-- @return string|nil argument_error
local function validate_argument_names(args, allowed)
	for name in pairs(args) do
		if not allowed[name] then
			return "Unknown argument: " .. tostring(name)
		end
	end

	return nil
end

--- Open a numeric prompt and keep an invalid value available for correction
-- @param caption string prompt caption
-- @param initial_value string prefilled value
-- @param initial_error string|nil validation error
-- @param parse_fn function numeric parser
-- @param submit_fn function receives the parsed value
local function open_number_prompt(caption, initial_value, initial_error, parse_fn, submit_fn)
	local opened = core.command_view:enter({
		caption = caption,
		submit = function(text)
			local value, value_error = parse_fn(text)

			if value_error then
				open_number_prompt(caption, text, value_error, parse_fn, submit_fn)

				return
			end

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

			return {}
		end,
	})

	if not opened then
		return
	end

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

--- Keep one repeating block of lines and remove the next block
-- @param lines table source lines
-- @param keep number lines kept per cycle
-- @param remove number lines removed per cycle
-- @return table result
local function stride_line_list(lines, keep, remove)
	local result = {}
	local cycle = keep + remove

	for index = 1, #lines do
		if (index - 1) % cycle < keep then
			result[#result + 1] = lines[index]
		end
	end

	return result
end

--- Keep an evenly spaced fraction of source lines
-- @param lines table source lines
-- @param factor number fraction kept from 0 through 1
-- @return table result
local function reduce_line_list(lines, factor)
	local count = math.floor(#lines * factor + 0.5)

	if count == 0 then
		return {}
	end
	if count >= #lines then
		return lines
	end
	if count == 1 then
		return { lines[1] }
	end

	local result = {}

	for index = 1, count do
		local source_index = math.floor((index - 1) * (#lines - 1) / (count - 1) + 0.5) + 1
		result[index] = lines[source_index]
	end

	return result
end

--- Sample source lines without replacement and preserve their original order
-- @param lines table source lines
-- @param count number lines to keep
-- @return table|nil result
-- @return string|nil sample_error
local function sample_line_list(lines, count)
	if count > #lines then
		return nil, "Sample count exceeds the number of target lines"
	end

	local indexes = {}

	-- Reservoir sampling avoids copying and shuffling every source line
	for index = 1, #lines do
		if index <= count then
			indexes[index] = index
		else
			local replacement = math.random(index)

			if replacement <= count then
				indexes[replacement] = index
			end
		end
	end

	table.sort(indexes)

	local result = {}

	for index = 1, #indexes do
		result[index] = lines[indexes[index]]
	end

	return result
end

--- Show a transform error without turning it into a critical failure
-- @param transform_error string|nil reported failure
-- @return boolean succeeded
local function report_transform_error(transform_error)
	if transform_error then
		core.warning("%s", transform_error)

		return false
	end

	return true
end

--- Run line stride directly or collect its values interactively
-- @param doc table active text document
-- @param args table|nil named command arguments
local function perform_stride(doc, args)
	local has_arguments = type(args) == "table" and next(args) ~= nil

	if has_arguments then
		local argument_error = validate_argument_names(args, { keep = true, remove = true })
		local keep, keep_error = parse_number(args.keep, "Keep", 1, nil, true)
		local remove, remove_error = parse_number(args.remove or 1, "Remove", 1, nil, true)

		if argument_error or keep_error or remove_error then
			core.warning("%s", argument_error or keep_error or remove_error)

			return false
		end

		return apply_line_list_transform(doc, function(lines)
			return stride_line_list(lines, keep, remove)
		end)
	end

	open_number_prompt("Keep N lines", "3", nil, function(value)
		return parse_number(value, "Keep", 1, nil, true)
	end, function(keep)
		open_number_prompt("Remove N lines", "1", nil, function(value)
			return parse_number(value, "Remove", 1, nil, true)
		end, function(remove)
			apply_line_list_transform(doc, function(lines)
				return stride_line_list(lines, keep, remove)
			end)
		end)
	end)

	return true
end

--- Run even line reduction directly or collect its factor interactively
-- @param doc table active text document
-- @param args table|nil named command arguments
local function perform_reduce(doc, args)
	local has_arguments = type(args) == "table" and next(args) ~= nil

	if has_arguments then
		local argument_error = validate_argument_names(args, { factor = true, percent = true })
		local has_factor = args.factor ~= nil
		local has_percent = args.percent ~= nil

		if argument_error or has_factor == has_percent then
			core.warning("%s", argument_error or "Use either factor or percent")

			return false
		end

		local value = has_factor and args.factor or args.percent
		local maximum = has_factor and 1 or 100
		local label = has_factor and "Factor" or "Percent"
		local factor, factor_error = parse_number(value, label, 0, maximum, false)

		if factor_error then
			core.warning("%s", factor_error)

			return false
		end

		if has_percent then
			factor = factor / 100
		end

		return apply_line_list_transform(doc, function(lines)
			return reduce_line_list(lines, factor)
		end)
	end

	open_number_prompt("Keep what fraction of lines (0..1)", "0.5", nil, function(value)
		return parse_number(value, "Factor", 0, 1, false)
	end, function(factor)
		apply_line_list_transform(doc, function(lines)
			return reduce_line_list(lines, factor)
		end)
	end)

	return true
end

--- Run random line sampling directly or collect its count interactively
-- @param doc table active text document
-- @param args table|nil named command arguments
local function perform_sample(doc, args)
	local has_arguments = type(args) == "table" and next(args) ~= nil

	if has_arguments then
		local argument_error = validate_argument_names(args, { count = true, percent = true })
		local has_count = args.count ~= nil
		local has_percent = args.percent ~= nil

		if argument_error or has_count == has_percent then
			core.warning("%s", argument_error or "Use either count or percent")

			return false
		end

		local maximum = has_percent and 100 or nil
		local label = has_percent and "Percent" or "Count"
		local value, value_error = parse_number(
			has_percent and args.percent or args.count,
			label,
			0,
			maximum,
			not has_percent
		)

		if value_error then
			core.warning("%s", value_error)

			return false
		end

		local _, transform_error = apply_line_list_transform(doc, function(lines)
			local count = has_percent and math.floor(#lines * value / 100 + 0.5) or value

			return sample_line_list(lines, count)
		end)

		return report_transform_error(transform_error)
	end

	open_number_prompt("Sample N lines", "10", nil, function(value)
		return parse_number(value, "Count", 0, nil, true)
	end, function(count)
		local _, transform_error = apply_line_list_transform(doc, function(lines)
			return sample_line_list(lines, count)
		end)

		report_transform_error(transform_error)
	end)

	return true
end

--- Remove every empty or whitespace-only targeted line
-- @param doc table text document
local function remove_empty_lines(doc)
	if not doc then
		return
	end

	local ranges = text_line.get_target_ranges(doc)
	local edit_time = nil

	-- Work from the bottom so deleted lines cannot shift a later selected block
	for index = #ranges, 1, -1 do
		local range = ranges[index]

		for line = range.line2, range.line1, -1 do
			local content = split_line_content_and_eol(doc.lines[line])

			if common.trim(content) == "" then
				edit_time = remove_full_line(doc, line, edit_time)
			end
		end
	end
end

--- Collapse adjacent empty targeted lines into one empty line per selected block
-- @param doc table text document
local function collapse_empty_lines(doc)
	if not doc then
		return
	end

	local ranges = text_line.get_target_ranges(doc)
	local edit_time = nil

	for index = #ranges, 1, -1 do
		local range = ranges[index]
		local previous_was_empty = false

		for line = range.line2, range.line1, -1 do
			local content = select(1, split_line_content_and_eol(doc.lines[line]))
			local is_empty = common.trim(content) == ""

			if is_empty then
				if previous_was_empty then
					edit_time = remove_full_line(doc, line, edit_time)
				else
					edit_time = replace_line_content(doc, line, function()
						return ""
					end, edit_time)
				end
			end

			previous_was_empty = is_empty
		end
	end
end

command.register({
	["lines:trim-whitespaces"] = {
		requires = { "doc", "edit" },
		perform = function()
			trim_whitespaces(core.active_view.doc)
		end,
	},
	["lines:trim-leading-whitespaces"] = {
		requires = { "doc", "edit" },
		perform = function()
			trim_leading_whitespaces(core.active_view.doc)
		end,
	},
	["lines:trim-trailing-whitespaces"] = {
		requires = { "doc", "edit" },
		perform = function()
			trim_trailing_whitespaces(core.active_view.doc)
		end,
	},
	["lines:sort"] = {
		requires = { "doc", "edit" },
		perform = function()
			sort_lines(core.active_view.doc, true)
		end,
	},
	["lines:sort-caseless"] = {
		requires = { "doc", "edit" },
		perform = function()
			sort_lines(core.active_view.doc, false)
		end,
	},
	["lines:unique"] = {
		requires = { "doc", "edit" },
		perform = function()
			unique_lines(core.active_view.doc, true)
		end,
	},
	["lines:unique-caseless"] = {
		requires = { "doc", "edit" },
		perform = function()
			unique_lines(core.active_view.doc, false)
		end,
	},
	["lines:reverse"] = {
		requires = { "doc", "edit" },
		perform = function()
			reverse_lines(core.active_view.doc)
		end,
	},
	["lines:shuffle"] = {
		requires = { "doc", "edit" },
		perform = function()
			shuffle_lines(core.active_view.doc)
		end,
	},
	["lines:stride"] = {
		requires = { "doc", "edit" },
		perform = function(_, args)
			return perform_stride(core.active_view.doc, args)
		end,
	},
	["lines:reduce"] = {
		requires = { "doc", "edit" },
		perform = function(_, args)
			return perform_reduce(core.active_view.doc, args)
		end,
	},
	["lines:sample"] = {
		requires = { "doc", "edit" },
		perform = function(_, args)
			return perform_sample(core.active_view.doc, args)
		end,
	},
	["lines:remove-empty"] = {
		requires = { "doc", "edit" },
		perform = function()
			remove_empty_lines(core.active_view.doc)
		end,
	},
	["lines:collapse-empty"] = {
		requires = { "doc", "edit" },
		perform = function()
			collapse_empty_lines(core.active_view.doc)
		end,
	},
})

-- Trim trailing whitespace just before save writes to disk
hook.add("doc.save", {
	id = "trim-whitespace",
	priority = 10,
	before = function(event)
		if event and event.doc then
			local _, _, cursor_line, cursor_col = event.doc:get_selection()
			local logical_line = select(1, event.doc:get_logical_line(cursor_line))
			local is_cursor_at_line_end = (cursor_col == #logical_line + 1)
			trim_trailing_whitespaces(event.doc, {
				skip_line = is_cursor_at_line_end and cursor_line or nil,
				whole_document = true,
			})
		end
	end,
})
