@tool
class_name BoneFlattener
extends SkeletonModifier3D

@export var _bones_to_flatten: Array[BoneToFlatten] = []
@export var _mirror_y_angle: float
@export var _editor_preview: bool = false

@export_group("Bone names")
@export_enum(" ") var _head_bone: String = "Head"
@export_enum(" ") var _mouth_bone: String = "Mouth_base"

@export_group("Mouth")
@export var _mouth_center_pos_z: float:
	set(value):
		update_gizmos()
		_mouth_center_pos_z = value

@export var _mouth_corner_pos_z: float:
	set(value):
		update_gizmos()
		_mouth_corner_pos_z = value

@export var _mouth_pos_z_curve: Curve:
	set(value):
		update_gizmos()
		_mouth_pos_z_curve = value

@export var _mouth_corner_pos_x: float:
	set(value):
		update_gizmos()
		_mouth_corner_pos_x = value

@export var _mouth_center_pos_y: float:
	set(value):
		update_gizmos()
		_mouth_center_pos_y = value

@export var _mouth_corner_pos_y: float:
	set(value):
		update_gizmos()
		_mouth_corner_pos_y = value

@export var _mouth_corner_rot_y: float:
	set(value):
		update_gizmos()
		_mouth_corner_rot_y = value

@export var _mouth_rot_y_curve: Curve:
	set(value):
		update_gizmos()
		_mouth_rot_y_curve = value

@export var _mouth_corner_rot_x: float:
	set(value):
		update_gizmos()
		_mouth_corner_rot_x = value

@export var _mouth_corner_rot_z: float:
	set(value):
		update_gizmos()
		_mouth_corner_rot_z = value

@export var _mouth_value_yaw_curve: Curve:
	set(value):
		update_gizmos()
		_mouth_value_yaw_curve = value

@export var _mouth_front_pitch_curve: Curve
@export var _mouth_front_yaw_curve: Curve
@export var _mouth_hide_rot_y: float

var _skeleton: Skeleton3D

var _angle: Vector3 = Vector3()


func _validate_property(property: Dictionary) -> void:
	if property.name.ends_with("bone"):
		if _skeleton:
			property.hint = PROPERTY_HINT_ENUM
			property.hint_string = _skeleton.get_concatenated_bone_names()


func _ready() -> void:
	_skeleton = get_skeleton()
	assert(_skeleton, str(self) + ": _skeleton missing!")
	assert(_skeleton.find_bone(_head_bone) != -1, str(self) + ": _head_bone missing!")
	assert(_skeleton.find_bone(_mouth_bone) != -1, str(self) + ": _mouth_bone missing!")

	for bone_to_flatten in _bones_to_flatten:
		if bone_to_flatten == null:
			continue
		for bone_name in bone_to_flatten.bone_names:
			var bone_idx := _skeleton.find_bone(bone_name)
			assert(bone_idx != -1, str(self) + ": bone " + bone_name + " missing!")


func _process_modification() -> void:
	if _skeleton == null or (Engine.is_editor_hint() and not _editor_preview):
		return

	_get_angle()
	_handle_mouth()

	for i in range(_bones_to_flatten.size()):
		var bone_to_flatten := _bones_to_flatten[i]
		if bone_to_flatten == null:
			continue
		for bone_name in bone_to_flatten.bone_names:
			var bone_idx: int = _skeleton.find_bone(bone_name)
			if bone_idx == -1:
				continue
			_handle_position(bone_to_flatten, bone_idx)
			_handle_rotation(bone_to_flatten, bone_idx)
			_handle_scale(bone_to_flatten, bone_idx)


func mouth_pose(normalized: Vector3) -> Vector3:
	var pos: Vector3 = Vector3.ZERO
	pos.x = normalized.y * _mouth_corner_pos_x
	var depth_value := _mouth_pos_z_curve.sample(absf(normalized.y))
	pos.z = remap(depth_value, 0, 1, _mouth_center_pos_z, _mouth_corner_pos_z)
	pos.y = remap(depth_value, 0, 1, _mouth_center_pos_y, _mouth_corner_pos_y)

	return pos


func global_bone_transform(bone_idx: int) -> Transform3D:
	var transform_local := _skeleton.get_bone_global_pose(bone_idx)
	return _skeleton.global_transform * transform_local


func _get_angle() -> void:
	var camera: Camera3D
	if Engine.is_editor_hint():
		var editor_interface := Engine.get_singleton("EditorInterface")
		camera = editor_interface.get_editor_viewport_3d().get_camera_3d()
	else:
		camera = get_tree().root.get_viewport().get_camera_3d()

	var head_bone_idx := _skeleton.find_bone(_head_bone)
	var head_transform := global_bone_transform(head_bone_idx)
	var head_transform_rotated := (
		head_transform
		. looking_at(
			camera.global_position,
			Vector3.UP,
			true,
		)
	)
	var head_rot := head_transform.basis.get_rotation_quaternion()
	var head_rotated_rot := head_transform_rotated.basis.get_rotation_quaternion()
	var diff := (head_rot.inverse() * head_rotated_rot).normalized()
	_angle = diff.get_euler()


