scripts/modules/rigging.gd

extends Node
##
## RiggingModule - armature creation and automatic vertex weighting (Features 41-60).
##
## Creates a Skeleton3D, lets the user add bones as a transform hierarchy, and
## binds skinning data (ARRAY_BONES + ARRAY_WEIGHTS) to a MeshInstance3D using
## inverse-distance weighting against bone segments. The result is a normal,
## engine-native skinned mesh - no custom shader needed.
##

## Bone definitions the tool builds. Each entry: { name, parent, position, tail }
var _bones: Array[Dictionary] = []
var _skeleton: Skeleton3D = null
var _host: Node = null
var _bound_mesh: MeshInstance3D = null


func _activate(host: Node) -> void:
	_host = host
	StudioBus.object_selected.connect(_on_object_selected)
	print("[Rigging] activated")


func _deactivate() -> void:
	if StudioBus.object_selected.is_connected(_on_object_selected):
		StudioBus.object_selected.disconnect(_on_object_selected)
	_host = null


func _on_object_selected(node: Node3D) -> void:
	_bound_mesh = node as MeshInstance3D


## --- Armature creation ------------------------------------------------------

## Add a bone. parent="" means root. Returns the new bone index.
func add_bone(name: String, parent: String, position: Vector3, tail: Vector3 = Vector3.ZERO) -> int:
	_ensure_skeleton()
	var parent_idx := -1
	if parent != "":
		parent_idx = _find_bone(parent)
	var idx := _skeleton.get_bone_count()
	_skeleton.add_bone(name)
	if parent_idx >= 0:
		_skeleton.set_bone_parent(idx, parent_idx)
	var rest := Transform3D()
	# Point the bone's -Z toward its tail so rotations look natural.
	var dir := (tail - position)
	if dir.length_squared() > 0.000001:
		rest = Transform3D(_look_at_basis(dir.normalized()), position)
	else:
		rest = Transform3D(Basis.IDENTITY, position)
	_skeleton.set_bone_rest(idx, rest)
	_skeleton.set_bone_pose_position(idx, position)
	_bones.append({"name": name, "parent": parent, "index": idx})
	StudioBus.scene_modified.emit()
	return idx


func _ensure_skeleton() -> void:
	if _skeleton != null and is_instance_valid(_skeleton):
		return
	_skeleton = Skeleton3D.new()
	_skeleton.name = "Armature"
	if _host != null:
		_host.get_node_or_null("Rig").add_child(_skeleton) if _host.has_node("Rig") else _host.add_child(_skeleton)


func _find_bone(name: String) -> int:
	if _skeleton == null:
		return -1
	for i in range(_skeleton.get_bone_count()):
		if _skeleton.get_bone_name(i) == name:
			return i
	return -1


func _look_at_basis(dir: Vector3) -> Basis:
	var up := Vector3.UP
	if absf(dir.dot(up)) > 0.99:
		up = Vector3.RIGHT
	var right := dir.cross(up).normalized()
	var real_up := right.cross(dir).normalized()
	return Basis(right, real_up, -dir)


## Build a simple starter humanoid-ish spine + limbs for a mesh's bounds.
func generate_default_armature(mesh: MeshInstance3D) -> void:
	if mesh == null or mesh.mesh == null:
		return
	var size := mesh.mesh.get_aabb().size
	var h := size.y
	add_bone("root", "", Vector3.ZERO, Vector3(0, h * 0.1, 0))
	add_bone("spine", "root", Vector3(0, h * 0.1, 0), Vector3(0, h * 0.35, 0))
	add_bone("chest", "spine", Vector3(0, h * 0.35, 0), Vector3(0, h * 0.6, 0))
	add_bone("neck", "chest", Vector3(0, h * 0.6, 0), Vector3(0, h * 0.72, 0))
	add_bone("head", "neck", Vector3(0, h * 0.72, 0), Vector3(0, h * 0.95, 0))
	add_bone("arm_l", "chest", Vector3(-size.x * 0.4, h * 0.6, 0), Vector3(-size.x * 0.9, h * 0.45, 0))
	add_bone("arm_r", "chest", Vector3(size.x * 0.4, h * 0.6, 0), Vector3(size.x * 0.9, h * 0.45, 0))
	add_bone("leg_l", "root", Vector3(-size.x * 0.2, 0, 0), Vector3(-size.x * 0.2, -h * 0.5, 0))
	add_bone("leg_r", "root", Vector3(size.x * 0.2, 0, 0), Vector3(size.x * 0.2, -h * 0.5, 0))
	bind_weights(mesh)


## --- Automatic weight binding -----------------------------------------------

## Compute skinning arrays using inverse-distance to bone segments and apply.
## `max_influences` must be <= 4 for the standard mobile skinning path.
func bind_weights(mesh: MeshInstance3D, max_influences: int = 4) -> void:
	if mesh == null or _skeleton == null or mesh.mesh == null or mesh.mesh.get_surface_count() == 0:
		return
	var arr := mesh.mesh.surface_get_arrays(0)
	var verts: PackedVector3Array = arr[Mesh.ARRAY_VERTEX]
	var n := verts.size()
	var bone_count := _skeleton.get_bone_count()
	var bone_segments: Array = []
	for i in range(bone_count):
		var pos := _skeleton.get_bone_global_rest(i).origin
		bone_segments.append({"idx": i, "pos": pos})

	var vbones := PackedInt32Array()
	var vweights := PackedFloat32Array()
	vbones.resize(n * 4)
	vweights.resize(n * 4)
	var to_mesh := mesh.global_transform.affine_inverse()

	for vi in range(n):
		var p := verts[vi]
		# Rank bones by distance to their head position.
		var dists: Array = []
		for b in bone_segments:
			var local_bone: Vector3 = to_mesh * (mesh.global_transform * b["pos"])
			dists.append({"idx": b["idx"], "d": p.distance_to(local_bone)})
		dists.sort_custom(func(a, b): return a["d"] < b["d"])
		var influ: int = mini(max_influences, dists.size())
		var wsum := 0.0
		var ws: Array[float] = []
		for k in range(influ):
			var d: float = maxf(0.0001, dists[k]["d"])
			var w := 1.0 / (d * d)
			ws.append(w)
			wsum += w
		for slot in range(4):
			var base := vi * 4 + slot
			if slot < influ and wsum > 0.0:
				vbones[base] = dists[slot]["idx"]
				vweights[base] = ws[slot] / wsum
			else:
				vbones[base] = 0
				vweights[base] = 0.0

	# Rebuild the mesh surface with skinning channels added.
	arr[Mesh.ARRAY_BONES] = vbones
	arr[Mesh.ARRAY_WEIGHTS] = vweights
	var new_mesh := MeshPool.acquire_mesh()
	new_mesh.add_surface_from_arrays(Mesh.PRIMITIVE_TRIANGLES, arr)
	mesh.mesh = new_mesh
	mesh.skeleton = mesh.get_path_to(_skeleton)
	# A Skin must be bound so the renderer knows which skeleton supplies the
	# bone transforms for these indices. add_bind(bone_index, rest_pose).
	var skin := Skin.new()
	for i in range(bone_count):
		skin.add_bind(i, _skeleton.get_bone_global_rest(i))
	mesh.skin = skin
	StudioBus.scene_modified.emit()
	print("[Rigging] bound %d verts to %d bones" % [n, bone_count])


func bone_count() -> int:
	return _skeleton.get_bone_count() if _skeleton != null else 0


func get_skeleton() -> Skeleton3D:
	return _skeleton