-- lx-bundle:true

--- Small raw-DEFLATE inflater for ZIP members.
-- Supports stored, fixed Huffman, and dynamic Huffman blocks.
-- @module core.lib.inflate

local inflate = {}

local function byte_at(data, index)
	return data:byte(index) or 0
end

local function reverse_bits(value, length)
	local out = 0

	for _ = 1, length do
		out = (out << 1) | (value & 1)
		value = value >> 1
	end

	return out
end

local BitReader = {}
BitReader.__index = BitReader

function BitReader:new(data)
	return setmetatable({
		data = data,
		pos = 1,
		bitbuf = 0,
		bitcnt = 0,
	}, self)
end

function BitReader:read_bits(count)
	while self.bitcnt < count do
		if self.pos > #self.data then
			return nil, "unexpected end of deflate stream"
		end

		self.bitbuf = self.bitbuf | (byte_at(self.data, self.pos) << self.bitcnt)
		self.bitcnt = self.bitcnt + 8
		self.pos = self.pos + 1
	end

	local mask = (1 << count) - 1
	local value = self.bitbuf & mask
	self.bitbuf = self.bitbuf >> count
	self.bitcnt = self.bitcnt - count

	return value
end

function BitReader:align_to_byte()
	self.bitbuf = 0
	self.bitcnt = 0
end

function BitReader:read_stored_bytes(length)
	if self.pos + length - 1 > #self.data then
		return nil, "stored block exceeds deflate stream"
	end

	local text = self.data:sub(self.pos, self.pos + length - 1)
	self.pos = self.pos + length
	return text
end

local function build_huffman(lengths)
	local max_bits = 0
	local bl_count = {}

	for symbol = 0, #lengths do
		local length = lengths[symbol] or 0

		if length > 0 then
			bl_count[length] = (bl_count[length] or 0) + 1
			if length > max_bits then
				max_bits = length
			end
		end
	end

	local code = 0
	local next_code = {}

	for bits = 1, max_bits do
		code = (code + (bl_count[bits - 1] or 0)) << 1
		next_code[bits] = code
	end

	local lookup = {}

	for symbol = 0, #lengths do
		local length = lengths[symbol] or 0

		if length > 0 then
			local symbol_code = next_code[length]
			next_code[length] = symbol_code + 1
			lookup[length .. ":" .. reverse_bits(symbol_code, length)] = symbol
		end
	end

	return {
		lookup = lookup,
		max_bits = max_bits,
	}
end

local function decode_symbol(reader, tree)
	local code = 0

	for length = 1, tree.max_bits do
		local bit, err = reader:read_bits(1)
		if not bit then
			return nil, err
		end

		code = code | (bit << (length - 1))

		local symbol = tree.lookup[length .. ":" .. code]
		if symbol ~= nil then
			return symbol
		end
	end

	return nil, "invalid deflate huffman code"
end

local fixed_literal_tree
local fixed_distance_tree

local function fixed_trees()
	if fixed_literal_tree and fixed_distance_tree then
		return fixed_literal_tree, fixed_distance_tree
	end

	local literal_lengths = {}
	for symbol = 0, 143 do literal_lengths[symbol] = 8 end
	for symbol = 144, 255 do literal_lengths[symbol] = 9 end
	for symbol = 256, 279 do literal_lengths[symbol] = 7 end
	for symbol = 280, 287 do literal_lengths[symbol] = 8 end

	local distance_lengths = {}
	for symbol = 0, 31 do distance_lengths[symbol] = 5 end

	fixed_literal_tree = build_huffman(literal_lengths)
	fixed_distance_tree = build_huffman(distance_lengths)

	return fixed_literal_tree, fixed_distance_tree
end

local LENGTH_BASE = {
	[257] = 3, [258] = 4, [259] = 5, [260] = 6, [261] = 7, [262] = 8, [263] = 9, [264] = 10,
	[265] = 11, [266] = 13, [267] = 15, [268] = 17, [269] = 19, [270] = 23, [271] = 27, [272] = 31,
	[273] = 35, [274] = 43, [275] = 51, [276] = 59, [277] = 67, [278] = 83, [279] = 99, [280] = 115,
	[281] = 131, [282] = 163, [283] = 195, [284] = 227, [285] = 258,
}

