extends SceneTree

## Headless script: loads animation-0.glb, extracts animations with
## original Blender bone names, applies rest-pose correction per bone,
## and saves animation-pack.res.
##
## Usage: godot --headless --path <project> --script tools/build_animation_pack.gd

const SOURCE_GLB := "res://assets/characters/animation-0.glb"
const TARGET_SKELETON_GLB := "res://assets/characters/Bob.glb"
const OUTPUT_PATH := "res://assets/characters/animations/animation-pack.res"

var _corrections: Dictionary = {}  # bone_name -> Quaternion correction
var _correction_src_rest: Dictionary = {}  # "Skeleton3D:bone" -> Vector3
var _correction_tgt_rest: Dictionary = {}  # "Skeleton3D:bone" -> Vector3

func _init() -> void:
	print("[build_animation_pack] Starting...")

	# Step 1: Load source skeleton (animation-0.glb)
	var gltf_doc := GLTFDocument.new()
	var gltf_state := GLTFState.new()
	var err := gltf_doc.append_from_file(
		ProjectSettings.globalize_path(SOURCE_GLB), gltf_state
	)
	if err != OK:
		push_error("[build_animation_pack] Failed to load source GLB: error %d" % err)
		quit(1)
		return

	var src_scene := gltf_doc.generate_scene(gltf_state)
	if not src_scene:
		push_error("[build_animation_pack] source generate_scene returned null")
		quit(1)
		return

	var src_skel: Skeleton3D = src_scene.find_child("Skeleton3D", true, false)
	if not src_skel:
		push_error("[build_animation_pack] No Skeleton3D in source")
		src_scene.queue_free()
		quit(1)
		return

	# Step 2: Load target skeleton (Bob, raw)
	var tgt_gltf_doc := GLTFDocument.new()
	var tgt_gltf_state := GLTFState.new()
	err = tgt_gltf_doc.append_from_file(
		ProjectSettings.globalize_path(TARGET_SKELETON_GLB), tgt_gltf_state
	)
	if err != OK:
		push_error("[build_animation_pack] Failed to load target GLB: error %d" % err)
		src_scene.queue_free()
		quit(1)
		return
	var tgt_scene: Node = tgt_gltf_doc.generate_scene(tgt_gltf_state)
	if not tgt_scene:
		push_error("[build_animation_pack] target generate_scene returned null")
		src_scene.queue_free()
		quit(1)
		return
	var tgt_skel: Skeleton3D = tgt_scene.find_child("Skeleton3D", true, false)
	if not tgt_skel:
		push_error("[build_animation_pack] No Skeleton3D in target")
		src_scene.queue_free()
		tgt_scene.queue_free()
		quit(1)
		return

	# Step 3: Build per-bone correction quaternions (original bone names, 1:1)
	var src_bone_index: Dictionary = {}
	for i in src_skel.get_bone_count():
		src_bone_index[src_skel.get_bone_name(i)] = i

	for i in tgt_skel.get_bone_count():
		var bone_name: String = tgt_skel.get_bone_name(i)
		if not src_bone_index.has(bone_name):
			continue
		var src_idx: int = src_bone_index[bone_name]
		var src_rest: Transform3D = src_skel.get_bone_rest(src_idx)
		var tgt_rest: Transform3D = tgt_skel.get_bone_rest(i)
		# correction = tgt_rest.rot * src_rest.rot.inverse()
		var correction: Quaternion = Quaternion(tgt_rest.basis) * Quaternion(src_rest.basis.inverse())
		_corrections[bone_name] = correction
		_correction_src_rest["Character/Skeleton3D:%s" % bone_name] = src_rest.origin
		_correction_tgt_rest["Character/Skeleton3D:%s" % bone_name] = tgt_rest.origin

	print("[build_animation_pack] Built %d rest-pose corrections" % _corrections.size())

	# Step 4: Extract animations with original bone names
	# We use the generated scene's AnimationPlayer, which contains the
	# already-stripped tracks (immutable ones removed by GLTF importer).
	var anim_player: AnimationPlayer = src_scene.find_child("AnimationPlayer", true, false)
	if not anim_player:
		push_error("[build_animation_pack] No AnimationPlayer in source")
		src_scene.queue_free()
		tgt_scene.queue_free()
		quit(1)
		return

	var out_lib := AnimationLibrary.new()
	var count := 0
	var total_tracks := 0
	var corrected_keys := 0

	for lib_name in anim_player.get_animation_library_list():
		var src_lib := anim_player.get_animation_library(lib_name)
		for anim_name in src_lib.get_animation_list():
			var src_anim: Animation = src_lib.get_animation(anim_name)
			if not src_anim:
				continue

			var new_anim := Animation.new()
			new_anim.length = src_anim.length
			new_anim.loop_mode = src_anim.loop_mode

			for i in src_anim.get_track_count():
				var path_str: String = str(src_anim.track_get_path(i))
				var track_type: int = src_anim.track_get_type(i)

				# Parse "retarget/Skeleton3D:bone_name" -> bone_name
				var bone_name := ""
				if ":" in path_str:
					bone_name = path_str.split(":")[1]
				else:
					# Non-bone track: skip GLB root "retarget"
					if "retarget" in path_str:
						continue
					_copy_track(src_anim, i, new_anim, path_str, Quaternion.IDENTITY, track_type)
					continue

				# Skip the GLB scene root "retarget" — not a real bone.
				if bone_name == "retarget":
					continue

				# Use original bone name with character-rig-relative path.
				# All characters have rig root renamed to "Character" via GLB modification,
				# so the path "Character/Skeleton3D:bone" is consistent across all of them.
				var new_path := "Character/Skeleton3D:%s" % bone_name
				var correction: Quaternion = _corrections.get(bone_name, Quaternion())
				_copy_track(src_anim, i, new_anim, new_path, correction, track_type)
				total_tracks += 1
				corrected_keys += _last_key_count

			out_lib.add_animation(anim_name, new_anim)
			count += 1
			print("  + %s (%d tracks)" % [anim_name, new_anim.get_track_count()])

	print("[build_animation_pack] Extracted %d animations, %d tracks, corrected %d rotation keys" % [count, total_tracks, corrected_keys])

	err = ResourceSaver.save(out_lib, OUTPUT_PATH)
	if err != OK:
		push_error("[build_animation_pack] Failed to save: error %d" % err)
		src_scene.queue_free()
		tgt_scene.queue_free()
		quit(1)
		return

	print("[build_animation_pack] Saved to %s" % OUTPUT_PATH)
	src_scene.queue_free()
	tgt_scene.queue_free()
	quit(0)

