class_name SonarReliefRenderer
extends SubViewport

const SonarModelScript = preload("res://src/sonar_model.gd")

@export var terrain_skirt_depth := 4.5

var _camera: Camera3D
var _current_mesh: MeshInstance3D
var _previous_mesh: MeshInstance3D
var _current_anomaly: MeshInstance3D
var _previous_anomaly: MeshInstance3D
var _submarine_marker: Node3D
var _current_cells: Array[Dictionary] = []
var _previous_cells: Array[Dictionary] = []
var _last_built_cell_count := -1


func _ready() -> void:
	size = Vector2i(960, 540)
	disable_3d = false
	render_target_update_mode = SubViewport.UPDATE_ALWAYS
	world_3d = World3D.new()
	_build_world()


func start_scan() -> void:
	_previous_cells = _current_cells.duplicate(true)
	_current_cells.clear()
	_last_built_cell_count = -1
	_rebuild_terrain(_previous_mesh, _previous_cells)
	_rebuild_terrain(_current_mesh, _current_cells)
	_update_anomaly(_previous_anomaly, _previous_cells)
	_update_anomaly(_current_anomaly, _current_cells)
	_previous_mesh.transparency = 0.48 if not _previous_cells.is_empty() else 1.0
	_previous_anomaly.transparency = _previous_mesh.transparency


func update_scan(
	current_cells: Array[Dictionary],
	previous_cells: Array[Dictionary],
	progress: float
) -> void:
	_current_cells = current_cells.duplicate(true)
	_previous_cells = previous_cells.duplicate(true)
	if _current_cells.size() != _last_built_cell_count:
		_last_built_cell_count = _current_cells.size()
		_rebuild_terrain(_current_mesh, _current_cells)
		_update_anomaly(_current_anomaly, _current_cells)
	if _previous_mesh.mesh == null and not _previous_cells.is_empty():
		_rebuild_terrain(_previous_mesh, _previous_cells)
		_update_anomaly(_previous_anomaly, _previous_cells)
	var previous_alpha := lerpf(0.48, 1.0, clampf(progress, 0.0, 1.0))
	_previous_mesh.transparency = previous_alpha
	_previous_anomaly.transparency = previous_alpha


func finish_scan(current_cells: Array[Dictionary]) -> void:
	_current_cells = current_cells.duplicate(true)
	_last_built_cell_count = _current_cells.size()
	_rebuild_terrain(_current_mesh, _current_cells)
	_update_anomaly(_current_anomaly, _current_cells)
	_previous_cells.clear()
	_previous_mesh.mesh = null
	_previous_mesh.transparency = 1.0
	_previous_anomaly.visible = false
	_previous_anomaly.transparency = 1.0


func update_submarine_pose(submarine_transform: Transform3D) -> void:
	if _camera == null:
		return
	var forward := -submarine_transform.basis.z
	forward.y = 0.0
	forward = forward.normalized()
	var camera_position := submarine_transform.origin - forward * 30.0 + Vector3.UP * 65.0
	var focus := submarine_transform.origin + forward * 43.0 + Vector3.DOWN * 8.0
	_camera.global_position = camera_position
	_camera.look_at(focus, Vector3.UP)
	_submarine_marker.global_position = submarine_transform.origin
	_submarine_marker.rotation = Vector3(0.0, submarine_transform.basis.get_euler().y, 0.0)


func current_face_count() -> int:
	if _current_mesh.mesh == null:
		return 0
	return int(_current_mesh.mesh.get_faces().size() / 3)


