-- 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.

--- Windows ICO and CUR metadata and frame loader
-- @module plugins.icopreview.ico

local ico = {}

local PNG_SIGNATURE = "\137PNG\r\n\26\n"
local MAX_ENTRY_COUNT = 256
local MAX_PAYLOAD_WORK = 64 * 1024 * 1024
local MAX_IMAGE_DIMENSION = 4096
local MAX_IMAGE_PIXELS = 16 * 1024 * 1024

--- Validate dimensions before they reach preview layout or native decoding
-- @param width number decoded width
-- @param height number decoded height
-- @param format string image format name
-- @return boolean valid
-- @return string|nil error_message
local function validate_dimensions(width, height, format)
	if width < 1 or height < 1 then
		return false, format .. " dimensions must be positive"
	end
	if width > MAX_IMAGE_DIMENSION or height > MAX_IMAGE_DIMENSION
		or width > MAX_IMAGE_PIXELS // height
	then
		return false, string.format(
			"%s dimensions exceed the preview limit: %dx%d",
			format,
			width,
			height
		)
	end

	return true
end

--- Reject a byte range that extends beyond image-set data
-- @param data string complete ICO or CUR data
-- @param offset number zero-based byte offset
-- @param size number byte count
-- @param description string range name for errors
local function check_range(data, offset, size, description)
	if offset < 0 or size < 0 or offset + size > #data then
		error(string.format(
			"%s is outside the file: offset=%d size=%d file_size=%d",
			description,
			offset,
			size,
			#data
		), 3)
	end
end

--- Read one little-endian unsigned 16-bit value
-- @param data string source bytes
-- @param offset number zero-based byte offset
-- @return number value
local function u16le(data, offset)
	check_range(data, offset, 2, "uint16")

	return string.unpack("<I2", data, offset + 1)
end

--- Read one little-endian unsigned 32-bit value
-- @param data string source bytes
-- @param offset number zero-based byte offset
-- @return number value
local function u32le(data, offset)
	check_range(data, offset, 4, "uint32")

	return string.unpack("<I4", data, offset + 1)
end

--- Read one little-endian signed 32-bit value
-- @param data string source bytes
-- @param offset number zero-based byte offset
-- @return number value
local function i32le(data, offset)
	check_range(data, offset, 4, "int32")

	return string.unpack("<i4", data, offset + 1)
end

--- Read one big-endian unsigned 32-bit value
-- @param data string source bytes
-- @param offset number zero-based byte offset
-- @return number value
local function u32be(data, offset)
	check_range(data, offset, 4, "big-endian uint32")

	return string.unpack(">I4", data, offset + 1)
end

--- Read one byte
-- @param data string source bytes
-- @param offset number zero-based byte offset
-- @return number value
local function byte(data, offset)
	check_range(data, offset, 1, "byte")

	return data:byte(offset + 1)
end

--- Read dimensions and color information from a PNG frame
-- @param payload string complete PNG payload
-- @return table|nil metadata
-- @return string|nil error_message
local function parse_png_metadata(payload)
	if #payload < 33 or payload:sub(1, 8) ~= PNG_SIGNATURE then
		return nil, "invalid or truncated PNG payload"
	end
	if u32be(payload, 8) ~= 13 or payload:sub(13, 16) ~= "IHDR" then
		return nil, "PNG does not begin with a normal IHDR chunk"
	end

	local width = u32be(payload, 16)
	local height = u32be(payload, 20)
	local valid, dimension_error = validate_dimensions(width, height, "PNG")
	if not valid then
		return nil, dimension_error
	end

	return {
		format = "png",
		width = width,
		height = height,
		bit_depth = byte(payload, 24),
		color_type = byte(payload, 25),
	}
end

