@tool
class_name TerrainGenerator
extends MeshInstance3D

enum PreviewMode { NONE, BIOME, SLOPE, MASK_INDEX, RAIL_DISTANCE, STATION_DISTANCE }

signal terrain_generated
signal props_data_ready(props_data: TerrainPropsData)
 
@export var noise_setting: FastNoiseLite:
	set(value):
		if noise_setting and noise_setting.changed.is_connected(_on_settings_changed):
			noise_setting.changed.disconnect(_on_settings_changed)
		noise_setting = value
		if noise_setting:
			noise_setting.changed.connect(_on_settings_changed)
		_request_update()

@export var material_setting: Material:
	set(value):
		material_setting = value
		if mesh:
			mesh.surface_set_material(0, material_setting)

@export_group("Terrain Size")
@export var terrain_size: Vector2 = Vector2(200, 200):
	set(value):
		terrain_size = value
		_request_update()

@export_range(2, 512) var subdivision: int = 128:
	set(value):
		subdivision = value
		_request_update()

@export var max_height: float = 50.0:
	set(value):
		max_height = value
		_request_update()

@export_group("Island Mask")
@export var island_mask: bool = true:
	set(value):
		island_mask = value
		_request_update()

@export_range(0.0, 1.0) var mask_falloff_start: float = 0.5:
	set(value):
		mask_falloff_start = value
		_request_update()

@export_range(0.1, 2.0) var mask_power: float = 1.5:
	set(value):
		mask_power = value
		_request_update()

@export_group("Control Layers")
@export var control_layers: TerrainControlLayers:
	set(value):
		if control_layers and control_layers.layers_changed.is_connected(_on_settings_changed):
			control_layers.layers_changed.disconnect(_on_settings_changed)
		control_layers = value
		if control_layers:
			control_layers.layers_changed.connect(_on_settings_changed)
		_request_update()

@export_group("Rail Constraints")
@export var rail_paths: Array[NodePath] = []:
	set(value):
		rail_paths = value
		_request_update()

@export var rail_influence_radius: float = 10.0:
	set(value):
		rail_influence_radius = value
		_request_update()

@export var rail_flatten_strength: float = 1.0:
	set(value):
		rail_flatten_strength = value
		_request_update()

@export_range(0.0, 45.0) var rail_max_slope: float = 5.0:
	set(value):
		rail_max_slope = value
		_request_update()

@export_group("Station Constraints")
@export var station_paths: Array[NodePath] = []:
	set(value):
		station_paths = value
		_request_update()

@export var station_influence_radius: float = 25.0:
	set(value):
		station_influence_radius = value
		_request_update()

@export var station_flatten_height: float = 0.0:
	set(value):
		station_flatten_height = value
		_request_update()

@export var use_station_y_as_height: bool = true:
	set(value):
		use_station_y_as_height = value
		_request_update()

@export_group("Preview")
@export var preview_mode: PreviewMode = PreviewMode.NONE:
	set(value):
		preview_mode = value
		_request_update()

@export var preview_mask_index: int = 0:
	set(value):
		preview_mask_index = value
		if preview_mode == PreviewMode.MASK_INDEX:
			_request_update()

@export_group("Actions")
@export var regenerate: bool = false:
	set(_value):
		generate_terrain()

@export var bake_collision: bool = false:
	set(_value):
		_bake_collision()

var _update_queued: bool = false
var _rail_spline_cache: Array[Curve3D] = []
var _station_positions: Array[Vector3] = []
var _props_data: TerrainPropsData


func _ready() -> void:
	if noise_setting:
		if not noise_setting.changed.is_connected(_on_settings_changed):
			noise_setting.changed.connect(_on_settings_changed)
	generate_terrain()


func _on_settings_changed() -> void:
	_request_update()


func _request_update() -> void:
	if _update_queued:
		return
	_update_queued = true
	call_deferred("_do_update")


func _do_update() -> void:
	_update_queued = false
	generate_terrain()


