extends SceneTree

# Reproducible offline bake: run after --editor --import. Runtime never opens FBX.
const BASE := "res://assets/animations/"
const MODELS := {"steve": "Steve", "gas_mask": "Gas_Mask", "pro_rifle": "Ch49_nonPBR", "basic_shooter": "Ch49_nonPBR", "slim_shooter": "Ch49_nonPBR"}
var manifest: Array[Dictionary] = []

func _initialize() -> void:
	call_deferred("bake")

func bake() -> void:
	DirAccess.make_dir_recursive_absolute("res://scenes/operators")
	var library := AnimationLibrary.new()
	for package: String in ["pro_rifle", "basic_shooter", "slim_shooter", "slide", "steve", "gas_mask"]:
		var files := DirAccess.get_files_at(BASE + package)
		files.sort()
		for file: String in files:
			if file.get_extension().to_lower() != "fbx": continue
			var path := BASE + package + "/" + file
			var source := (load(path) as PackedScene).instantiate()
			var skeleton := source.find_child("Skeleton3D", true, false) as Skeleton3D
			assert(skeleton != null, "Missing skeleton: " + path)
			for child in source.find_children("*", "AnimationPlayer", true, false):
				var player := child as AnimationPlayer
				for clip: String in player.get_animation_list():
					if clip == "RESET": continue
					var key := package + "__" + file.get_basename().to_lower().replace(" ", "_").replace("(", "").replace(")", "")
					if MODELS.get(package, "") == file.get_basename(): key += "__" + clip.to_lower().replace(" ", "_")
					var animation := player.get_animation(clip).duplicate(true) as Animation
					animation.resource_name = key
					# All exported clips use the same named Mixamo hierarchy. Keep their
					# native poses, bind data and source translation for retargeting.
					var rests := {}
					for bone in skeleton.get_bone_count():
						rests[str(skeleton.get_bone_name(bone))] = skeleton.get_bone_rest(bone)
					animation.set_meta("source_rests", rests)
					for track in animation.get_track_count():
						var original := animation.track_get_path(track)
						if original.get_subname_count() > 0:
							animation.track_set_path(track, NodePath("Skeleton3D:" + str(original.get_subname(0))))
					library.add_animation(key, animation)
					manifest.append({"key": key, "source": path, "clip": clip, "seconds": animation.length, "tracks": animation.get_track_count()})
			if MODELS.get(package, "") == file.get_basename():
				bake_model(source, package)
			source.free()
	assert(ResourceSaver.save(library, BASE + "operators.res") == OK)
	var report := FileAccess.open(BASE + "manifest.json", FileAccess.WRITE)
	report.store_string(JSON.stringify(manifest, "\t"))
	print("OPERATOR BAKE: ", manifest.size(), " animations, ", MODELS.size(), " models")
	quit()