local LENGTH_EXTRA = {
	[257] = 0, [258] = 0, [259] = 0, [260] = 0, [261] = 0, [262] = 0, [263] = 0, [264] = 0,
	[265] = 1, [266] = 1, [267] = 1, [268] = 1, [269] = 2, [270] = 2, [271] = 2, [272] = 2,
	[273] = 3, [274] = 3, [275] = 3, [276] = 3, [277] = 4, [278] = 4, [279] = 4, [280] = 4,
	[281] = 5, [282] = 5, [283] = 5, [284] = 5, [285] = 0,
}

local DIST_BASE = {
	[0] = 1, [1] = 2, [2] = 3, [3] = 4, [4] = 5, [5] = 7, [6] = 9, [7] = 13,
	[8] = 17, [9] = 25, [10] = 33, [11] = 49, [12] = 65, [13] = 97, [14] = 129, [15] = 193,
	[16] = 257, [17] = 385, [18] = 513, [19] = 769, [20] = 1025, [21] = 1537, [22] = 2049, [23] = 3073,
	[24] = 4097, [25] = 6145, [26] = 8193, [27] = 12289, [28] = 16385, [29] = 24577,
}

local DIST_EXTRA = {
	[0] = 0, [1] = 0, [2] = 0, [3] = 0, [4] = 1, [5] = 1, [6] = 2, [7] = 2,
	[8] = 3, [9] = 3, [10] = 4, [11] = 4, [12] = 5, [13] = 5, [14] = 6, [15] = 6,
	[16] = 7, [17] = 7, [18] = 8, [19] = 8, [20] = 9, [21] = 9, [22] = 10, [23] = 10,
	[24] = 11, [25] = 11, [26] = 12, [27] = 12, [28] = 13, [29] = 13,
}

local CODE_LENGTH_ORDER = {
	16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15,
}

local function dynamic_trees(reader)
	local hlit, err = reader:read_bits(5)
	if not hlit then return nil, nil, err end
	local hdist
	hdist, err = reader:read_bits(5)
	if not hdist then return nil, nil, err end
	local hclen
	hclen, err = reader:read_bits(4)
	if not hclen then return nil, nil, err end

	hlit = hlit + 257
	hdist = hdist + 1
	hclen = hclen + 4

	local code_lengths = {}
	for index = 1, hclen do
		local length
		length, err = reader:read_bits(3)
		if not length then return nil, nil, err end
		code_lengths[CODE_LENGTH_ORDER[index]] = length
	end

	for _, symbol in ipairs(CODE_LENGTH_ORDER) do
		code_lengths[symbol] = code_lengths[symbol] or 0
	end

	local code_tree = build_huffman(code_lengths)
	local lengths = {}
	local total = hlit + hdist
	local index = 0

	while index < total do
		local symbol
		symbol, err = decode_symbol(reader, code_tree)
		if not symbol then return nil, nil, err end

		if symbol <= 15 then
			lengths[index] = symbol
			index = index + 1
		elseif symbol == 16 then
			if index == 0 then
				return nil, nil, "deflate repeat length has no previous value"
			end
			local extra
			extra, err = reader:read_bits(2)
			if not extra then return nil, nil, err end
			local repeat_count = 3 + extra
			local previous = lengths[index - 1] or 0
			for _ = 1, repeat_count do
				lengths[index] = previous
				index = index + 1
			end
		elseif symbol == 17 then
			local extra
			extra, err = reader:read_bits(3)
			if not extra then return nil, nil, err end
			for _ = 1, 3 + extra do
				lengths[index] = 0
				index = index + 1
			end
		elseif symbol == 18 then
			local extra
			extra, err = reader:read_bits(7)
			if not extra then return nil, nil, err end
			for _ = 1, 11 + extra do
				lengths[index] = 0
				index = index + 1
			end
		else
			return nil, nil, "invalid code-length huffman symbol"
		end

		if index > total then
			return nil, nil, "deflate code lengths exceed header count"
		end
	end

	local literal_lengths = {}
	for symbol = 0, hlit - 1 do
		literal_lengths[symbol] = lengths[symbol] or 0
	end

	local distance_lengths = {}
	for symbol = 0, hdist - 1 do
		distance_lengths[symbol] = lengths[hlit + symbol] or 0
	end

	if not literal_lengths[256] or literal_lengths[256] == 0 then
		return nil, nil, "deflate literal tree lacks end-of-block symbol"
	end

	return build_huffman(literal_lengths), build_huffman(distance_lengths)
