extends SceneTree

# Headless: convert each dasher_*.glb into a dasher_<name>.res with proper retargeting.
# Reads rest poses from both source GLB skeleton and target GeneralSkeleton (player.tscn)
# to compute per-bone correction quaternions.


const DASHER_DIR := "res://assets/characters/dashers"
const ANIM_DIR := "res://assets/characters/animations"

const BONE_REMAP := {
	"bob-hold": "Hips",
	"spine": "Hips",
	"spine.001": "Spine",
	"head": "Head",
	"hand.L": "LeftHand",
	"hand.R": "RightHand",
	"forearm.L": "LeftLowerArm",
	"forearm.R": "RightLowerArm",
	"upper_arm.L": "LeftUpperArm",
	"upper_arm.R": "RightUpperArm",
	"shoulder.L": "LeftShoulder",
	"shoulder.R": "RightShoulder",
	"leg.L": "LeftLowerLeg",
	"leg.R": "RightLowerLeg",
	"thigh.L": "LeftUpperLeg",
	"thigh.R": "RightUpperLeg",
	"foot.L": "LeftFoot",
	"foot.R": "RightFoot",
	"toe.L": "LeftToeBase",
	"toe.R": "RightToeBase",
}

const ANIM_PICK := {
	"dasher_getting_hit": "Getting Hit",
	"dasher_hit": "Hit",
	"dasher_hold": "Hold",
	"dasher_put": "Put",
	"dasher_stun": "Stun",
	"dasher_take": "bob ani",
}

# Source bone rest transforms (from GLB)
var src_rest: Dictionary = {} # bone_name -> Transform3D
# Target bone rest transforms (from GeneralSkeleton)
var tgt_rest: Dictionary = {} # bone_name -> Transform3D
# Precomputed correction quaternions
var corrections: Dictionary = {} # bone_name -> Quaternion

func _init() -> void:
	print("[Convert] Headless run starting...")
	
	# Step 1: Load target skeleton rest poses from player scene
	_load_target_rest()
	
	# Step 2: Convert each GLB
	var dasher_dir := DirAccess.open(DASHER_DIR)
	if not dasher_dir:
		push_error("[Convert] Cannot open %s" % DASHER_DIR)
		quit(1); return

	var converted: Array[String] = []
	for fname in dasher_dir.get_files():
		if not fname.ends_with(".glb"): continue
		if not fname.begins_with("dasher_"): continue
		var stem := fname.get_basename()
		var out_path := "%s/%s.res" % [ANIM_DIR, stem]
		var ok := _convert_one("%s/%s" % [DASHER_DIR, fname], stem, out_path)
		if ok:
			converted.append(out_path)
		else:
			push_warning("[Convert] Skipped %s" % fname)

	print("[Convert] Done. %d dasher glb(s) converted:" % converted.size())
	for p in converted:
		print("  - %s" % p)
	quit(0)

func _load_target_rest() -> void:
	# Load Masbro.glb directly to extract GeneralSkeleton rest poses
	# (player.tscn may fail to load if dasher-pack.res doesn't exist yet)
	var doc := GLTFDocument.new()
	var state := GLTFState.new()
	var err := doc.append_from_file("res://assets/characters/Masbro.glb", state)
	if err != OK:
		push_error("[Convert] Cannot load Masbro.glb: %d" % err)
		return
	
	var instance = doc.generate_scene(state)
	if not instance:
		push_error("[Convert] generate_scene returned null for Masbro.glb")
		return
	
	# Find the Skeleton3D (named GeneralSkeleton in the scene tree)
	var skeleton: Skeleton3D = instance.find_child("GeneralSkeleton", true, false)
	if not skeleton:
		# Try finding any Skeleton3D
		skeleton = instance.find_child("Skeleton3D", true, false)
	if not skeleton:
		push_error("[Convert] No Skeleton3D in Masbro.glb")
		instance.free()
		return
	
	print("[Convert] Target skeleton: '%s' with %d bones" % [skeleton.name, skeleton.get_bone_count()])
	for bone_idx in skeleton.get_bone_count():
		var bone_name := skeleton.get_bone_name(bone_idx)
		tgt_rest[bone_name] = skeleton.get_bone_rest(bone_idx)
	
	instance.free()
	print("[Convert] Loaded %d target bone rest poses" % tgt_rest.size())

