@icon("res://addons/at-icons/node3d/steering_wheel.svg")
class_name Steering
extends Node3D

signal fallback_direction_used

enum DangerComputationMethod { RAY, MULTIRAY, SHAPE }

@export var body: CharacterBody3D
@export var gravity: bool = true
@export_category("Movement")
@export_range(1, 10, 0.1, "suffix:m/s") var movement_speed: float = 4.0
@export_range(0.5, 15, 0.1, "How many times per second the gap to target gets closed") var steering_speed: float = 0.2
## Apply forces to keep bodies apart
@export_category("Separation")
@export var separation_radius: float = 1.0
@export var separation_weight: float = 0.0
@export_category("Context Aware Steering")
@export var danger_method: DangerComputationMethod = DangerComputationMethod.RAY
@export var collision_shape: Shape3D
@export var ray_count: int = 16:
	set = set_ray_count
@export var look_ahead: float = 2.0
@export var ray_origin_offset: Vector3 = Vector3.UP
@export var steering_bias_seed: float = randf_range(-1.0, 1.0)
@export_group("Validation", "validation_")
## Project the actual body shape towards the chosen direction and validate any collision.
@export var validation_use_body_shape: bool = true
@export_category("Modifiers")
@export var _modifiers: Array[SpeedModifier] = []
@export_category("Debug")
@export var debug_draw_enabled := false:
	set(value):
		debug_draw_enabled = value
		if not value and _debug_draw:
			_debug_draw.clear()

var current_velocity: Vector3 = Vector3.ZERO
var _ray_directions: Array[Vector3] = []
var _interest: Array[float] = []
var _danger: Array[float] = []
var _last_chosen_direction := Vector3.MODEL_FRONT

@onready var _debug_draw: SteeringDebugDraw = $MeshInstance3D
@onready var area_3d: Area3D = $Area3D


func _ready() -> void:
	if not body:
		body = get_parent()
	set_ray_count(ray_count)


func update(delta: float, target_position: Vector3) -> Vector3:
	_modifiers = _modifiers.filter(func(mod: SpeedModifier) -> bool: return mod.tick(delta))

	var direction := target_position - body.global_position

	_compute_interest(direction)
	_compute_danger()

	var chosen_direction := _choose_direction(direction)
	if validation_use_body_shape:
		chosen_direction = _validate_chosen_direction(chosen_direction)

	var effective_speed := _get_effective_speed()
	var separation_force := get_separation_force(area_3d.get_overlapping_bodies(), separation_radius)

	var desired_velocity := chosen_direction.normalized() * effective_speed
	current_velocity = current_velocity.lerp(desired_velocity, 1.0 - exp(-steering_speed * delta))
	current_velocity += separation_force * separation_weight
	current_velocity += body.get_gravity() * delta

	if debug_draw_enabled and _debug_draw:
		_debug_draw.draw_steering(
			_ray_directions,
			_interest,
			_danger,
			chosen_direction,
			direction,
			look_ahead,
		)

	return current_velocity


func get_separation_force(neighbors: Array[Node3D], radius: float) -> Vector3:
	var force := Vector3.ZERO
	for neighbor in neighbors:
		var dist := body.global_position.distance_to(neighbor.global_position)
		if dist < radius and dist > 0.01:
			var away := (body.global_position - neighbor.global_position).normalized()
			force += away * (1.0 - dist / radius) # stronger when closer
	force.y = 0
	return force.normalized()


func add_modifier(modifier: SpeedModifier) -> void:
	_modifiers.append(modifier)


func remove_modifier(modifier: SpeedModifier) -> void:
	_modifiers.erase(modifier)


func set_ray_count(value: int) -> void:
	ray_count = value

	_ray_directions.resize(ray_count)

	var half_step := PI / ray_count # offset so no ray sits at exactly forward/back
	for i in ray_count:
		var angle := i * 2 * PI / ray_count + half_step
		_ray_directions[i] = Vector3.MODEL_FRONT.rotated(Vector3.UP, angle).normalized()


func _compute_interest(desired_direction: Vector3) -> void:
	_interest.resize(ray_count)
	for i in ray_count:
		var d := _ray_directions[i].dot(desired_direction)
		# To prevent left/right stalemates when obstacles straight ahead
		var bias := _ray_directions[i].cross(Vector3.UP).dot(Vector3.MODEL_FRONT) \
				* steering_bias_seed * 0.05

		_interest[i] = max(0, d + bias)


func _compute_danger() -> void:
	_danger.resize(ray_count)
	match danger_method:
		DangerComputationMethod.RAY:
			_compute_danger_ray_cast()
		DangerComputationMethod.MULTIRAY:
			_compute_danger_multi_ray()
		DangerComputationMethod.SHAPE:
			_compute_danger_shape_cast()


