class_name TileGrid 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 _floor_meshes: Array[Mesh] = []

	var _multimesh_instances: Array[MultiMeshInstance] = []

	var _wall_meshes: Array[Mesh] = []

	func _init() -> void:
		var buffer_size := _CHUNK_SIZE * _CHUNK_SIZE

		_floor_meshes.resize(buffer_size)
		_wall_meshes.resize(buffer_size)

	##
	## Invalidates the mesh block, re-baking its contents from the current mesh data set.
	##
	func invalidate(tile_grid: TileGrid, 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 := tile_grid.size

		for i in _floor_meshes.size():
			var mesh := _floor_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))))

		for i in _wall_meshes.size():
			var mesh := _wall_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 := tile_grid.get_world_3d().scenario
		var global_transform := tile_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 floor mesh in the chunk at the location 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_floor_mesh(coordinates: Vector2i, mesh: Mesh) -> void:
		_floor_meshes[(_CHUNK_SIZE * coordinates.y) + coordinates.x] = mesh

	##
	## Sets the wall mesh in the chunk at the location 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_wall_mesh(coordinates: Vector2i, mesh: Mesh) -> void:
		_wall_meshes[(_CHUNK_SIZE * coordinates.y) + coordinates.x] = mesh

##
## Specialized multi-mesh instance convenience for use within the tile 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 tile 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_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 entirety of the tile grid to only contain [code]tile[/code].
##
## For clearing a specific region of the mesh grid, see [method fill_mesh].
##
func clear_tiles(tile: Tile) -> void:
	if tile == null:
		for y in size.y:
			for x in size.x:
				var coordinates := Vector2i(x, y)
				var chunk := _get_chunk(coordinates / _CHUNK_SIZE)
				var chunk_coordinates := coordinates % _CHUNK_SIZE

				chunk.set_floor_mesh(chunk_coordinates, null)
				chunk.set_wall_mesh(chunk_coordinates, null)

	else:
		var mesh := tile.mesh

		match tile.kind:
			Tile.Kind.FLOOR:
				for y in size.y:
					for x in size.x:
						var coordinate := Vector2i(x, y)

						_get_chunk(coordinate / _CHUNK_SIZE).\
							set_floor_mesh(coordinate % _CHUNK_SIZE, mesh)

			Tile.Kind.WALL:
				for y in size.y:
					for x in size.x:
						var coordinate := Vector2i(x, y)

						_get_chunk(coordinate / _CHUNK_SIZE).\
							set_floor_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 tile grid at [code]area[/code] to only contain [code]tile[/code].
##
## *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_tiles(area: Rect2i, tile: Tile) -> void:
	assert(get_area().encloses(area), "area must be within grid")

	var filled_chunks := BitMap.new()

	filled_chunks.resize(_chunks_size)

	if tile == null:
		for y in range(area.position.y, area.end.y):
			for x in range(area.position.x, area.end.x):
				var coordinates := Vector2i(x, y)
				var chunk_coordinates := coordinates / _CHUNK_SIZE
				var chunk := _get_chunk(coordinates / _CHUNK_SIZE)
				var local_coordinates := coordinates % _CHUNK_SIZE

				chunk.set_floor_mesh(local_coordinates, null)
				chunk.set_wall_mesh(local_coordinates, null)
				filled_chunks.set_bitv(chunk_coordinates, true)

	else:
		var mesh := tile.mesh

		match tile.kind:
			Tile.Kind.FLOOR:
				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_floor_mesh(coordinate % _CHUNK_SIZE, mesh)
						filled_chunks.set_bitv(chunk_coordinate, true)

			Tile.Kind.WALL:
				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_wall_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 get_area() -> Rect2i:
	return Rect2i(Vector2i.ZERO, size)

##
## Plots a single tile at [code]coordinates[/code] to be [code]tile[/code], overwriting whatever it
## previously contained.
##
## *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_tile(coordinates: Vector2i, tile: Tile) -> void:
	assert(get_area().has_point(coordinates), "coordinate must be within grid")

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

	if tile == null:
		var local_coordinates := coordinates % _CHUNK_SIZE

		chunk.set_floor_mesh(local_coordinates, null)
		chunk.set_wall_mesh(local_coordinates, null)
		chunk.invalidate(self, chunk_coordinates)

	else:
		var mesh := tile.mesh

		match tile.kind:
			Tile.Kind.FLOOR:
				chunk.set_floor_mesh(coordinates % _CHUNK_SIZE, mesh)
				chunk.invalidate(self, chunk_coordinates)

			Tile.Kind.WALL:
				chunk.set_wall_mesh(coordinates % _CHUNK_SIZE, mesh)
				chunk.invalidate(self, chunk_coordinates)