func _convert_one(glb_path: String, anim_name: String, out_path: String) -> bool:
	print("[Convert] %s -> %s" % [glb_path, out_path])
	var doc := GLTFDocument.new()
	var state := GLTFState.new()
	var err := doc.append_from_file(glb_path, state)
	if err != OK:
		push_error("[Convert] append_from_file failed: %d" % err)
		return false

	var scene: Node = doc.generate_scene(state)
	if not scene:
		push_error("[Convert] generate_scene returned null")
		return false

	# Extract source skeleton rest poses
	src_rest.clear()
	var src_skeleton: Skeleton3D = scene.find_child("Skeleton3D", true, false)
	if src_skeleton:
		for bone_idx in src_skeleton.get_bone_count():
			var bone_name := src_skeleton.get_bone_name(bone_idx)
			src_rest[bone_name] = src_skeleton.get_bone_rest(bone_idx)
		print("  Source skeleton: %d bones" % src_rest.size())
	else:
		push_error("[Convert] No Skeleton3D in %s" % glb_path)
		scene.queue_free()
		return false

	# Compute per-bone correction quaternions
	corrections.clear()
	for src_bone in BONE_REMAP:
		var tgt_bone: String = BONE_REMAP[src_bone]
		if src_rest.has(src_bone) and tgt_rest.has(tgt_bone):
			var src_basis: Basis = src_rest[src_bone].basis
			var tgt_basis: Basis = tgt_rest[tgt_bone].basis
			# correction = tgt_rest.inverse() * src_rest
			# Applied as: final_q = correction * keyframe_q
			corrections[src_bone] = tgt_basis.inverse() * src_basis
			# print("  Correction for %s -> %s: %s" % [src_bone, tgt_bone, str(corrections[src_bone])])
	
	var src_player: AnimationPlayer = scene.find_child("AnimationPlayer", true, false)
	if not src_player:
		push_error("[Convert] No AnimationPlayer in %s" % glb_path)
		scene.queue_free()
		return false

	# Pick the animation whose middle segment matches the keyword
	var pick_keyword: String = ANIM_PICK.get(anim_name, anim_name.replace("dasher_", ""))
	var picked: Animation = null
	for lib_name in src_player.get_animation_library_list():
		var src_lib: AnimationLibrary = src_player.get_animation_library(lib_name)
		for src_anim_name in src_lib.get_animation_list():
			var segments := src_anim_name.split("|")
			var mid_segment: String = segments[1] if segments.size() >= 3 else src_anim_name
			if pick_keyword.to_lower() == mid_segment.to_lower():
				picked = src_lib.get_animation(src_anim_name)
				break
		if picked: break
	if not picked:
		for lib_name in src_player.get_animation_library_list():
			var src_lib: AnimationLibrary = src_player.get_animation_library(lib_name)
			if src_lib.get_animation_list().size() > 0:
				picked = src_lib.get_animation(src_lib.get_animation_list()[0])
				break
	if not picked:
		push_error("[Convert] No animations in %s" % glb_path)
		scene.queue_free()
		return false

	var retargeted := _retarget(picked)
	scene.queue_free()

	if retargeted.get_track_count() == 0:
		push_error("[Convert] Retarget produced 0 tracks for %s" % glb_path)
		return false

	var out_lib := AnimationLibrary.new()
	out_lib.add_animation(anim_name, retargeted)
	print("  + %s (%d tracks from %s)" % [anim_name, retargeted.get_track_count(), picked.resource_name if picked.resource_name else "<anon>"])

	var save_err := ResourceSaver.save(out_lib, out_path)
	if save_err != OK:
		push_error("[Convert] ResourceSaver.save failed: %d" % save_err)
		return false
	return true

func _retarget(src: Animation) -> Animation:
	var dst := Animation.new()
	dst.length = src.length
	dst.loop_mode = src.loop_mode
	dst.step = src.step

	var usable: Array[int] = []
	for i in src.get_track_count():
		var orig_path: String = src.track_get_path(i)
		var bone := _extract_bone_name(orig_path)
		if BONE_REMAP.has(bone):
			usable.append(i)

	for i in usable:
		var orig_path: String = src.track_get_path(i)
		var bone: String = _extract_bone_name(orig_path)
		var new_bone: String = BONE_REMAP[bone]
		var new_path := "%%GeneralSkeleton:%s" % new_bone

		var track_type := src.track_get_type(i)

		match track_type:
			Animation.TYPE_POSITION_3D:
				if new_bone != "Hips": continue
				dst.add_track(Animation.TYPE_POSITION_3D)
			Animation.TYPE_ROTATION_3D:
				dst.add_track(Animation.TYPE_ROTATION_3D)
			Animation.TYPE_SCALE_3D:
				if new_bone != "Hips": continue
				dst.add_track(Animation.TYPE_SCALE_3D)
			Animation.TYPE_BLEND_SHAPE: continue
			_: continue

		var new_idx := dst.get_track_count() - 1
		dst.track_set_path(new_idx, NodePath(new_path))
		dst.track_set_interpolation_type(new_idx, src.track_get_interpolation_type(i))
		dst.track_set_imported(new_idx, src.track_is_imported(i))
		dst.track_set_enabled(new_idx, src.track_is_enabled(i))

		var key_count := src.track_get_key_count(i)
		for k in key_count:
			var t: float = src.track_get_key_time(i, k)
			var v: Variant = src.track_get_key_value(i, k)
			var trans := src.track_get_key_transition(i, k)
			
			# Apply rest-pose correction to rotation tracks
			if track_type == Animation.TYPE_ROTATION_3D and corrections.has(bone):
				var orig_q: Quaternion = v
				var corr_q: Quaternion = corrections[bone]
				v = corr_q * orig_q
			
			dst.track_insert_key(new_idx, t, v)
			dst.track_set_key_transition(new_idx, k, trans)
	return dst

func _extract_bone_name(path: String) -> String:
	var parts := path.split(":")
	if parts.size() < 2:
		return ""
	return parts[-1]