class_name MeshGrid extends Node3D

const _CHUNK_SIZE := 32

const _GRID_ORIGIN := Vector2(0.5, 0.5)

##
## Baked block of meshes making up a segment of the grid.
##
class Chunk:
	var _multimesh_instances: Array[MultiMeshInstance] = []

	var _meshes: Array[Mesh] = []

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

	##
	## Invalidates the mesh block, re-baking its contents from the current mesh data set.
	##
	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))))

		# (Re)-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)

	##
	## Sets the mesh location in the chunk at the relative [code]coordinatess[/code] to
	## [code]mesh[/code].
	##
	## *Note* that [method Chunk.invalidate] must be called at some point after to visually update
	## the chunk.
	##
	func set_mesh(coordinates: Vector2i, mesh: Mesh) -> void:
		_meshes[(_CHUNK_SIZE * coordinates.y) + coordinates.x] = mesh

##
## Specialized multi-mesh instance convenience for use within the mesh grid and its chunks.
##
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])

	##
	## Sets the parent transform of all mesh instances under it to [code]offset[/code].
	##
	func set_offset(offset: Transform3D) -> void:
		RenderingServer.instance_set_transform(_instance_rid, offset)

var _chunks: Array[Chunk] = []

var _chunks_size := Vector2i.ZERO

##
## Size of the mesh grid (in engine units).
##
@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_area() -> Rect2i:
	return Rect2i(Vector2i.ZERO, size)

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)

##
## Clears the entire mesh grid to only contain [code]mesh[/code].
##
## [code]null[/code] may be past to [code]mesh[/code] for clearing the mesh grid to nothing.
##
## For clearing a specific region of the mesh grid, see [method fill_mesh].
##
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)

##
## Clears the region of the mesh grid at [code]area[/code] to only contain [code]mesh[/code].
##
## [code]null[/code] may be past to [code]mesh[/code] for filling the region to nothing.
##
## *Note* that [code]area[/code] *must* be within the area of the of the mesh grid, where
## [code]Vector2i.ZERO[/code] is the top-left and [member size] [code]- 1[/code] is the bottom-
## right.
##
## For clearing the whole of the mesh grid, see [method clear_mesh].
##
func fill_mesh(area: Rect2i, mesh: Mesh) -> void:
	assert(_get_area().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)

##
## Plots a single mesh at [code]coordinates[/code] to be [code]mesh[/code], overwriting whatever it
## previously contained.
##
## [code]null[/code] may be past to [code]mesh[/code] for clearing the mesh grid to nothing.
##
## *Note* that [code]coordinates[/code] *must* be within the area of the of the mesh grid, where
## [code]Vector2i.ZERO[/code] is the top-left and [member size] [code]- 1[/code] is the bottom-
## right.
##
## For bulk-setting many meshes at once, see [method fill_mesh] and [method clear_mesh].
##
func plot_mesh(coordinates: Vector2i, mesh: Mesh) -> void:
	assert(_get_area().has_point(coordinates), "coordinate must be within grid")

	var chunk_coordinates := coordinates / _CHUNK_SIZE
	var chunk := _get_chunk(chunk_coordinates)

	chunk.set_mesh(coordinates % _CHUNK_SIZE, mesh)
	chunk.invalidate(self, chunk_coordinates)

##
## 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())