local zprite = require("zprite.zprite")

---@class zchunk
---@field private _chunks {integer: {z: zprite, x: integer, y: integer}}
---@field private _chunk_size integer
---@field private _chunk_height integer
---@field private _chunk_wrap integer
---@field private _texture love.Image
---@field private _height_scale number
local zchunk = {}
zchunk.__index = zchunk

---@param texture love.Image
---@param chunk_size integer?
---@param chunk_height integer?
---@param chunk_wrap integer?
function zchunk.new(texture, chunk_size, chunk_height, chunk_wrap)
	chunk_size = chunk_size or 16
	chunk_height = chunk_height or 32
	if chunk_size < 1 then
		error("chunk size must be at least 1")
	end
	if chunk_height < 1 then
		error("chunk height must be at least 1")
	end
	if chunk_wrap and chunk_wrap < 1 then
		error("chunk wrap must be at least 1")
	end
	if not texture then
		error("must have a texture")
	end

	local obj = setmetatable({}, zchunk)

	obj._chunks = {}
	obj._chunk_size = chunk_size
	obj._chunk_height = chunk_height
	obj._chunk_wrap = chunk_wrap

	obj._texture = texture

	obj._height_scale = 4

	return obj
end

---@param x integer
---@param y integer
---@return integer
function zchunk:_chunk_index(x, y)
	local cx = math.floor(x / self._chunk_size)
	local cy = math.floor(y / self._chunk_size)
	if self._chunk_wrap then
		cx = cx % self._chunk_wrap
		cy = cy % self._chunk_wrap
	end

	return cx + cy * 1e7
end

---@param x integer?
---@param y integer?
function zchunk:remove(x, y)
	if not x then
		self._chunks = {}
	else
		local index = self:_chunk_index(x, y)
		self._chunks[index] = nil
	end
end

---@param x integer?
---@param y integer?
function zchunk:clear(x, y)
	if not x then
		for _, chunk in pairs(self._chunks) do
			chunk.z:clear()
		end
	else
		local index = self:_chunk_index(x, y)
		if not self._chunks[index] then
			return
		end
		self._chunks[index].z:clear()
	end
end

--- Add an 'entity' to be rendered
---@param x number
---@param y number
---@param quad love.Quad
---@param layer_count integer
---@param color_map fun(layer: integer): number, number, number
---@param base_layer integer?
function zchunk:put(x, y, quad, layer_count, color_map, base_layer)
	local index = self:_chunk_index(x, y)
	if not self._chunks[index] then
		self._chunks[index] = {
			z = zprite.new(self._texture, self._chunk_size * self._chunk_size * self._chunk_height),
		}
		self._chunks[index].z._height_scale = self._height_scale
	end
	self._chunks[index].z:put(x, y, quad, layer_count, color_map, base_layer)
	self._chunks[index].x = math.floor(x / self._chunk_size)
	self._chunks[index].y = math.floor(y / self._chunk_size)
end

local transform = love.math.newTransform()
local function sort_fn(a, b)
	local _, ay_rotated = transform:transformPoint(a.x, a.y)
	local _, by_rotated = transform:transformPoint(b.x, b.y)

	return ay_rotated < by_rotated
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 zchunk:draw(x, y, r, ...)
	local chunks = {}
	for _, chunk in pairs(self._chunks) do
		table.insert(chunks, chunk)
	end

	transform:reset()
	if type(x) ~= "number" then
		transform:apply(x)
	else
		transform:rotate(r and r or 0)
	end

	table.sort(chunks, sort_fn)

	for _, chunk in ipairs(chunks) do
		chunk.z:draw(x, y, r, ...)
	end
end

return zchunk