local unicode = require("luacheck.unicode")
local utils = require("luacheck.utils")
local decoder = {}
local sbyte = string.byte
local sfind = string.find
local sgsub = string.gsub
local ssub = string.sub
-- `LatinChars` and `UnicodeChars` objects represent source strings
-- and provide Unicode-aware access to them with a common interface.
-- Source bytes should not be accessed directly.
-- Provided methods are:
-- `Chars:get_codepoint(index)`: returns codepoint at given index as integer or nil if index is out of range.
-- `Chars:get_substring(from, to)`: returns substring of original bytes corresponding to characters from `from` to `to`.
-- `Chars:get_printable_substring(from. to)`: like get_substring but escapes not printable characters.
-- `Chars:get_length()`: returns total number of characters.
-- `Chars:find(pattern, from)`: `string.find` but `from` is in characters. Return values are still in bytes.
-- `LatinChars` is an optimized special case for latin1 strings.
local LatinChars = utils.class()

function LatinChars:__init(bytes)
	self._bytes = bytes
end

function LatinChars:get_codepoint(index)
	return sbyte(self._bytes, index)
end

function LatinChars:get_substring(from, to)
	return ssub(self._bytes, from, to)
end

local function hexadecimal_escaper(byte)
	return ("\\x%02X"):format(sbyte(byte))
end

function LatinChars:get_printable_substring(from, to)
	return (sgsub(ssub(self._bytes, from, to), "[^\32-\126]", hexadecimal_escaper))
end

function LatinChars:get_length()
	return #self._bytes
end

function LatinChars:find(pattern, from)
	return sfind(self._bytes, pattern, from)
end

-- Decodes `bytes` as UTF8. Returns arrays of codepoints as integers and their byte offsets.
-- Byte offsets have one extra item pointing to one byte past the end of `bytes`.
-- On decoding error returns nothing.
local function get_codepoints_and_byte_offsets(bytes)
	local codepoints = {}
	local byte_offsets = {}
	local byte_index = 1
	local codepoint_index = 1

	while true do
		byte_offsets[codepoint_index] = byte_index
		-- Attempt to decode the next codepoint from UTF8.
		local codepoint = sbyte(bytes, byte_index)

		if not codepoint then return codepoints, byte_offsets end

		byte_index = byte_index + 1

		if codepoint >= 0x80 then
			-- Not ASCII.
			if codepoint < 0xC0 then return end

			local cont = (sbyte(bytes, byte_index) or 0) - 0x80

			if cont < 0 or cont >= 0x40 then return end

			byte_index = byte_index + 1

			if codepoint < 0xE0 then
				-- Two bytes.
				codepoint = cont + (codepoint - 0xC0) * 0x40
			elseif codepoint < 0xF0 then
				-- Three bytes.
				codepoint = cont + (codepoint - 0xE0) * 0x40
				cont = (sbyte(bytes, byte_index) or 0) - 0x80

				if cont < 0 or cont >= 0x40 then return end

				byte_index = byte_index + 1
				codepoint = cont + codepoint * 0x40
			elseif codepoint < 0xF8 then
				-- Four bytes.
				codepoint = cont + (codepoint - 0xF0) * 0x40
				cont = (sbyte(bytes, byte_index) or 0) - 0x80

				if cont < 0 or cont >= 0x40 then return end

				byte_index = byte_index + 1
				codepoint = cont + codepoint * 0x40
				cont = (sbyte(bytes, byte_index) or 0) - 0x80

				if cont < 0 or cont >= 0x40 then return end

				byte_index = byte_index + 1
				codepoint = cont + codepoint * 0x40

				if codepoint > 0x10FFFF then return end
			else
				return
			end
		end

		codepoints[codepoint_index] = codepoint
		codepoint_index = codepoint_index + 1
	end
end

-- `UnicodeChars` is the general case for non-latin1 strings.
-- Assumes UTF8, on decoding error falls back to latin1.
local UnicodeChars = utils.class()

function UnicodeChars:__init(bytes, codepoints, byte_offsets)
	self._bytes = bytes
	self._codepoints = codepoints
	self._byte_offsets = byte_offsets
end

function UnicodeChars:get_codepoint(index)
	return self._codepoints[index]
end

function UnicodeChars:get_substring(from, to)
	local byte_offsets = self._byte_offsets
	return ssub(self._bytes, byte_offsets[from], byte_offsets[to + 1] - 1)
end

function UnicodeChars:get_printable_substring(from, to)
	-- This is only called on syntax error, it's okay to be slow.
	local parts = {}

	for index = from, to do
		local codepoint = self._codepoints[index]

		if unicode.is_printable(codepoint) then
			table.insert(parts, self:get_substring(index, index))
		else
			table.insert(parts, (codepoint > 255 and "\\u{%X}" or "\\x%02X"):format(codepoint))
		end
	end

	return table.concat(parts)
end

function UnicodeChars:get_length()
	return #self._codepoints
end

function UnicodeChars:find(pattern, from)
	return sfind(self._bytes, pattern, self._byte_offsets[from])
end

function decoder.decode(bytes)
	-- Only use UnicodeChars if necessary. LatinChars isn't much faster but noticeably more memory efficient.
	if sfind(bytes, "[\128-\255]") then
		local codepoints, byte_offsets = get_codepoints_and_byte_offsets(bytes)

		if codepoints then return UnicodeChars(bytes, codepoints, byte_offsets) end
	end

	return LatinChars(bytes)
end

return decoder