end

local function append_string_as_bytes(output, text)
	for index = 1, #text do
		output[#output + 1] = text:byte(index)
	end
end

local function bytes_to_string(bytes)
	local chunks = {}
	local chunk = {}

	for index = 1, #bytes do
		chunk[#chunk + 1] = string.char(bytes[index])
		if #chunk >= 8192 then
			chunks[#chunks + 1] = table.concat(chunk)
			chunk = {}
		end
	end

	if #chunk > 0 then
		chunks[#chunks + 1] = table.concat(chunk)
	end

	return table.concat(chunks)
end

local function inflate_huffman_block(reader, output, literal_tree, distance_tree)
	while true do
		local symbol, err = decode_symbol(reader, literal_tree)
		if not symbol then
			return nil, err
		end

		if symbol < 256 then
			output[#output + 1] = symbol
		elseif symbol == 256 then
			return true
		elseif symbol <= 285 then
			local extra_bits = LENGTH_EXTRA[symbol]
			local extra_value = 0

			if extra_bits and extra_bits > 0 then
				extra_value, err = reader:read_bits(extra_bits)
				if not extra_value then return nil, err end
			elseif extra_bits == nil then
				return nil, "invalid deflate length symbol"
			end

			local length = LENGTH_BASE[symbol] + extra_value
			local distance_symbol
			distance_symbol, err = decode_symbol(reader, distance_tree)
			if not distance_symbol then return nil, err end

			local distance_extra_bits = DIST_EXTRA[distance_symbol]
			if distance_extra_bits == nil then
				return nil, "invalid deflate distance symbol"
			end

			local distance_extra_value = 0
			if distance_extra_bits > 0 then
				distance_extra_value, err = reader:read_bits(distance_extra_bits)
				if not distance_extra_value then return nil, err end
			end

			local distance = DIST_BASE[distance_symbol] + distance_extra_value
			if distance < 1 or distance > #output then
				return nil, "invalid deflate back-reference distance"
			end

			for _ = 1, length do
				output[#output + 1] = output[#output - distance + 1]
			end
		else
			return nil, "invalid deflate literal symbol"
		end
	end
end

function inflate.inflate_raw(data)
	if type(data) ~= "string" then
		return nil, "deflate input must be a string"
	end

	local reader = BitReader:new(data)
	local output = {}
	local final = false

	while not final do
		local bit, err = reader:read_bits(1)
		if not bit then
			return nil, err
		end
		final = bit == 1

		local block_type
		block_type, err = reader:read_bits(2)
		if not block_type then
			return nil, err
		end

		if block_type == 0 then
			reader:align_to_byte()
			local len1 = byte_at(data, reader.pos)
			local len2 = byte_at(data, reader.pos + 1)
			local nlen1 = byte_at(data, reader.pos + 2)
			local nlen2 = byte_at(data, reader.pos + 3)
			reader.pos = reader.pos + 4

			local length = len1 | (len2 << 8)
			local nlength = nlen1 | (nlen2 << 8)
			if ((length ~ nlength) & 0xffff) ~= 0xffff then
				return nil, "invalid deflate stored-block length"
			end

			local stored
			stored, err = reader:read_stored_bytes(length)
			if not stored then
				return nil, err
			end
			append_string_as_bytes(output, stored)
		elseif block_type == 1 then
			local literal_tree, distance_tree = fixed_trees()
			local ok
			ok, err = inflate_huffman_block(reader, output, literal_tree, distance_tree)
			if not ok then
				return nil, err
			end
		elseif block_type == 2 then
			local literal_tree, distance_tree
			literal_tree, distance_tree, err = dynamic_trees(reader)
			if not literal_tree then
				return nil, err
			end
			local ok
			ok, err = inflate_huffman_block(reader, output, literal_tree, distance_tree)
			if not ok then
				return nil, err
			end
		else
			return nil, "unsupported reserved deflate block type"
		end
	end

	return bytes_to_string(output)
end

return inflate