func _build_world() -> void:
	var environment_node := WorldEnvironment.new()
	var environment := Environment.new()
	environment.background_mode = Environment.BG_COLOR
	environment.background_color = Color("020407")
	environment.ambient_light_source = Environment.AMBIENT_SOURCE_COLOR
	environment.ambient_light_color = Color("8fc9d1")
	environment.ambient_light_energy = 0.42
	environment.tonemap_mode = Environment.TONE_MAPPER_FILMIC
	environment_node.environment = environment
	add_child(environment_node)

	_previous_mesh = _make_terrain_instance("PreviousBathymetry")
	_current_mesh = _make_terrain_instance("CurrentBathymetry")
	_previous_anomaly = _make_anomaly_instance("PreviousAnomaly")
	_current_anomaly = _make_anomaly_instance("CurrentAnomaly")

	var key_light := DirectionalLight3D.new()
	key_light.name = "ReliefKeyLight"
	key_light.rotation_degrees = Vector3(-58.0, -34.0, 0.0)
	key_light.light_color = Color("d8fbff")
	key_light.light_energy = 1.25
	key_light.shadow_enabled = true
	add_child(key_light)

	var fill_light := DirectionalLight3D.new()
	fill_light.name = "ReliefFillLight"
	fill_light.rotation_degrees = Vector3(-35.0, 145.0, 0.0)
	fill_light.light_color = Color("3d86ad")
	fill_light.light_energy = 0.38
	add_child(fill_light)

	_camera = Camera3D.new()
	_camera.name = "BathymetryCamera"
	_camera.projection = Camera3D.PROJECTION_ORTHOGONAL
	_camera.size = 116.0
	_camera.far = 240.0
	_camera.current = true
	add_child(_camera)

	_build_submarine_marker()


func _make_terrain_instance(node_name: String) -> MeshInstance3D:
	var instance := MeshInstance3D.new()
	instance.name = node_name
	var material := StandardMaterial3D.new()
	material.vertex_color_use_as_albedo = true
	material.roughness = 0.84
	material.metallic = 0.0
	material.transparency = BaseMaterial3D.TRANSPARENCY_ALPHA
	# The scan may be inspected while the submarine pitches or passes above it.
	# Rendering both sides keeps the reconstructed sheet readable from every
	# permitted sonar-camera angle.
	material.cull_mode = BaseMaterial3D.CULL_DISABLED
	instance.material_override = material
	add_child(instance)
	return instance


func _make_anomaly_instance(node_name: String) -> MeshInstance3D:
	var instance := MeshInstance3D.new()
	instance.name = node_name
	var crate := BoxMesh.new()
	crate.size = Vector3(8.0, 4.5, 8.0)
	instance.mesh = crate
	var material := StandardMaterial3D.new()
	material.shading_mode = BaseMaterial3D.SHADING_MODE_UNSHADED
	material.albedo_color = Color("ff9d1f")
	material.emission_enabled = true
	material.emission = Color("ff6800")
	material.emission_energy_multiplier = 4.0
	material.roughness = 0.32
	material.transparency = BaseMaterial3D.TRANSPARENCY_ALPHA
	instance.material_override = material
	instance.visible = false
	add_child(instance)
	return instance


func _build_submarine_marker() -> void:
	_submarine_marker = Node3D.new()
	_submarine_marker.name = "SubmarineReference"
	add_child(_submarine_marker)

	var body := MeshInstance3D.new()
	var body_mesh := CapsuleMesh.new()
	body_mesh.radius = 1.15
	body_mesh.height = 5.8
	body_mesh.radial_segments = 16
	body_mesh.rings = 6
	body.mesh = body_mesh
	body.rotation_degrees.x = 90.0
	var body_material := StandardMaterial3D.new()
	body_material.albedo_color = Color("eaffff")
	body_material.emission_enabled = true
	body_material.emission = Color("a7f8ff")
	body_material.emission_energy_multiplier = 1.45
	body.material_override = body_material
	_submarine_marker.add_child(body)

	var nose := MeshInstance3D.new()
	var nose_mesh := CylinderMesh.new()
	nose_mesh.top_radius = 0.0
	nose_mesh.bottom_radius = 0.72
	nose_mesh.height = 1.6
	nose_mesh.radial_segments = 16
	nose.mesh = nose_mesh
	nose.position.z = -3.5
	nose.rotation_degrees.x = -90.0
	nose.material_override = body_material
	_submarine_marker.add_child(nose)