var _last_key_count: int = 0

func _copy_track(src_anim: Animation, src_idx: int, dst_anim: Animation, new_path: String, correction: Quaternion, track_type: int = -1) -> void:
	if track_type < 0:
		track_type = src_anim.track_get_type(src_idx)
	var dst_idx := dst_anim.add_track(track_type)
	dst_anim.track_set_path(dst_idx, NodePath(new_path))
	dst_anim.track_set_interpolation_type(dst_idx, src_anim.track_get_interpolation_type(src_idx))

	var key_count := src_anim.track_get_key_count(src_idx)
	_last_key_count = 0
	for k in key_count:
		var time := src_anim.track_get_key_time(src_idx, k)
		var value = src_anim.track_get_key_value(src_idx, k)
		var transition := src_anim.track_get_key_transition(src_idx, k)

		match track_type:
			Animation.TYPE_POSITION_3D:
				var src_pos: Vector3 = value
				var src_rest_pos: Vector3 = _correction_src_rest.get(new_path, Vector3.ZERO)
				var tgt_rest_pos: Vector3 = _correction_tgt_rest.get(new_path, Vector3.ZERO)
				var corrected_pos: Vector3 = correction * (src_pos - src_rest_pos) + tgt_rest_pos
				dst_anim.position_track_insert_key(dst_idx, time, corrected_pos)
			Animation.TYPE_ROTATION_3D:
				var src_q: Quaternion = value
				var corrected_q: Quaternion = correction * src_q
				dst_anim.rotation_track_insert_key(dst_idx, time, corrected_q)
				_last_key_count += 1
			Animation.TYPE_SCALE_3D:
				dst_anim.scale_track_insert_key(dst_idx, time, value)
			Animation.TYPE_BLEND_SHAPE:
				dst_anim.blend_shape_track_insert_key(dst_idx, time, value)
			_:
				dst_anim.track_insert_key(dst_idx, time, value, transition)