class_name MeshGrid extends Node3D

const _CHUNK_SIZE := 32

const _GRID_ORIGIN := Vector2(0.5, 0.5)

##
##
##
class Chunk:
	var multimesh_instances: Array[MultiMeshInstance] = []

	var meshes: Array[Mesh] = []

	func _init() -> void:
		meshes.resize(_CHUNK_SIZE * _CHUNK_SIZE)

	##
	##
	##
	func invalidate(mesh_grid: MeshGrid, coordinate: Vector2i) -> void:
		# TODO: Once this is all lowered into native code, look for ways to parallelize the loops.
		for multimesh_instance in multimesh_instances:
			RenderingServer.free_rid(multimesh_instance._instance_rid)
			RenderingServer.free_rid(multimesh_instance._multimesh_rid)

		multimesh_instances.clear()
		
		# Normalize mesh instance data for the chunk.
		var transforms_by_mesh := {}
		var grid_size := mesh_grid.size

		for i in meshes.size():
			var mesh := meshes[i]

			if mesh != null:
				if not(mesh in transforms_by_mesh):
					transforms_by_mesh[mesh] = []

				transforms_by_mesh[mesh].append(Transform3D(Basis.IDENTITY, Vector3(
					(float(coordinate.x * _CHUNK_SIZE) + (i % _CHUNK_SIZE)) -
						(float(grid_size.x) * _GRID_ORIGIN.x), 0.0,
							(float(coordinate.y * _CHUNK_SIZE) + (i / _CHUNK_SIZE)) - 
								(float(grid_size.y) * _GRID_ORIGIN.y))))

		# Bake into multimesh instances for the chunk.
		var scenario_rid := mesh_grid.get_world_3d().scenario
		var global_transform := mesh_grid.global_transform

		for chunk_mesh in transforms_by_mesh:
			var multimesh_instance := MultiMeshInstance.new(
				scenario_rid, chunk_mesh.get_rid(), transforms_by_mesh[chunk_mesh])

			multimesh_instance.set_offset(global_transform)
			multimesh_instances.append(multimesh_instance)

	##
	##
	##
	func set_mesh(coordinate: Vector2i, mesh: Mesh) -> void:
		meshes[(_CHUNK_SIZE * coordinate.y) + coordinate.x] = mesh

##
##
##
class MultiMeshInstance:
	var _instance_rid := RID()

	var _multimesh_rid := RID()

	func _init(scenario_rid: RID, mesh_rid: RID, transforms: Array) -> void:
		_multimesh_rid = RenderingServer.multimesh_create()
		_instance_rid = RenderingServer.instance_create2(_multimesh_rid, scenario_rid)

		RenderingServer.multimesh_set_mesh(_multimesh_rid, mesh_rid)

		var transform_count := transforms.size()

		RenderingServer.multimesh_allocate_data(_multimesh_rid,
			transform_count, RenderingServer.MULTIMESH_TRANSFORM_3D, true)

		for i in transform_count:
			RenderingServer.multimesh_instance_set_transform(_multimesh_rid, i, transforms[i])

	##
	##
	##
	func set_offset(offset: Transform3D) -> void:
		RenderingServer.instance_set_transform(_instance_rid, offset)

var _chunks: Array[Chunk] = []

var _chunks_size := Vector2i.ZERO

##
##
##
@export
var size: Vector2i:
	set(value):
		var chunk_size_factor := float(_CHUNK_SIZE)

		_chunks.resize(int(ceilf(value.x / chunk_size_factor) * ceilf(value.y / chunk_size_factor)))

		for i in _chunks.size():
			_chunks[i] = Chunk.new()

		_chunks_size = Vector2i((value / float(_CHUNK_SIZE)).ceil())
		size = value

func _get_chunk(chunk_coordinate: Vector2i) -> Chunk:
	return _chunks[(_chunks_size.x * chunk_coordinate.y) + chunk_coordinate.x]

func _notification(what: int) -> void:
	match what:
		NOTIFICATION_TRANSFORM_CHANGED:
			for chunk in _chunks:
				for multimesh_instance in chunk.multimesh_instances:
					multimesh_instance.set_offset_transform(global_transform)

##
##
##
func clear_mesh(mesh: Mesh) -> void:
	for y in size.y:
		for x in size.x:
			var coordinate := Vector2i(x, y)

			_get_chunk(coordinate / _CHUNK_SIZE).set_mesh(coordinate % _CHUNK_SIZE, mesh)

	for y in _chunks_size.y:
		for x in _chunks_size.x:
			var chunk_coordinate := Vector2i(x, y)

			_get_chunk(chunk_coordinate).invalidate(self, chunk_coordinate)

##
##
##
func fill_mesh(area: Rect2i, mesh: Mesh) -> void:
	assert(Rect2i(Vector2i.ZERO, size).encloses(area), "area must be within grid")

	var filled_chunks := BitMap.new()

	filled_chunks.resize(_chunks_size)

	for y in range(area.position.y, area.end.y):
		for x in range(area.position.x, area.end.x):
			var coordinate := Vector2i(x, y)
			var chunk_coordinate := coordinate / _CHUNK_SIZE

			_get_chunk(chunk_coordinate).set_mesh(coordinate % _CHUNK_SIZE, mesh)
			filled_chunks.set_bitv(chunk_coordinate, true)

	for y in _chunks_size.y:
		for x in _chunks_size.x:
			if filled_chunks.get_bit(x, y):
				var chunk_coordinate := Vector2i(x, y)

				_get_chunk(chunk_coordinate).invalidate(self, chunk_coordinate)

##
##
##
func plot_mesh(coordinate: Vector2i, mesh: Mesh) -> void:
	assert(Rect2i(Vector2i.ZERO, size).has_point(coordinate), "coordinate must be within grid")

	var chunk_coordinate := coordinate / _CHUNK_SIZE
	var chunk := _get_chunk(chunk_coordinate)

	chunk.set_mesh(coordinate % _CHUNK_SIZE, mesh)
	chunk.invalidate(self, chunk_coordinate)

##
## Returns [code]world_position[/code] converted into a coordinate aligned with the [MeshGrid].
##
## Note that [code]world_position[/code] values not within the [MeshGrid] will produce grid
## coordinates outside of the [MeshGrid] bounds as well.
##
func world_to_grid(world_position: Vector2) -> Vector2i:
	return Vector2i((world_position + (Vector2(size) * _GRID_ORIGIN)).floor())