class_name OperatorPose
extends RefCounted

# Shared immutable clips, per-instance skeleton and blend state. CharacterBody3D
# remains authoritative: animation root motion must never move a network actor.
static var library: AnimationLibrary
var skeleton: Skeleton3D
var rotations: Array[Quaternion] = []
var rest_rotations: Array[Quaternion] = []
var rest_positions: Array[Vector3] = []
var upper_body: PackedByteArray = []
var track_cache: Dictionary = {}
var global_poses: Array[Transform3D] = []

func setup(rig: Skeleton3D) -> void:
	skeleton = rig
	if library == null:
		library = load("res://assets/animations/operators.res") as AnimationLibrary
	for bone in rig.get_bone_count():
		var rest := rig.get_bone_rest(bone)
		rest_positions.append(rest.origin)
		rest_rotations.append(rest.basis.get_rotation_quaternion())
		rotations.append(rest.basis.get_rotation_quaternion())
		global_poses.append(Transform3D.IDENTITY)
		var bone_name := str(rig.get_bone_name(bone))
		upper_body.append(0 if bone_name.contains("Hips") or bone_name.contains("Leg") or bone_name.contains("Foot") or bone_name.contains("Toe") else 1)

func clip(key: String) -> Animation:
	return library.get_animation(key) if library.has_animation(key) else library.get_animation("pro_rifle__idle")

func tracks(key: String) -> PackedInt32Array:
	if track_cache.has(key): return track_cache[key]
	var result := PackedInt32Array()
	var animation := clip(key)
	for track in animation.get_track_count():
		var path := animation.track_get_path(track)
		result.append(skeleton.find_bone(path.get_subname(0)) if path.get_subname_count() > 0 else -1)
	track_cache[key] = result
	return result

func sample(key: String, time: float, looping: bool, upper_only: bool, weight: float, positions: Array[Vector3]) -> void:
	var animation := clip(key)
	var stamp := fposmod(time, animation.length) if looping else clampf(time, 0, animation.length)
	var indices := tracks(key)
	for track in animation.get_track_count():
		var bone := indices[track]
		if bone < 0 or (upper_only and upper_body[bone] == 0): continue
		if animation.track_get_type(track) == Animation.TYPE_ROTATION_3D:
			# Mixamo clips encode local bone rotations in the same named axes.
			# Keep target bone lengths; never replace the target bind skeleton.
			var rotation := animation.rotation_track_interpolate(track, stamp).normalized()
			rotations[bone] = rotations[bone].slerp(rotation, weight)
		elif animation.track_get_type(track) == Animation.TYPE_POSITION_3D and not upper_only:
			var value := animation.position_track_interpolate(track, stamp)
			if str(skeleton.get_bone_name(bone)).ends_with("Hips"):
				var source_rests: Dictionary = animation.get_meta("source_rests", {})
				var source_rest: Transform3D = source_rests.get(str(skeleton.get_bone_name(bone)), skeleton.get_bone_rest(bone))
				value *= rest_positions[bone].y / maxf(0.01, source_rest.origin.y)
				value.x = rest_positions[bone].x
				value.z = rest_positions[bone].z
				if "jump" in key:
					value.y = minf(value.y, rest_positions[bone].y)
				positions[bone] = positions[bone].lerp(value, weight)

func apply(base: String, time: float, looping: bool, action: String, action_time: float, action_weight: float, delta: float, pitch: float) -> void:
	rotations.assign(rest_rotations)
	var positions: Array[Vector3] = rest_positions.duplicate()
	sample(base, time, looping, false, 1.0, positions)
	if not action.is_empty(): sample(action, action_time, false, true, action_weight, positions)
	var spine := skeleton.find_bone("mixamorig_Spine2")
	if spine >= 0:
		rotations[spine] = Quaternion(Vector3.RIGHT, -clampf(pitch, -0.8, 0.8)) * rotations[spine]
	var blend := 1.0 if delta <= 0 else 1.0 - exp(-delta * 22.0)
	for bone in skeleton.get_bone_count():
		skeleton.set_bone_pose_rotation(bone, skeleton.get_bone_pose_rotation(bone).slerp(rotations[bone], blend))
		skeleton.set_bone_pose_position(bone, skeleton.get_bone_pose_position(bone).lerp(positions[bone], blend))
	update_globals()

func update_globals() -> void:
	# Skeleton3D's global pose cache can still reflect the previous render frame.
	# Compose the freshly written local poses for same-frame sockets and IK.
	for bone in skeleton.get_bone_count():
		var parent := skeleton.get_bone_parent(bone)
		global_poses[bone] = skeleton.get_bone_pose(bone)
		if parent >= 0: global_poses[bone] = global_poses[parent] * global_poses[bone]

func move_shoulder(side: String, target: Vector3) -> void:
	var bone := skeleton.find_bone("mixamorig_" + side + "Arm")
	var parent := skeleton.get_bone_parent(bone)
	skeleton.set_bone_pose_position(bone, global_poses[parent].affine_inverse() * target)
	update_globals()

func reach_arm(side: String, target: Vector3) -> void:
	var arm := skeleton.find_bone("mixamorig_" + side + "Arm")
	var elbow := skeleton.find_bone("mixamorig_" + side + "ForeArm")
	var hand := skeleton.find_bone("mixamorig_" + side + "Hand")
	if arm < 0 or elbow < 0 or hand < 0: return
	var palm := global_poses[hand].basis.get_rotation_quaternion()
	var start := global_poses[arm].origin
	var joint := global_poses[elbow].origin
	var end := global_poses[hand].origin
	var upper_length := start.distance_to(joint)
	var lower_length := joint.distance_to(end)
	var axis := (target - start).normalized()
	var distance := clampf(start.distance_to(target), absf(upper_length - lower_length) + 0.001, upper_length + lower_length - 0.001)
	var along := (upper_length * upper_length + distance * distance - lower_length * lower_length) / (2 * distance)
	var bend := (joint - start) - axis * axis.dot(joint - start)
	if bend.length_squared() < 0.00001: bend = Vector3.DOWN - axis * axis.dot(Vector3.DOWN)
	var desired_elbow := start + axis * along + bend.normalized() * sqrt(maxf(0, upper_length * upper_length - along * along))
	rotate_towards(arm, joint - start, desired_elbow - start)
	rotate_towards(elbow, global_poses[hand].origin - global_poses[elbow].origin, start + axis * distance - global_poses[elbow].origin)
	var parent := skeleton.get_bone_parent(hand)
	skeleton.set_bone_pose_rotation(hand, global_poses[parent].basis.get_rotation_quaternion().inverse() * palm)
	update_globals()

func rotate_towards(bone: int, from: Vector3, to: Vector3) -> void:
	if from.length_squared() < 0.000001 or to.length_squared() < 0.000001: return
	var parent := skeleton.get_bone_parent(bone)
	var parent_rotation := global_poses[parent].basis.get_rotation_quaternion()
	var delta := Quaternion(from.normalized(), to.normalized())
	skeleton.set_bone_pose_rotation(bone, (parent_rotation.inverse() * delta * parent_rotation * skeleton.get_bone_pose_rotation(bone)).normalized())
	update_globals()