func bake_model(source: Node, package: String) -> void:
	var material_dir := BASE + package + "/materials/"
	DirAccess.make_dir_recursive_absolute(material_dir)
	var materials := {}
	for child in source.find_children("*", "MeshInstance3D", true, false):
		var mesh := child as MeshInstance3D
		for surface in mesh.mesh.get_surface_count():
			var original := mesh.get_active_material(surface) as StandardMaterial3D
			if original == null: continue
			var key := original.resource_name.validate_filename()
			if not materials.has(key):
				var mat := original.duplicate() as StandardMaterial3D
				# FBX diffuse-factor alpha is not opacity (these exports use 0.8).
				mat.albedo_color = Color.WHITE
				mat.metallic = 0.0
				mat.roughness = 0.82
				if mat.albedo_texture != null:
					var roughness_path := mat.albedo_texture.resource_path.replace("_Diffuse.png", "_Roughness.png")
					if roughness_path != mat.albedo_texture.resource_path and ResourceLoader.exists(roughness_path):
						mat.roughness_texture = load(roughness_path) as Texture2D
						mat.roughness_texture_channel = BaseMaterial3D.TEXTURE_CHANNEL_RED
						mat.roughness = 1.0
				mat.normal_enabled = mat.normal_texture != null
				mat.normal_scale = 1.0
				mat.emission_enabled = false
				mat.transparency = BaseMaterial3D.TRANSPARENCY_ALPHA_SCISSOR if "eyelash" in key.to_lower() else BaseMaterial3D.TRANSPARENCY_DISABLED
				mat.alpha_scissor_threshold = 0.45
				var path := material_dir + key + ".tres"
				assert(ResourceSaver.save(mat, path) == OK)
				mat.take_over_path(path)
				materials[key] = mat
			mesh.set_surface_override_material(surface, materials[key])
		mesh.extra_cull_margin = 1.5
	for child in source.find_children("*", "AnimationPlayer", true, false):
		child.get_parent().remove_child(child)
		child.free()
	assign_owners(source, source)
	var packed := PackedScene.new()
	assert(packed.pack(source) == OK)
	assert(ResourceSaver.save(packed, "res://scenes/operators/" + package + ".scn") == OK)
	print("BAKED MODEL ", package, " materials ", materials.size())
	if package in ["steve", "gas_mask"]: bake_arms(source, package)

func bake_arms(source: Node, package: String) -> void:
	var skeleton := source.find_child("Skeleton3D", true, false) as Skeleton3D
	for child in source.find_children("*", "MeshInstance3D", true, false):
		var instance := child as MeshInstance3D
		var mesh := ArrayMesh.new()
		var materials: Array[Material] = []
		for surface in instance.mesh.get_surface_count():
			var arrays := instance.mesh.surface_get_arrays(surface)
			var vertices: PackedVector3Array = arrays[Mesh.ARRAY_VERTEX]
			var bones: PackedInt32Array = arrays[Mesh.ARRAY_BONES]
			var weights: PackedFloat32Array = arrays[Mesh.ARRAY_WEIGHTS]
			var source_indices: PackedInt32Array = arrays[Mesh.ARRAY_INDEX]
			var indices := PackedInt32Array()
			var keep := PackedByteArray()
			var influences := bones.size() / vertices.size()
			for vertex in vertices.size():
				var arm_weight := 0.0
				for influence in int(influences):
					var offset := vertex * int(influences) + influence
					var bind := bones[offset]
					var name := str(instance.skin.get_bind_name(bind))
					if name.is_empty(): name = str(skeleton.get_bone_name(instance.skin.get_bind_bone(bind)))
					if "Arm" in name or "Hand" in name: arm_weight += weights[offset]
				keep.append(1 if arm_weight > 0.7 else 0)
			for triangle in range(0, source_indices.size(), 3):
				if keep[source_indices[triangle]] and keep[source_indices[triangle + 1]] and keep[source_indices[triangle + 2]]:
					indices.append_array(source_indices.slice(triangle, triangle + 3))
			if indices.is_empty(): continue
			arrays[Mesh.ARRAY_INDEX] = indices
			mesh.add_surface_from_arrays(Mesh.PRIMITIVE_TRIANGLES, arrays)
			materials.append(instance.get_active_material(surface))
		if mesh.get_surface_count() == 0:
			instance.get_parent().remove_child(instance)
			instance.free()
		else:
			instance.mesh = mesh
			for surface in materials.size(): instance.set_surface_override_material(surface, materials[surface])
			instance.cast_shadow = GeometryInstance3D.SHADOW_CASTING_SETTING_OFF
	var packed := PackedScene.new()
	assert(packed.pack(source) == OK)
	assert(ResourceSaver.save(packed, "res://scenes/operators/" + package + "_arms.scn") == OK)
	print("BAKED FPS ARMS ", package)

func assign_owners(node: Node, scene: Node) -> void:
	for child in node.get_children():
		child.owner = scene
		assign_owners(child, scene)