func generate_terrain() -> void:
	if not noise_setting:
		return
	
	_cache_constraints()
	
	var st := SurfaceTool.new()
	st.begin(Mesh.PRIMITIVE_TRIANGLES)
	
	var step_x: float = terrain_size.x / float(subdivision - 1)
	var step_z: float = terrain_size.y / float(subdivision - 1)
	var half_x: float = terrain_size.x / 2.0
	var half_z: float = terrain_size.y / 2.0
	
	var heights: Array[float] = []
	var biome_ids: Array[int] = []
	var slope_values: Array[float] = []
	var rail_distances: Array[float] = []
	var station_distances: Array[float] = []
	var mask_values: Array[Array] = []
	
	heights.resize(subdivision * subdivision)
	biome_ids.resize(subdivision * subdivision)
	slope_values.resize(subdivision * subdivision)
	rail_distances.resize(subdivision * subdivision)
	station_distances.resize(subdivision * subdivision)
	
	var mask_count := 0
	if control_layers:
		mask_count = control_layers.mask_layers.size()
	for i in range(mask_count):
		var arr: Array[float] = []
		arr.resize(subdivision * subdivision)
		mask_values.append(arr)
	
	# Pass 1: Generate base heights
	for z in range(subdivision):
		for x in range(subdivision):
			var world_x: float = x * step_x - half_x
			var world_z: float = z * step_z - half_z
			var idx: int = z * subdivision + x
			var uv := Vector2(float(x) / float(subdivision - 1), float(z) / float(subdivision - 1))
			
			var base_noise: float = (noise_setting.get_noise_2d(world_x, world_z) + 1.0) * 0.5
			var height: float = base_noise * max_height
			
			# Apply height profile from control layers
			if control_layers and control_layers.height_profile:
				height = control_layers.get_height_profile_value(base_noise) * max_height
			
			# Apply island mask
			if island_mask:
				height *= _calculate_island_mask(world_x, world_z, half_x, half_z)
			
			# Sample biome
			var biome_color := Color.WHITE
			if control_layers:
				biome_color = control_layers.get_biome_value(uv)
			biome_ids[idx] = _color_to_biome_id(biome_color)
			
			# Sample slope limit
			var slope_limit := 45.0
			if control_layers:
				slope_limit = control_layers.get_slope_limit(uv)
			slope_values[idx] = slope_limit
			
			# Sample masks
			for m in range(mask_count):
				mask_values[m][idx] = control_layers.get_mask_value(m, uv)
			
			heights[idx] = height
	
	# Pass 2: Apply rail constraints
	for z in range(subdivision):
		for x in range(subdivision):
			var world_x: float = x * step_x - half_x
			var world_z: float = z * step_z - half_z
			var idx: int = z * subdivision + x
			var world_pos := Vector3(world_x, heights[idx], world_z)
			
			var rail_result := _get_rail_constraint(world_pos)
			rail_distances[idx] = rail_result.distance
			
			if rail_result.distance < rail_influence_radius and rail_result.distance >= 0.0:
				var t: float = 1.0 - (rail_result.distance / rail_influence_radius)
				t = t * t * rail_flatten_strength
				heights[idx] = lerpf(heights[idx], rail_result.height, t)
	
	# Pass 3: Apply station constraints
	for z in range(subdivision):
		for x in range(subdivision):
			var world_x: float = x * step_x - half_x
			var world_z: float = z * step_z - half_z
			var idx: int = z * subdivision + x
			var world_pos := Vector3(world_x, heights[idx], world_z)
			
			var station_result := _get_station_constraint(world_pos)
			station_distances[idx] = station_result.distance
			
			if station_result.distance < station_influence_radius and station_result.distance >= 0.0:
				var t: float = 1.0 - (station_result.distance / station_influence_radius)
				t = t * t
				heights[idx] = lerpf(heights[idx], station_result.height, t)
	
	# Pass 4: Apply slope limits
	_apply_slope_limits(heights, slope_values, step_x, step_z)
	
	# Build props data
	_props_data = TerrainPropsData.new()
	_props_data.resolution = Vector2i(subdivision, subdivision)
	_props_data.terrain_size = terrain_size
	_props_data.heights = heights
	_props_data.biome_ids = biome_ids
	_props_data.slope_values = slope_values
	_props_data.rail_distances = rail_distances
	_props_data.station_distances = station_distances
	_props_data.mask_values = mask_values
	
	# Build mesh
	for z in range(subdivision - 1):
		for x in range(subdivision - 1):
			var idx00: int = z * subdivision + x
			var idx10: int = z * subdivision + (x + 1)
			var idx01: int = (z + 1) * subdivision + x
			var idx11: int = (z + 1) * subdivision + (x + 1)
			
			var x0: float = x * step_x - half_x
			var x1: float = (x + 1) * step_x - half_x
			var z0: float = z * step_z - half_z
			var z1: float = (z + 1) * step_z - half_z
			
			var v00 := Vector3(x0, heights[idx00], z0)
			var v10 := Vector3(x1, heights[idx10], z0)
			var v01 := Vector3(x0, heights[idx01], z1)
			var v11 := Vector3(x1, heights[idx11], z1)
			
			var uv00 := Vector2(x0 / terrain_size.x + 0.5, z0 / terrain_size.y + 0.5)
			var uv10 := Vector2(x1 / terrain_size.x + 0.5, z0 / terrain_size.y + 0.5)
			var uv01 := Vector2(x0 / terrain_size.x + 0.5, z1 / terrain_size.y + 0.5)
			var uv11 := Vector2(x1 / terrain_size.x + 0.5, z1 / terrain_size.y + 0.5)
			
			# Preview mode coloring
			var c00 := _get_preview_color(idx00, biome_ids, slope_values, rail_distances, station_distances, mask_values)
			var c10 := _get_preview_color(idx10, biome_ids, slope_values, rail_distances, station_distances, mask_values)
			var c01 := _get_preview_color(idx01, biome_ids, slope_values, rail_distances, station_distances, mask_values)
			var c11 := _get_preview_color(idx11, biome_ids, slope_values, rail_distances, station_distances, mask_values)
			
			# Triangle 1: v00, v10, v01 (CCW winding, normal pointing up)
			var normal1: Vector3 = (v01 - v00).cross(v10 - v00).normalized()
			st.set_normal(normal1)
			st.set_color(c00)
			st.set_uv(uv00)
			st.add_vertex(v00)
			st.set_color(c10)
			st.set_uv(uv10)
			st.add_vertex(v10)
			st.set_color(c01)
			st.set_uv(uv01)
			st.add_vertex(v01)
			
			# Triangle 2: v10, v11, v01 (CCW winding, normal pointing up)
			var normal2: Vector3 = (v01 - v10).cross(v11 - v10).normalized()
			st.set_normal(normal2)
			st.set_color(c10)
			st.set_uv(uv10)
			st.add_vertex(v10)
			st.set_color(c11)
			st.set_uv(uv11)
			st.add_vertex(v11)
			st.set_color(c01)
			st.set_uv(uv01)
			st.add_vertex(v01)
	
	mesh = st.commit()
	
	if material_setting:
		mesh.surface_set_material(0, material_setting)
	
	terrain_generated.emit()
	props_data_ready.emit(_props_data)