func _handle_mouth() -> void:
	if (
		_mouth_value_yaw_curve == null
		or _mouth_rot_y_curve == null
		or _mouth_pos_z_curve == null
		or _mouth_front_pitch_curve == null
		or _mouth_front_yaw_curve == null
	):
		return

	var mouth_bone_idx := _skeleton.find_bone(_mouth_bone)
	var bone_transform := _skeleton.get_bone_rest(mouth_bone_idx)
	var bone_rot := bone_transform.basis.get_euler()
	var bone_pos := bone_transform.origin
	var normalized := Vector3.ZERO
	normalized.y = clampf(_angle.y / (PI / 2), -1, 1)
	normalized.x = clampf(_angle.x / (PI / 2), -1, 1)
	normalized.y = (
		_mouth_value_yaw_curve.sample(absf(normalized.y)) * signf(normalized.y)
	)

	bone_pos += mouth_pose(normalized)
	_skeleton.set_bone_pose_position(mouth_bone_idx, bone_pos)

	var corner_angle_value := ease(absf(normalized.y), 1) * signf(normalized.y)
	bone_rot.x += (
		(
			_mouth_rot_y_curve
			. sample(
				absf(corner_angle_value),
			)
		)
		* _mouth_corner_rot_x
	)
	var rot_y := (
		(
			_mouth_rot_y_curve
			. sample(
				absf(corner_angle_value),
			)
		)
		* signf(normalized.y)
		* _mouth_corner_rot_y
	)
	bone_rot = (
		(
			Quaternion(
				Vector3.FORWARD,
				_mouth_corner_rot_z * normalized.y,
			)
			* Quaternion(
				Vector3.UP,
				rot_y,
			)
			* Quaternion.from_euler(bone_rot)
		)
		. get_euler()
	)

	_skeleton.set_bone_pose_rotation(mouth_bone_idx, Quaternion.from_euler(bone_rot))

	var bone_scale := Vector3.ONE
	var scale_x_front := (
		_mouth_front_pitch_curve
		. sample(
			inverse_lerp(-1, 1, normalized.x),
		)
	)
	bone_scale.x = lerpf(
		1,
		scale_x_front,
		_mouth_front_yaw_curve.sample(absf(normalized.y)),
	)
	if _angle.y > _mirror_y_angle:
		bone_scale.x *= -1
	if absf(_angle.y) > _mouth_hide_rot_y:
		bone_scale = Vector3.ZERO
	_skeleton.set_bone_pose_scale(mouth_bone_idx, bone_scale)


func _handle_position(bone_to_flatten: BoneToFlatten, bone: int) -> void:
	if not bone_to_flatten.do_position:
		return

	var bone_pos := _skeleton.get_bone_rest(bone).origin
	bone_pos.x += (
		bone_to_flatten.get_amount(_angle) * bone_to_flatten.position_x_amount
	)
	_skeleton.set_bone_pose_position(bone, bone_pos)


func _handle_rotation(bone_to_flatten: BoneToFlatten, bone: int) -> void:
	if not bone_to_flatten.do_rotation:
		return

	var bone_transform := _skeleton.get_bone_rest(bone)
	var bone_rot := bone_transform.basis.get_euler()
	var bone_pos := bone_transform.origin
	var amount := bone_to_flatten.get_amount(_angle)
	var normalized := bone_to_flatten.get_angle_normalized(_angle)

	bone_rot.x += amount * (PI / 2) * bone_to_flatten.rotation_x_amount
	if bone_to_flatten.mirror_rot_x and bone_pos.x > 0:
		bone_rot.x *= -1
	if bone_to_flatten.rotation_x_curve != null:
		bone_rot.x *= (bone_to_flatten.rotation_x_curve.sample(absf(normalized.y)))
	if bone_to_flatten.consider_side_rot_x:
		bone_rot.x *= -signf(normalized.y)

	bone_rot.y += amount * (PI / 2) * bone_to_flatten.rotation_y_amount
	if bone_to_flatten.mirror_rot_y and bone_pos.x > 0:
		bone_rot.y *= -1
	if bone_to_flatten.rotation_y_curve != null:
		bone_rot.y *= (bone_to_flatten.rotation_y_curve.sample(absf(normalized.y)))
	if bone_to_flatten.consider_side_rot_y:
		bone_rot.y *= -signf(normalized.y)

	bone_rot.z += amount * (PI / 2) * bone_to_flatten.rotation_z_amount
	if bone_to_flatten.mirror_rot_z and bone_pos.x > 0:
		bone_rot.z *= -1
	if bone_to_flatten.rotation_z_curve != null:
		bone_rot.z *= (bone_to_flatten.rotation_z_curve.sample(absf(normalized.y)))
	if bone_to_flatten.consider_side_rot_z:
		bone_rot.z *= -signf(normalized.y)

	_skeleton.set_bone_pose_rotation(bone, Quaternion.from_euler(bone_rot))


func _handle_scale(bone_to_flatten: BoneToFlatten, _bone: int) -> void:
	if not bone_to_flatten.do_scale:
		return