func _compute_danger_ray_cast() -> void:
	var origin := body.global_position + ray_origin_offset
	# Cast rays to find danger directions
	var space_state := get_world_3d().direct_space_state
	for i in ray_count:
		var target := origin + _ray_directions[i] * look_ahead
		var query := PhysicsRayQueryParameters3D.create(
			origin,
			target,
			Globals.collision_mask_units,
			[body],
		)
		var result := space_state.intersect_ray(query)
		if result.is_empty():
			_danger[i] = 0.0
		else:
			@warning_ignore("unsafe_call_argument")
			var hit_dist := origin.distance_to(result.position)
			_danger[i] = 1.0 - clampf(hit_dist / look_ahead, 0.0, 1.0)
			#_danger[i] = 1.0


# Doesn't seem to work that well
func _compute_danger_multi_ray() -> void:
	var space_state := get_world_3d().direct_space_state
	var origin := body.global_position + Vector3.UP

	for i in ray_count:
		var dir := _ray_directions[i]
		var side := dir.cross(Vector3.UP).normalized() * 0.4
		var max_danger := 0.0

		for offset: Vector3 in [Vector3.ZERO, side, -side]: # center, left edge, right edge
			var o := origin + ray_origin_offset
			var query := PhysicsRayQueryParameters3D.create(
				o,
				o + dir * look_ahead,
				Globals.collision_mask_units,
				[body],
			)
			var result := space_state.intersect_ray(query)
			if not result.is_empty():
				@warning_ignore("unsafe_call_argument")
				var hit_dist := o.distance_to(result.position)
				max_danger = max(max_danger, 1.0 - clampf(hit_dist / look_ahead, 0.0, 1.0))
				#max_danger = 1.0
		_danger[i] = max_danger


func _compute_danger_shape_cast() -> void:
	var origin := body.global_position + ray_origin_offset

	for i in ray_count:
		var result := _shape_cast_body(origin, _ray_directions[i])
		# result[0] = safe fraction of motion before collision, result[1] = fraction at contact
		var hit_fraction: float = result[0] if result.size() > 0 else 1.0
		_danger[i] = 1.0 - hit_fraction


func _choose_direction(original_direction: Vector3) -> Vector3:
	# Eliminate interest in slots with danger
	for i in ray_count:
		if _danger[i] > 0.0:
			_interest[i] = 0.0
		#_interest[i] = max(0.0, _interest[i] - _danger[i])
		else:
			# small continuity bonus for directions close to what we picked last frame
			_interest[i] += max(0.0, _ray_directions[i].dot(_last_chosen_direction)) * 0.15

	# Choose direction based on remaining interest
	var chosen_dir := Vector3.ZERO
	for i in ray_count:
		chosen_dir += _ray_directions[i] * _interest[i]

	if chosen_dir.is_zero_approx():
		return _fallback_direction(original_direction)

	_last_chosen_direction = chosen_dir
	return chosen_dir.normalized()


## Check if the chosen direction is actually valid based on the actual shape
## of the agent/body. Cheaper than converting rays
func _validate_chosen_direction(chosen_dir: Vector3) -> Vector3:
	@warning_ignore("unsafe_call_argument")
	var result := _shape_cast_body(
		body.global_position + ray_origin_offset,
		chosen_dir,
		min(look_ahead, 1.0), # short-range
	)
	if result.size() > 0 and result[0] < 0.5:
		# chosen direction is actually blocked for our body => fallback
		return _fallback_direction(chosen_dir)
	return chosen_dir


func _shape_cast_body(
		origin: Vector3,
		direction: Vector3,
		look_ahead_: float = look_ahead,
) -> PackedFloat32Array:
	var space_state := get_world_3d().direct_space_state
	var shape: Shape3D
	if collision_shape:
		shape = collision_shape
	else:
		var capsule_shape := CapsuleShape3D.new()
		capsule_shape.radius = 0.4
		capsule_shape.height = 1.8
		shape = capsule_shape

	var query := PhysicsShapeQueryParameters3D.new()
	query.shape = shape
	query.transform = Transform3D(Basis(), origin)
	query.motion = direction * look_ahead_
	query.collision_mask = Globals.collision_mask_units
	query.exclude = [body]
	return space_state.cast_motion(query)


func _fallback_direction(_original_direction: Vector3) -> Vector3:
	#return _last_chosen_direction

	#var left_right := randi() % 2 # 0 or 1
	#var base_angle := left_right * PI + PI / 2 # -PI/2 for right PI/2 for left
	#var angle := base_angle + randf_range(deg_to_rad(-10), deg_to_rad(10))
	#return original_direction.rotated(Vector3.UP, angle)

	# find the ray with lowest danger; if tied, prefer the one closest to last direction
	var best_i := -1
	var best_danger := INF
	for i in ray_count:
		var d: float = _danger[i]
		if d < best_danger:
			best_danger = d
			best_i = i
	return _ray_directions[best_i] if best_i != -1 else -_last_chosen_direction


func _get_effective_speed() -> float:
	var total_multiplier := 1.0
	var total_flat := 0.0
	for mod in _modifiers:
		total_multiplier *= mod.multiplier
		total_flat += mod.flat_bonus
	return movement_speed * total_multiplier + total_flat