func _rebuild_terrain(instance: MeshInstance3D, cells: Array[Dictionary]) -> void:
	var samples: Dictionary = {}
	for cell: Dictionary in cells:
		if String(cell.get("sample_type", "bathymetry")) != "bathymetry":
			continue
		samples[Vector2i(int(cell["column"]), int(cell["row"]))] = cell
	if samples.size() < 4:
		instance.mesh = null
		return

	var keys: Array[Vector2i] = []
	for key: Vector2i in samples:
		keys.append(key)
	keys.sort_custom(func(a: Vector2i, b: Vector2i) -> bool:
		return a.y < b.y or (a.y == b.y and a.x < b.x)
	)

	var vertices: Array[Vector3] = []
	var colors: Array[Color] = []
	var normals: Array[Vector3] = []
	var indices: Array[int] = []
	var vertex_for_key: Dictionary = {}
	for key: Vector2i in keys:
		var cell: Dictionary = samples[key]
		vertex_for_key[key] = vertices.size()
		vertices.append(cell["position"])
		colors.append(_bathymetry_color(cell))
		normals.append(Vector3.ZERO)

	for key: Vector2i in keys:
		var right_key := key + Vector2i.RIGHT
		var far_key := key + Vector2i.DOWN
		var far_right_key := key + Vector2i(1, 1)
		if not (
			vertex_for_key.has(right_key)
			and vertex_for_key.has(far_key)
			and vertex_for_key.has(far_right_key)
		):
			continue
		_append_top_triangle(
			int(vertex_for_key[key]),
			int(vertex_for_key[right_key]),
			int(vertex_for_key[far_key]),
			vertices,
			normals,
			indices
		)
		_append_top_triangle(
			int(vertex_for_key[right_key]),
			int(vertex_for_key[far_right_key]),
			int(vertex_for_key[far_key]),
			vertices,
			normals,
			indices
		)

	if indices.is_empty():
		instance.mesh = null
		return
	for normal_index in normals.size():
		normals[normal_index] = normals[normal_index].normalized()

	_append_terrain_skirt(samples, vertices, normals, colors, indices)
	var arrays: Array = []
	arrays.resize(Mesh.ARRAY_MAX)
	arrays[Mesh.ARRAY_VERTEX] = PackedVector3Array(vertices)
	arrays[Mesh.ARRAY_NORMAL] = PackedVector3Array(normals)
	arrays[Mesh.ARRAY_COLOR] = PackedColorArray(colors)
	arrays[Mesh.ARRAY_INDEX] = PackedInt32Array(indices)
	var mesh := ArrayMesh.new()
	mesh.add_surface_from_arrays(Mesh.PRIMITIVE_TRIANGLES, arrays)
	instance.mesh = mesh


func _append_top_triangle(
	first: int,
	second: int,
	third: int,
	vertices: Array[Vector3],
	normals: Array[Vector3],
	indices: Array[int]
) -> void:
	var face_normal := (vertices[second] - vertices[first]).cross(
		vertices[third] - vertices[first]
	).normalized()
	if face_normal.y < 0.0:
		var swapped := second
		second = third
		third = swapped
		face_normal = -face_normal
	indices.append_array([first, second, third])
	normals[first] += face_normal
	normals[second] += face_normal
	normals[third] += face_normal


