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
	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)
		)
		for x_index in width_segments + 1:
			var across := float(x_index) / float(width_segments)
			var world_x := center_x + lerpf(-half_width, half_width, across)
			var world_sample := Vector3(world_x, 0.0, world_z)
			var terrain_height := terrain.data.get_height(world_sample)
			if is_nan(terrain_height):
				terrain_height = global_position.y
			vertices.append(
				to_local(Vector3(world_sample.x, terrain_height + water_offset, 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 _align_pool_to_terrain(terrain: Terrain3D) -> void:
	var pool := $PlungePool as MeshInstance3D
	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)