---@class zprite
---@field private _texture love.Image
---@field private _instance_vertex_count integer
---@field private _meshes love.Mesh[]
---@field private _instance_buffers love.Buffer[]
---@field private _instances {integer: {}[]}
---@field private _top_layer integer
---@field private _height_scale number
---@field private _changed boolean
local zprite = {}
zprite.__index = zprite

local ZPRITE_SHADER = love.graphics.newShader([[
#pragma language glsl4

varying vec4 VsColorOut;

#ifdef VERTEX
layout (location = 3) attribute uvec4 InstanceData;

uniform float HeightScale;

vec4 position(mat4 transform_projection, vec4 vertex_position)
{
  uint position_x = bitfieldExtract(InstanceData.x, 0, 16);
  uint position_y = bitfieldExtract(InstanceData.x, 16, 16);
  vec2 position = vec2(position_x, position_y) - vec2(32768, 32768);

  uint uv_x = bitfieldExtract(InstanceData.y, 0, 10);
  uint uv_y = bitfieldExtract(InstanceData.y, 10, 10);
  vec2 uv = vec2(uv_x, uv_y) / vec2(1024, 1024);

  uint scale_x = bitfieldExtract(InstanceData.y, 20, 6);
  uint scale_y = bitfieldExtract(InstanceData.y, 26, 6);
  vec2 scale = vec2(scale_x, scale_y);

  uint r = bitfieldExtract(InstanceData.z, 0, 8);
  uint g = bitfieldExtract(InstanceData.z, 8, 8);
  uint b = bitfieldExtract(InstanceData.z, 16, 8);
  uint a = bitfieldExtract(InstanceData.z, 24, 8);
  vec4 color = vec4(r, g, b, a) / vec4(256, 256, 256, 256);

  uint layer = bitfieldExtract(InstanceData.w, 0, 8);
  uint offset_x = bitfieldExtract(InstanceData.w, 8, 4);
  uint offset_y = bitfieldExtract(InstanceData.w, 12, 4);
  vec2 offset = vec2(offset_x, offset_y) - vec2(8, 8);

  vec2 pixel_offset = offset * vec2(layer * HeightScale, layer * HeightScale) / love_ScreenSize.xy;

  vertex_position.xy *= scale;
  vertex_position.xy += position;
  VaryingTexCoord.xy *= scale;
  VaryingTexCoord.xy += uv;
  VsColorOut = color;
  return (transform_projection * vertex_position) + vec4(pixel_offset, 0, 0);
}
#endif

#ifdef PIXEL
vec4 effect(vec4 color, Image tex, vec2 texture_coords, vec2 screen_coords)
{
  vec4 texture_color = Texel(tex, texture_coords);
  return texture_color * color * VsColorOut;
}
#endif
]])

local function ZPRITE_DEFAULT_COLORMAP(_)
	return 1, 1, 1, 1
end

local function ZPRITE_DEFAULT_OFFSETMAP(_)
	return 0, 2
end

---@param texture love.Image
---@param max_instances integer?
function zprite.new(texture, max_instances)
	if not texture then
		error("must have a texture")
	end
	if max_instances and max_instances < 1 then
		error("max_instances must be at least 1")
	end

	local obj = setmetatable({}, zprite)

	obj._texture = texture
	obj._instance_vertex_count = max_instances or 1024
	obj._meshes = {}
	obj._instances = {}
	obj._instance_buffers = {}
	obj._height_scale = 4
	obj._changed = false

	obj._top_layer = 0

	return obj
end

--- Clear all drawn objects
function zprite:clear()
	self._changed = true
	self._instances = {}
	self._top_layer = 0
end

