-- SPDX-License-Identifier: GPL-3.0-or-later
-- © 2020 Georgi Kirilov

require("vis")
local vis = vis

local l = require("lpeg")

local M

local builtin_textobjects = {
	["["] = { "[" , "]" },
	["{"] = { "{" , "}" },
	["<"] = { "<" , ">" },
	["("] = { "(" , ")" },
	['"'] = { '"' , '"' },
	["'"] = { "'" , "'" },
	["`"] = { "`" , "`" },
}

local builtin_motions = {
	["["] = { ["("] = true, ["{"] = true },
	["]"] = { [")"] = true, ["}"] = true },
}

local aliases = {}
for key, pair in pairs(builtin_textobjects) do aliases[pair[2]] = key ~= pair[2] and pair or nil end
for alias, pair in pairs(aliases) do builtin_textobjects[alias] = pair end
for alias, key in pairs({
	B = "{",
	b = "(",
}) do builtin_textobjects[alias] = builtin_textobjects[key] end

local function get_pair(key)
	return M.map[vis.win.syntax] and M.map[vis.win.syntax][key]
		or M.map[1] and M.map[1][key]
		or builtin_textobjects[key]
		or not key:match("%w") and {key, key}
end

local function at_pos(t, pos)
	if pos >= t[1] and pos < t[#t] then return t end
end

local function asymmetric(d, pos)
	local p
	local I = l.Cp()
	if #d == 1 then
		p = (d - l.B"\\") * I * ("\\" * l.P(1) + (l.P(1) - d))^0 * I * d
	else
		p = d * I * (l.P(1) - d)^0 * I * d
	end
	return l.Ct(I * p * I) * l.Cc(pos) / at_pos
end

local function symmetric(d1, d2, escaped, pos)
	local I = l.Cp()
	local skip = escaped and escaped + l.P(1) or l.P(1)
	return l.P{l.Ct(I * d1 * I * ((skip - d1 - d2) + l.V(1))^0 * I * d2 * I) * l.Cc(pos) / at_pos}
end

local function nth(t)
	if type(t) == "table" then
		local start, finish, c = 0, 0, vis.count or 1
		if #t == 5 then
			start, finish, c = nth(t[3])
		end
		if c then
			return {t[1], t[2]}, {t[#t - 1], t[#t]}, c > 1 and c - 1 or nil
		end
		return start, finish
	end
end

local function any_captures(_, position, t)
	if type(t) == "table" then
		return position, t
	end
end

local precedence = {
	[vis.lexers.COMMENT] = {vis.lexers.STRING},
	[vis.lexers.STRING] = {},
}

local function selection_range(pos)
	for selection in vis.win:selections_iterator() do
		if selection.pos == pos then
			return selection.range
		end
	end
end

local function past(_, position, pos)
	return position > pos
end

local function match_at(str, pattern, pos)
	local string_pos = pos + 1
	local I = l.Cp()
	local p = l.P{ I * (pattern/0) * I + 1 * (l.V(1) - l.Cmt(l.Cc(string_pos), past)) }
	local s, e = 1
	while true do
		s, e = p:match(str, s)
		if not s then return nil end
		if s <= string_pos and string_pos < e then
			return s - 1, e - 1
		end
		s = e
	end
end

local function escaping_context(data, range)
	if not vis.win.syntax then return {} end
	local rules = vis.lexers.lexers[vis.win.syntax]._RULES
	local p
	for _, name in ipairs({vis.lexers.COMMENT, vis.lexers.STRING}) do
		if rules[name] then
			p = p and p + rules[name] / 0 or rules[name] / 0
		end
	end
	if not p then return {} end
	if not range then return {escape = p} end    -- means we are retrying with a "fake" pos
	local e1 = {match_at(data, p, range.start)}
	local e2 = range.finish - range.start > 1 and {match_at(data, p, range.finish - 1)} or e1
	if #e1 == 0 and #e2 == 0 then return {escape = p} end
	if #e1 == 0 and #e2 > 0 then return {escape = p, newpos = range.start} end
	if #e2 == 0 and #e1 > 0 then return {escape = p, newpos = range.finish - 1} end
	p = nil
	local escaped_range = {e1[1] + 1, e1[2]}
	local escaped_data = data:sub(e1[1] + 1, e1[2])
	for _, level in ipairs({vis.lexers.COMMENT, vis.lexers.STRING}) do
		if l.match(rules[level] / 0 * -1, escaped_data) then
			for _, name in ipairs(precedence[level]) do
				if rules[name] then
					p = p and p + rules[name] / 0 or rules[name] / 0
				end
			end
			return {escape = p, range = escaped_range}
		end
	end
end

local function get_range(key, file_data, pos)
	local d = get_pair(key)
	if not d then return end
	local offsets, correction
	repeat
		local sel_range = selection_range(pos)
		local c = escaping_context(file_data, sel_range)
		local range = c.range or {1, #file_data}
		if c.newpos then
			pos = c.newpos
		else
			pos = pos - (range[1] - 1)
		end
		correction = range[1] - 1
		local p = d[1] ~= d[2] and symmetric(d[1], d[2], c.escape, pos + 1) or asymmetric(d[1], pos + 1)
		local skip = c.escape and c.escape + 1 or 1
		local data = c.range and file_data:sub(unpack(c.range)) or file_data
		local pattern = l.P{p + skip * (l.V(1) - l.Cmt(l.Cc(pos + 1), past))}
		local result
		local start = 1
		repeat
			result = pattern:match(data, start)
			if not result then break end
			start = result
		until type(result) == "table"
		if not result then
			pos = correction - 1
		else
			offsets = {nth(result)}
			offsets[3] = nil
		end
	until offsets or pos < 0
	if not offsets then return end
	for _, o in ipairs(offsets) do
		for i, v in ipairs(o) do
			o[i] = v - 1 + correction
		end
	end
	return unpack(offsets)
end

local function get_delimiters(key, pos)
	local d = get_pair(key)
	if not d or type(d[1]) == "string" and type(d[2]) == "string" then return d end
	local content = vis.win.file:content(0, vis.win.file.size)
	local start, finish = get_range(key, content, pos)
	if start and finish then
		return {vis.win.file:content(start[1], start[2] - start[1]), vis.win.file:content(finish[1], finish[2] - finish[1]), d[3], d[4]}
	elseif #d > 2 then
		return {nil, nil, d[3], d[4]}
	end
end

local function outer(content, pos)
	local start, finish = get_range(M.key, content, pos)
	if not (start and finish) then return end
	return start[1], finish[2]
end

local function inner(content, pos)
	local start, finish = get_range(M.key, content, pos)
	if not (start and finish) then return end
	return start[2], finish[1]
end

local function opening(content, pos)
	local start, _ = get_range(M.key, content, pos)
	if start then
		if pos == start[2] - 1 then
			start, _ = get_range(M.key, content, start[1] - 1)
		end
		if start then
			local exclusive = (vis.mode == vis.modes.OPERATOR_PENDING or vis.mode == vis.modes.VISUAL and pos < start[2] - 1) and 1 or 0
			return start[2] - 1 + exclusive
		end
	end
	return pos
end

local function closing(content, pos)
	local _, finish = get_range(M.key, content, pos)
	if finish then
		if pos == finish[1] then
			_, finish = get_range(M.key, content, finish[2])
		end
		if finish then
			local exclusive = (vis.mode == vis.modes.VISUAL and pos > finish[1]) and 1 or 0
			return finish[1] - exclusive
		end
	end
	return pos
end

local done_once

local function bail_early()
	if vis.count and vis.count > 1 then
		if done_once then
			done_once = nil
			return true
		else
			done_once = true
		end
	end
	return false
end

local function handler(func)
	return function(win, pos)
		if bail_early() then return pos end
		local content = win.file:content(0, win.file.size)
		return func(content, pos)
	end
end

local function new(execute, register, prefix, handler, help)
	local id = register(vis, handler)
	if id < 0 then
		return false
	end
	if prefix then
		local binding = function(keys)
			if #keys < 1 then return -1 end
			if #keys == 1 then
				M.key = keys
				execute(vis, id)
			end
			return #keys
		end
		if execute ~= vis.textobject then
			vis:map(vis.modes.NORMAL, prefix, binding, help)
		end
		vis:map(vis.modes.VISUAL, prefix, binding, help)
		vis:map(vis.modes.OPERATOR_PENDING, prefix, binding, help)
		local builtin = execute == vis.motion and builtin_motions[prefix] or builtin_textobjects
		for key, _ in pairs(builtin) do
			if execute ~= vis.textobject then
				vis:unmap(vis.modes.NORMAL, prefix..key)
			end
			vis:unmap(vis.modes.VISUAL, prefix..key)
			vis:unmap(vis.modes.OPERATOR_PENDING, prefix..key)
		end
	end
	return id
end

vis.events.subscribe(vis.events.INIT, function()
	M.motion = {
		opening = new(vis.motion, vis.motion_register, M.prefix.opening, handler(opening), "Move cursor to the beginning of a delimited block"),
		closing = new(vis.motion, vis.motion_register, M.prefix.closing, handler(closing), "Move cursor to the end of a delimited block"),
	}
	M.textobject = {
		inner = new(vis.textobject, vis.textobject_register, M.prefix.inner, handler(inner), "Delimited block (inner variant)"),
		outer = new(vis.textobject, vis.textobject_register, M.prefix.outer, handler(outer), "Delimited block (outer variant)"),
	}

	local tex_environment = {"\\begin{" * l.Cg(l.R("az", "AZ")^1, "t") * "}", l.Cmt("\\end{" * l.Cb("t") * l.C((1 - l.P"}")^1) * "}", function(_, _, c1, c2) return c1 == c2 end), {"\\begin{\xef\xbf\xbd}", "\\end{\xef\xbf\xbd}"}, "environment name"}
	local function is(_, _, v) return v ~= 1 end
	local tag_name = (l.S"_:" + l.R("az", "AZ")) * (l.R("az", "AZ", "09") + l.S"_:.-")^0
	local function closing(cmpfunc) return l.Cmt(l.Cb("t") * l.C((1 - l.P">")^1), cmpfunc) end
	local noslash = l.Cmt(l.Cb("t") / string.lower / {--[[implicit:]] p=1, dt=1, dd=1, li=1, --[[void:]] img=1, input=1, hr=1, br=1, link=1, meta=1}, is)
	local function casecmp(_, _, c1, c2) return c1:lower() == c2:lower() end
	local html_tag = {"<" * l.Cg(tag_name, "t") * noslash * (1 - l.S"><")^0 * (">" - l.B"/"), "</" * closing(casecmp) * ">", {"<\xef\xbf\xbd>", "</\xef\xbf\xbd>"}, "tag name"}
	local function cmp(_, _, c1, c2) return c1 == c2 end
	local xml_tag = {"<" * l.Cg(tag_name, "t") * (1 - l.S"><")^0 * (">" - l.B"/"), "</" * closing(cmp) * ">", {"<\xef\xbf\xbd>", "</\xef\xbf\xbd>"}, "tag name"}
	local function any_pair(set, default) return {l.Cg(l.S(set), "s"), l.Cmt(l.Cb("s") * l.C(1), function(_, _, c1, c2) return builtin_textobjects[c1][2] == c2 end), builtin_textobjects[default]} end
	local any_bracket = any_pair("({[", "(")
	local presets = {
		html = {t = html_tag},
		xml = {t = xml_tag},
		scheme = {b = any_bracket},
		lisp = {b = any_bracket},
		clojure = {b = any_bracket},
		latex = {t = tex_environment},
	}
	for syntax, bindings in pairs(presets) do
		if not M.map[syntax] then
			M.map[syntax] = bindings
		else
			for key, pattern in pairs(bindings) do
				if not M.map[syntax][key] then M.map[syntax][key] = pattern end
			end
		end
	end

end)

M = {
	map = {},
	get_pair = get_delimiters,
	prefix = {outer = "a", inner = "i", opening = "[", closing = "]"},
}

return M