func _cache_constraints() -> void:
	_rail_spline_cache.clear()
	_station_positions.clear()
	
	for path in rail_paths:
		if path.is_empty():
			continue
		var node := get_node_or_null(path)
		if node is Path3D:
			var path3d := node as Path3D
			if path3d.curve:
				_rail_spline_cache.append(path3d.curve)
	
	for path in station_paths:
		if path.is_empty():
			continue
		var node := get_node_or_null(path)
		if node is Node3D:
			_station_positions.append((node as Node3D).global_position)


func _get_rail_constraint(world_pos: Vector3) -> Dictionary:
	var min_dist := INF
	var rail_height := world_pos.y
	
	for curve in _rail_spline_cache:
		var closest := curve.get_closest_point(world_pos)
		var dist := Vector2(world_pos.x - closest.x, world_pos.z - closest.z).length()
		if dist < min_dist:
			min_dist = dist
			rail_height = closest.y
	
	return {"distance": min_dist, "height": rail_height}


func _get_station_constraint(world_pos: Vector3) -> Dictionary:
	var min_dist := INF
	var target_height := station_flatten_height
	
	for station_pos in _station_positions:
		var dist := Vector2(world_pos.x - station_pos.x, world_pos.z - station_pos.z).length()
		if dist < min_dist:
			min_dist = dist
			if use_station_y_as_height:
				target_height = station_pos.y
	
	return {"distance": min_dist, "height": target_height}


