-- lx-bundle:true
local core = require("core")
local command = require("core.command")
local hook = require("core.hook")
local common = require("core.common")
local system = require("system")

-- If selection ends at column 1 on a later line, treat that line as unselected
local function effective_last_line(line1, line2, col2)
	if line2 > line1 and col2 == 1 then
		return line2 - 1
	end
	return line2
end

-- Returns selected line range or full-doc line range
local function get_target_line_range(doc)
	if doc:has_selection() then
		local line1, _, line2, col2 = doc:get_selection(true)
		local last_line = effective_last_line(line1, line2, col2)
		if last_line < line1 then
			last_line = line1
		end
		return line1, last_line, true
	end
	return 1, #doc.lines, false
end

-- Applies a text transform to selected lines or whole document
local function apply_lines_transform(doc, transform_fn)
	if not doc or not transform_fn then
		return
	end
	local line1, line2, has_selection = get_target_line_range(doc)
	if has_selection then
		doc:replace_range(line1, 1, line2, math.huge, function(text)
			return transform_fn(text)
		end)
		return
	end
	doc:replace(function(text)
		return transform_fn(text)
	end)
end

-- Splits text into logical lines preserving all empty lines
local function split_lines(text)
	local lines = {}
	local start_pos = 1
	while true do
		local nl = text:find("\n", start_pos, true)
		if not nl then
			lines[#lines + 1] = text:sub(start_pos)
			break
		end
		lines[#lines + 1] = text:sub(start_pos, nl - 1)
		start_pos = nl + 1
	end
	return lines
end

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

-- Applies an array-of-lines transform to selected lines or whole document
local function apply_line_list_transform(doc, list_fn)
	apply_lines_transform(doc, function(text)
		local lines = split_lines(text)
		local out = list_fn(lines)
		return join_lines(out)
	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.redo_stack = { idx = 1 }
	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

-- Trims trailing whitespace on every line
local function trim_trailing_whitespaces(doc, opts)
	if not doc then
		return
	end
	opts = opts or {}
	local line1, line2 = get_target_line_range(doc)
	local edit_time = nil
	for line = line1, line2 do
		if line ~= opts.skip_line then
			edit_time = replace_line_content(doc, line, common.rtrim, edit_time)
		end
	end
end

-- Trims leading whitespace on every line
local function trim_leading_whitespaces(doc)
	if not doc then
		return
	end
	local line1, line2 = get_target_line_range(doc)
	local edit_time = nil
	for line = line1, line2 do
		edit_time = replace_line_content(doc, line, common.ltrim, edit_time)
	end
end

-- Trims leading and trailing whitespace on every line
local function trim_whitespaces(doc)
	if not doc then
		return
	end
	local line1, line2 = get_target_line_range(doc)
	local edit_time = nil
	for line = line1, line2 do
		edit_time = replace_line_content(doc, line, common.trim, edit_time)
	end
end

-- Returns a comparator for sort mode
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 = a:lower()
		local b_folded = b:lower()
		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

-- Sorts lines then removes adjacent duplicates in selected compare mode
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 line:lower()
			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

-- Removes all empty or whitespace-only lines
local function remove_empty_lines(doc)
	if not doc then
		return
	end
	local line1, line2 = get_target_line_range(doc)
	local edit_time = nil
	for line = line2, 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

-- Collapses adjacent empty or whitespace-only lines into one empty line
local function collapse_empty_lines(doc)
	if not doc then
		return
	end
	local line1, line2 = get_target_line_range(doc)
	local edit_time = nil
	local prev_was_empty = false
	for line = line2, 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 prev_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
		prev_was_empty = is_empty
	end
end

command.add({
	["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: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,
			})
		end
	end,
})
