extends Node3D

const Watercourse := preload("res://world/jajce/JajceWatercourse.gd")

@export var terrain_path := NodePath("../../TerrainRoot/Terrain3D")
@export var river_material: Material
@export_range(2, 12, 1) var width_segments := 6
@export_range(8, 96, 1) var length_segments := 36
@export_range(8, 64, 1) var upper_length_segments := 18
@export_range(0.05, 0.5, 0.01) var water_offset := Watercourse.WATER_SURFACE_OFFSET

@onready var river_ribbon: MeshInstance3D = $RiverRibbon
@onready var upper_river_ribbon: MeshInstance3D = $UpperRiverRibbon


func _ready() -> void:
	call_deferred("_build_terrain_ribbon")


func _build_terrain_ribbon() -> void:
	var terrain := get_node_or_null(terrain_path) as Terrain3D
	if terrain == null or terrain.data == null or river_material == null:
		push_warning("TerrainRiverRibbon: terrain or material is unavailable")
		return

	river_ribbon.mesh = _build_ribbon(
		terrain,
		Watercourse.DOWNSTREAM_START_Z,
		Watercourse.DOWNSTREAM_END_Z,
		length_segments,
		false
	)
	upper_river_ribbon.mesh = _build_ribbon(
		terrain,
		Watercourse.UPPER_STREAM_START_Z,
		Watercourse.UPPER_STREAM_END_Z,
		upper_length_segments,
		true
	)
	_align_pool_to_terrain(terrain)


func _build_ribbon(
	terrain: Terrain3D, start_z: float, end_z: float, segment_count: int, is_upper: bool
) -> ArrayMesh:
	var vertices := PackedVector3Array()
	var normals := PackedVector3Array()
	var uvs := PackedVector2Array()
	var indices := PackedInt32Array()
	var row_size := width_segments + 1
	var pool := $PlungePool as MeshInstance3D
	var pool_mesh := pool.mesh as PlaneMesh
	var pool_radii := pool_mesh.size * Vector2(pool.scale.x, pool.scale.z) * 0.5
	for z_index in segment_count + 1:
		var progress := float(z_index) / float(segment_count)
		var world_z := lerpf(start_z, end_z, progress)
		var center_x := (
			Watercourse.upper_stream_center_x(world_z)
			if is_upper
			else Watercourse.downstream_center_x(world_z)
		)
		var half_width := (
			Watercourse.upper_stream_half_width(world_z)
			if is_upper
			else Watercourse.downstream_half_width(world_z)
		)
		if not is_upper:
			# The outlet continues the existing pool silhouette without drawing a
			# second, coplanar sheet on top of the pool's transparent material.
			var pool_progress := (world_z - Watercourse.PLUNGE_POOL_CENTER.y) / pool_radii.y
			if absf(pool_progress) < 1.0:
				var pool_width := pool_radii.x * sqrt(1.0 - pool_progress * pool_progress)
				half_width = maxf(half_width, pool_width)
		# A stream has one surface across its width. Sampling each bank separately
		# wraps the water mesh up the terrain and makes it read as a solid ramp.
		var terrain_height := terrain.data.get_height(Vector3(center_x, 0.0, world_z))
		if is_nan(terrain_height):
			terrain_height = global_position.y
		var surface_height := terrain_height + water_offset
		var left_x := _shoreline_x(
			terrain, center_x, center_x - half_width, world_z, surface_height
		)
		var right_x := _shoreline_x(
			terrain, center_x, center_x + half_width, world_z, surface_height
		)
		for x_index in width_segments + 1:
			var across := float(x_index) / float(width_segments)
			var world_x := lerpf(left_x, right_x, across)
			var world_sample := Vector3(world_x, 0.0, world_z)
			vertices.append(to_local(Vector3(world_sample.x, surface_height, world_sample.z)))
			normals.append(Vector3.UP)
			uvs.append(Vector2(across, progress))

	for z_index in segment_count:
		for x_index in width_segments:
			var top_left := z_index * row_size + x_index
			var top_right := top_left + 1
			var bottom_left := top_left + row_size
			var bottom_right := bottom_left + 1
			indices.append_array(
				[top_left, bottom_left, top_right, top_right, bottom_left, bottom_right]
			)

	var arrays := []
	arrays.resize(Mesh.ARRAY_MAX)
	arrays[Mesh.ARRAY_VERTEX] = vertices
	arrays[Mesh.ARRAY_NORMAL] = normals
	arrays[Mesh.ARRAY_TEX_UV] = uvs
	arrays[Mesh.ARRAY_INDEX] = indices
	var mesh := ArrayMesh.new()
	mesh.add_surface_from_arrays(Mesh.PRIMITIVE_TRIANGLES, arrays)
	mesh.surface_set_material(0, river_material)
	return mesh


func _shoreline_x(
	terrain: Terrain3D, center_x: float, edge_x: float, world_z: float, surface_height: float
) -> float:
	var edge_height := terrain.data.get_height(Vector3(edge_x, 0.0, world_z))
	if is_nan(edge_height) or edge_height <= surface_height:
		return edge_x
	# The carved basin is wider on one side. Stop its visible surface at the
	# shore instead of extending the pool's ideal ellipse metres under a hill.
	var wet_x := center_x
	var dry_x := edge_x
	for step in range(1, 17):
		var sample_x := lerpf(center_x, edge_x, float(step) / 16.0)
		var sample_height := terrain.data.get_height(Vector3(sample_x, 0.0, world_z))
		if is_nan(sample_height) or sample_height > surface_height:
			dry_x = sample_x
			break
		wet_x = sample_x
	for _step in 10:
		var sample_x := (wet_x + dry_x) * 0.5
		var sample_height := terrain.data.get_height(Vector3(sample_x, 0.0, world_z))
		if is_nan(sample_height) or sample_height > surface_height:
			dry_x = sample_x
		else:
			wet_x = sample_x
	return wet_x


func _align_pool_to_terrain(terrain: Terrain3D) -> void:
	var pool := $PlungePool as MeshInstance3D
	var pool_material := pool.mesh.surface_get_material(0) as ShaderMaterial
	pool_material.set_shader_parameter("pool_outlet_z", Watercourse.DOWNSTREAM_START_Z)
	var pool_sample := Vector3(
		Watercourse.PLUNGE_POOL_CENTER.x, 0.0, Watercourse.PLUNGE_POOL_CENTER.y
	)
	var pool_height := terrain.data.get_height(pool_sample)
	if is_nan(pool_height):
		return
	var local_surface := to_local(Vector3(pool_sample.x, pool_height + water_offset, pool_sample.z))
	pool.position = local_surface
	var foam_ring := $FoamRing as MeshInstance3D
	foam_ring.position = local_surface + Vector3(0.0, 0.1, 0.0)
	var foam_at_fall := $FoamAtFall as MeshInstance3D
	foam_at_fall.position = local_surface + Vector3(0.0, 0.13, -2.7)