--- Read dimensions and color information from a DIB frame
-- @param payload string complete DIB payload
-- @return table|nil metadata
-- @return string|nil error_message
local function parse_dib_metadata(payload)
	if #payload < 4 then
		return nil, "truncated DIB payload"
	end

	local header_size = u32le(payload, 0)

	if header_size == 12 then
		if #payload < 12 then
			return nil, "truncated BITMAPCOREHEADER"
		end

		local width = u16le(payload, 4)
		local stored_height = u16le(payload, 6)
		local height = stored_height // 2
		local bit_count = u16le(payload, 10)
		local valid, dimension_error = validate_dimensions(width, height, "DIB")
		if not valid then
			return nil, dimension_error
		end
		if stored_height % 2 ~= 0 or bit_count < 1 or bit_count > 32 then
			return nil, "invalid DIB height or bit depth"
		end

		local color_stride = ((width * bit_count + 31) // 32) * 4
		local mask_stride = ((width + 31) // 32) * 4
		local required_size = header_size + (color_stride + mask_stride) * height
		if required_size > #payload then
			return nil, string.format(
				"truncated DIB pixels: need at least %d bytes, have %d",
				required_size,
				#payload
			)
		end

		return {
			format = "dib",
			header_size = header_size,
			width = width,
			height = height,
			bit_count = bit_count,
		}
	end
	if header_size < 40 then
		return nil, string.format("unsupported or invalid DIB header size: %d", header_size)
	end
	if #payload < header_size then
		return nil, "truncated Windows bitmap header"
	end

	local width = math.abs(i32le(payload, 4))
	local stored_height = math.abs(i32le(payload, 8))
	local height = stored_height // 2
	local bit_count = u16le(payload, 14)
	local compression = u32le(payload, 16)
	local valid, dimension_error = validate_dimensions(width, height, "DIB")
	if not valid then
		return nil, dimension_error
	end
	if stored_height % 2 ~= 0 or bit_count < 1 or bit_count > 32 then
		return nil, "invalid DIB height or bit depth"
	end

	if compression == 0 then
		local color_stride = ((width * bit_count + 31) // 32) * 4
		local mask_stride = ((width + 31) // 32) * 4
		local required_size = header_size + (color_stride + mask_stride) * height
		if required_size > #payload then
			return nil, string.format(
				"truncated DIB pixels: need at least %d bytes, have %d",
				required_size,
				#payload
			)
		end
	end

	return {
		format = "dib",
		header_size = header_size,
		width = width,
		height = height,
		bit_count = bit_count,
		compression = compression,
	}
end

--- Return effective total bits per pixel for one parsed frame
-- @param image table parsed payload metadata
-- @param fallback number directory bit depth
-- @return number bits_per_pixel
local function get_bits_per_pixel(image, fallback)
	if image.bit_count then
		return image.bit_count
	end

	local channel_counts = { [0] = 1, [2] = 3, [3] = 1, [4] = 2, [6] = 4 }
	local channel_count = channel_counts[image.color_type]

	return channel_count and image.bit_depth * channel_count or fallback
end

--- Parse one complete ICO or CUR file
-- @param data string complete ICO or CUR data
-- @return table image_set
function ico.parse(data)
	assert(type(data) == "string", "ICO or CUR data must be a string")
	check_range(data, 0, 6, "ICO or CUR header")

	local reserved = u16le(data, 0)
	local resource_type = u16le(data, 2)
	local count = u16le(data, 4)

	if reserved ~= 0 then
		error("invalid ICO or CUR reserved field: " .. reserved)
	end
	if resource_type ~= 1 and resource_type ~= 2 then
		error(string.format("unsupported resource type: %d", resource_type))
	end

	local type_name = resource_type == 1 and "icon" or "cursor"
	local type_label = resource_type == 1 and "ICO" or "CUR"

	if count == 0 then
		error(type_label .. " contains no images")
	end
	if count > MAX_ENTRY_COUNT then
		error(string.format("%s contains too many images: %d (maximum %d)", type_label, count, MAX_ENTRY_COUNT))
	end

	check_range(data, 6, count * 16, type_label .. " directory table")

	local payload_work = 0

	-- Count overlapping ranges each time because every directory entry is decoded separately
	for index = 1, count do
		local entry_offset = 6 + (index - 1) * 16
		local payload_size = u32le(data, entry_offset + 8)
		local payload_offset = u32le(data, entry_offset + 12)

		check_range(data, payload_offset, payload_size, string.format("image payload %d", index - 1))

		if payload_size > MAX_PAYLOAD_WORK - payload_work then
			error(string.format("%s image payloads exceed the %d MiB work limit", type_label, MAX_PAYLOAD_WORK // 1024 // 1024))
		end

		payload_work = payload_work + payload_size
	end

	local image_set = {
		type = type_name,
		resource_type = resource_type,
		count = count,
		file_size = #data,
		entries = {},
	}

	for index = 1, count do
		local entry_offset = 6 + (index - 1) * 16
		local width_byte = byte(data, entry_offset)
		local height_byte = byte(data, entry_offset + 1)
		local payload_size = u32le(data, entry_offset + 8)
		local payload_offset = u32le(data, entry_offset + 12)

		check_range(data, payload_offset, payload_size, string.format("image payload %d", index - 1))

		local payload = data:sub(payload_offset + 1, payload_offset + payload_size)
		local image, metadata_error

		if payload:sub(1, 8) == PNG_SIGNATURE then
			image, metadata_error = parse_png_metadata(payload)
		else
			image, metadata_error = parse_dib_metadata(payload)
		end

		local entry = {
			index = index - 1,
			width_byte = width_byte,
			height_byte = height_byte,
			width = width_byte == 0 and 256 or width_byte,
			height = height_byte == 0 and 256 or height_byte,
			color_count = byte(data, entry_offset + 2),
			reserved = byte(data, entry_offset + 3),
			payload = payload,
			image = image,
			metadata_error = metadata_error,
		}

		if resource_type == 1 then
			entry.planes = u16le(data, entry_offset + 4)
			entry.bit_count = u16le(data, entry_offset + 6)
		else
			entry.hotspot_x = u16le(data, entry_offset + 4)
			entry.hotspot_y = u16le(data, entry_offset + 6)
		end

		if image then
			entry.width = image.width
			entry.height = image.height
			entry.format = image.format:upper()
			entry.bits_per_pixel = get_bits_per_pixel(image, entry.bit_count)

			if resource_type == 2
				and (entry.hotspot_x >= entry.width or entry.hotspot_y >= entry.height)
			then
				entry.hotspot_warning = "hotspot lies outside image"
			end
		end

		image_set.entries[index] = entry
	end

	return image_set
end

--- Read and parse one ICO or CUR file
-- @param path string ICO or CUR path
-- @return table|nil image_set
-- @return string|nil error_message
function ico.read_file(path)
	local handle, open_error = io.open(path, "rb")

	if not handle then
		return nil, open_error
	end

	local data = handle:read("*a")
	handle:close()

	local loaded, image_set = pcall(ico.parse, data)

	return loaded and image_set or nil, loaded and nil or tostring(image_set)
end

--- Decode one ICO or CUR entry through existing image loader
-- CUR hotspot words cannot be passed as ICO color fields, so cursor frames use decoded values
-- @param entry table parsed ICO or CUR entry
-- @return userdata|nil image
-- @return string|nil error_message
function ico.load_entry_image(entry)
	local planes = entry.planes or 1
	local bit_count = entry.bit_count or entry.bits_per_pixel or 32
	local single_image = string.pack(
		"<I2I2I2BBBBI2I2I4I4",
		0, 1, 1,
		entry.width_byte, entry.height_byte, entry.color_count, entry.reserved,
		planes, bit_count, #entry.payload, 22
	) .. entry.payload
	local loaded, image, load_error = pcall(canvas.load_image_data, single_image)

	if not loaded then
		return nil, tostring(image)
	end

	return image, load_error
end

return ico