--- Add an 'entity' to be rendered
---@param x number
---@param y number
---@param quad love.Quad|fun(layer: integer): love.Quad
---@param layer_count integer
---@param color_map (fun(layer: integer): number, number, number, number)?
---@param offset_map (fun(layer: integer): integer, integer)?
function zprite:put(x, y, quad, layer_count, color_map, offset_map)
	if #self._instances + layer_count > self._instance_vertex_count then
		return
	end

	color_map = color_map and color_map or ZPRITE_DEFAULT_COLORMAP
	offset_map = offset_map and offset_map or ZPRITE_DEFAULT_OFFSETMAP

	self._changed = true

	local pos_x = bit.band(x + 32768, 0xFFFF)
	local pos_y = bit.lshift(bit.band(y + 32768, 0xFFFF), 16)
	local pos = bit.bor(pos_x, pos_y)

	for i = 1, layer_count do
		if not self._instances[i] then
			self._instances[i] = {}
			if i > self._top_layer then
				self._top_layer = i
			end
		end
		local r, g, b, a = color_map(i)
		if r > 1 then
			r = 1
		end
		if r < 0 then
			r = 0
		end
		if g > 1 then
			g = 1
		end
		if g < 0 then
			g = 0
		end
		if b > 1 then
			b = 1
		end
		if b < 0 then
			b = 0
		end
		if a > 1 then
			a = 1
		end
		if a < 0 then
			a = 0
		end
		local rb = bit.lshift(bit.band(r * 0xFF, 0xFF), 0)
		local gb = bit.lshift(bit.band(g * 0xFF, 0xFF), 8)
		local bb = bit.lshift(bit.band(b * 0xFF, 0xFF), 16)
		local ab = bit.lshift(bit.band(a * 0xFF, 0xFF), 24)

		local rgba = bit.bor(rb, gb, bb, ab)

		local q = type(quad) == "function" and quad(i) or quad
		local qx, qy, qw, qh = q:getViewport()
		local tw, th = self._texture:getPixelDimensions()

		local uv_x = bit.band(1024 * qx / tw, 0x3FF)
		local uv_y = bit.lshift(bit.band(1024 * qy / th, 0x3FF), 10)
		local uv = bit.bor(uv_x, uv_y)

		local scale_x = bit.lshift(bit.band(qw, 0x3F), 20)
		local scale_y = bit.lshift(bit.band(qh, 0x3F), 26)
		local scale = bit.bor(scale_x, scale_y)

		local uv_scale = bit.bor(uv, scale)

		local offset_x, offset_y = offset_map(i)
		local offset_x_b = bit.lshift(bit.band(offset_x - 8, 0xF), 8)
		local offset_y_b = bit.lshift(bit.band(offset_y - 8, 0xF), 12)

		local layer = bit.band(i, 0xFF)

		local layer_offset = bit.bor(layer, offset_x_b, offset_y_b)

		table.insert(self._instances[i], pos)
		table.insert(self._instances[i], uv_scale)
		table.insert(self._instances[i], rgba)
		table.insert(self._instances[i], layer_offset)
	end
end

--- Draw to the screen / current canvas
---@overload fun(self, x: number, y: number, r: number?, sx: number?, sy: number?, ox: number?, oy: number?, kx: number?, ky: number?)
---@overload fun(self, transform: love.Transform)
function zprite:draw(...)
	zprite._set_shader(self._height_scale)

	self:_upload()

	for i = 1, self._top_layer do
		if self._instances[i] then
			love.graphics.drawInstanced(self._meshes[i], #self._instances[i] / 4, ...)
		end
	end
end

function zprite._set_shader(height_scale)
	love.graphics.setShader(ZPRITE_SHADER)
	ZPRITE_SHADER:send("HeightScale", height_scale)
end

function zprite:_upload()
	if self._changed then
		for i = 1, self._top_layer do
			if self._instances[i] then
				if not self._instance_buffers[i] then
					local tw, th = self._texture:getPixelDimensions()
					local vertices = {
						{ 0, 0, 0, 0 },
						{ 1, 0, 1 / tw, 0 },
						{ 0, 1, 0, 1 / th },
						{ 1, 1, 1 / tw, 1 / th },
					}
					self._meshes[i] = love.graphics.newMesh({
						{ name = "VertexPosition", format = "floatvec2", location = 0 },
						{ name = "VertexTexCoord", format = "floatvec2", location = 1 },
					}, vertices, "strip", "static")
					self._meshes[i]:setTexture(self._texture)
					self._instance_buffers[i] = love.graphics.newBuffer({
						{ name = "InstanceData", format = "uint32vec4", location = 3 },
					}, self._instance_vertex_count, { vertex = true })
					self._meshes[i]:attachAttribute(3, self._instance_buffers[i], "perinstance")
				end
				self._instance_buffers[i]:setArrayData(self._instances[i], 1, 1, #self._instances[i] / 4)
			end
		end
		self._changed = false
	end
end

return zprite