func _apply_slope_limits(heights: Array[float], slope_limits: Array[float], step_x: float, step_z: float) -> void:
	var max_iterations := 10
	for _iter in range(max_iterations):
		var changed := false
		for z in range(1, subdivision - 1):
			for x in range(1, subdivision - 1):
				var idx: int = z * subdivision + x
				var slope_limit_rad: float = deg_to_rad(slope_limits[idx])
				var max_delta_x: float = tan(slope_limit_rad) * step_x
				var max_delta_z: float = tan(slope_limit_rad) * step_z
				
				var h := heights[idx]
				var h_xm := heights[idx - 1]
				var h_xp := heights[idx + 1]
				var h_zm := heights[idx - subdivision]
				var h_zp := heights[idx + subdivision]
				
				var target := h
				if abs(h - h_xm) > max_delta_x:
					target = minf(target, h_xm + max_delta_x) if h > h_xm else maxf(target, h_xm - max_delta_x)
					changed = true
				if abs(h - h_xp) > max_delta_x:
					target = minf(target, h_xp + max_delta_x) if h > h_xp else maxf(target, h_xp - max_delta_x)
					changed = true
				if abs(h - h_zm) > max_delta_z:
					target = minf(target, h_zm + max_delta_z) if h > h_zm else maxf(target, h_zm - max_delta_z)
					changed = true
				if abs(h - h_zp) > max_delta_z:
					target = minf(target, h_zp + max_delta_z) if h > h_zp else maxf(target, h_zp - max_delta_z)
					changed = true
				
				heights[idx] = lerpf(h, target, 0.5)
		
		if not changed:
			break


func _color_to_biome_id(color: Color) -> int:
	return int(color.r * 255.0) + int(color.g * 255.0) * 256


func _get_preview_color(idx: int, biome_ids: Array[int], slope_values: Array[float], 
		rail_distances: Array[float], station_distances: Array[float], mask_values: Array[Array]) -> Color:
	match preview_mode:
		PreviewMode.NONE:
			return Color.WHITE
		PreviewMode.BIOME:
			var biome := biome_ids[idx]
			return Color.from_hsv(fmod(float(biome) * 0.618033988749, 1.0), 0.7, 0.9)
		PreviewMode.SLOPE:
			var slope := slope_values[idx] / 90.0
			return Color(slope, 1.0 - slope, 0.0)
		PreviewMode.MASK_INDEX:
			if preview_mask_index < mask_values.size():
				var v: float = mask_values[preview_mask_index][idx]
				return Color(v, v, v)
			return Color.BLACK
		PreviewMode.RAIL_DISTANCE:
			var d: float = clampf(rail_distances[idx] / (rail_influence_radius * 2.0), 0.0, 1.0)
			return Color(1.0 - d, 0.2, d)
		PreviewMode.STATION_DISTANCE:
			var d: float = clampf(station_distances[idx] / (station_influence_radius * 2.0), 0.0, 1.0)
			return Color(0.2, 1.0 - d, d)
	return Color.WHITE


func get_props_data() -> TerrainPropsData:
	return _props_data


func _calculate_island_mask(x: float, z: float, half_x: float, half_z: float) -> float:
	var nx: float = x / half_x
	var nz: float = z / half_z
	var dist: float = sqrt(nx * nx + nz * nz)
	
	if dist < mask_falloff_start:
		return 1.0
	elif dist > 1.0:
		return 0.0
	else:
		var t: float = (dist - mask_falloff_start) / (1.0 - mask_falloff_start)
		return pow(1.0 - t, mask_power)


func _bake_collision() -> void:
	if not mesh:
		push_warning("TerrainGenerator: No mesh to bake collision from")
		return
	
	var static_body: StaticBody3D
	var collision_shape: CollisionShape3D
	
	for child in get_children():
		if child is StaticBody3D:
			static_body = child
			break
	
	if not static_body:
		static_body = StaticBody3D.new()
		static_body.name = "TerrainCollision"
		add_child(static_body)
		if Engine.is_editor_hint():
			static_body.owner = get_tree().edited_scene_root
	
	for child in static_body.get_children():
		if child is CollisionShape3D:
			collision_shape = child
			break
	
	if not collision_shape:
		collision_shape = CollisionShape3D.new()
		collision_shape.name = "CollisionShape"
		static_body.add_child(collision_shape)
		if Engine.is_editor_hint():
			collision_shape.owner = get_tree().edited_scene_root
	
	collision_shape.shape = mesh.create_trimesh_shape()
	print("TerrainGenerator: Collision baked successfully")