func _append_terrain_skirt(
	samples: Dictionary,
	vertices: Array[Vector3],
	normals: Array[Vector3],
	colors: Array[Color],
	indices: Array[int]
) -> void:
	var row_bounds := _sample_bounds(samples)
	if row_bounds.is_empty():
		return
	var minimum_row: int = row_bounds["min_row"]
	var maximum_row: int = row_bounds["max_row"]
	var minimum_column: int = row_bounds["min_column"]
	var maximum_column: int = row_bounds["max_column"]
	for column in range(minimum_column, maximum_column):
		_append_skirt_between(
			samples,
			Vector2i(column, minimum_row),
			Vector2i(column + 1, minimum_row),
			vertices,
			normals,
			colors,
			indices
		)
		_append_skirt_between(
			samples,
			Vector2i(column + 1, maximum_row),
			Vector2i(column, maximum_row),
			vertices,
			normals,
			colors,
			indices
		)
	for row in range(minimum_row, maximum_row):
		_append_skirt_between(
			samples,
			Vector2i(minimum_column, row + 1),
			Vector2i(minimum_column, row),
			vertices,
			normals,
			colors,
			indices
		)
		_append_skirt_between(
			samples,
			Vector2i(maximum_column, row),
			Vector2i(maximum_column, row + 1),
			vertices,
			normals,
			colors,
			indices
		)


func _append_skirt_between(
	samples: Dictionary,
	first_key: Vector2i,
	second_key: Vector2i,
	vertices: Array[Vector3],
	normals: Array[Vector3],
	colors: Array[Color],
	indices: Array[int]
) -> void:
	if not samples.has(first_key) or not samples.has(second_key):
		return
	var first: Vector3 = samples[first_key]["position"]
	var second: Vector3 = samples[second_key]["position"]
	var first_bottom := first + Vector3.DOWN * terrain_skirt_depth
	var second_bottom := second + Vector3.DOWN * terrain_skirt_depth
	var side_normal := (second - first).cross(first_bottom - first).normalized()
	var base_index := vertices.size()
	vertices.append_array([first, second, first_bottom, second_bottom])
	normals.append_array([side_normal, side_normal, side_normal, side_normal])
	var side_color := Color("152a32")
	colors.append_array([side_color, side_color, side_color, side_color])
	indices.append_array([
		base_index,
		base_index + 1,
		base_index + 2,
		base_index + 1,
		base_index + 3,
		base_index + 2,
	])


func _sample_bounds(samples: Dictionary) -> Dictionary:
	var minimum_row := 1 << 30
	var maximum_row := -(1 << 30)
	var minimum_column := 1 << 30
	var maximum_column := -(1 << 30)
	for key: Vector2i in samples:
		minimum_row = mini(minimum_row, key.y)
		maximum_row = maxi(maximum_row, key.y)
		minimum_column = mini(minimum_column, key.x)
		maximum_column = maxi(maximum_column, key.x)
	return {
		"min_row": minimum_row,
		"max_row": maximum_row,
		"min_column": minimum_column,
		"max_column": maximum_column,
	}


func _bathymetry_color(cell: Dictionary) -> Color:
	if int(cell["material"]) == SonarModelScript.MaterialCategory.METAL:
		return Color("ff6b22")
	var elevation := clampf(float(cell.get("elevation", 0.0)), 0.0, 1.0)
	var color: Color
	if elevation < 0.28:
		color = Color("173f91").lerp(Color("168fc1"), elevation / 0.28)
	elif elevation < 0.52:
		color = Color("168fc1").lerp(Color("36b979"), (elevation - 0.28) / 0.24)
	elif elevation < 0.76:
		color = Color("36b979").lerp(Color("d7d93d"), (elevation - 0.52) / 0.24)
	else:
		color = Color("d7d93d").lerp(Color("ee5734"), (elevation - 0.76) / 0.24)
	var return_strength := clampf(float(cell.get("intensity", 0.5)), 0.0, 1.0)
	return color.lerp(Color.WHITE, return_strength * 0.08)


func _update_anomaly(instance: MeshInstance3D, cells: Array[Dictionary]) -> void:
	var position_sum := Vector3.ZERO
	var anomaly_count := 0
	for cell: Dictionary in cells:
		if int(cell.get("material", SonarModelScript.MaterialCategory.ROCK)) != SonarModelScript.MaterialCategory.METAL:
			continue
		position_sum += cell["position"]
		anomaly_count += 1
	instance.visible = anomaly_count > 0
	if anomaly_count > 0:
		instance.global_position = position_sum / float(anomaly_count) + Vector3.UP * 1.4