extends Object
class_name SkinJointHelper

## Articulates a joint through a CHAIN of helper bones so the limb cannot pinch.
##
## Linear-blend skinning averages TRANSFORMS, not shapes. Blending two rotations
## that are θ apart shortens anything perpendicular to the bend by cos(θ/2), so a
## vertex sitting between the thigh and the shin is pulled inside the leg the
## moment the knee folds. That is the thigh going thin during a jump.
##
## The collapse depends only on the WIDEST angle any single vertex has to blend
## across, so the fix is to make that angle small. Insert SEGMENTS-1 helper bones
## through the joint, each driven to its fraction of the child's rotation, and
## give every blend-zone vertex to the two ADJACENT bones it falls between. No
## vertex then blends across more than θ/SEGMENTS, and the worst collapse drops
## from cos(θ/2) to cos(θ/2·SEGMENTS):
##
##     θ = 120°     1 step (raw)  0.50     2 steps (half-angle)  0.87
##                  4 steps       0.97     6 steps               0.985
##
## Nothing is scaled and nothing is snapped rigid — the limb is simply
## articulated finely enough that averaging no longer eats it. This is why the
## fix is more joints rather than volume-correction: a bone scale fattens every
## vertex on the bone whether it was collapsing or not (tried; it ballooned parts
## of the thigh to 167% while the worst vertices stayed put).
##
## KNEES ONLY. Hip and ankle helpers were tried twice and measured worse both
## times — the hip swings rather than folds, so it was adding a blend for
## nothing, and it dragged the skirt down with it (0.80 -> 0.72). Re-measure
## before adding a joint here.
##
## install() returns [child_bone, helper_bone, fraction] triples; the caller MUST
## drive them every frame from inside the skeleton's modification pass (see
## SkinnedPlayerModel.ShooterPoseModifier). A helper that is not updated in step
## with its child deforms the limb instead of saving it.

## The CHILD bone of each joint to subdivide. The parent is whatever the
## skeleton says it is, not a second hardcoded name.
##
## It used to be a [parent, child] pair of ["DEF-thigh.L", "DEF-shin.L"], which
## silently did nothing on a rig with limb twist bones: Taila's shin hangs off
## DEF-thigh.L.001, so the knee vertices are weighted across THAT and the shin,
## and the pass found no vertex holding both named bones. The knee measured 0.76
## with this "installed" and doing nothing at all.
const JOINT_CHILDREN := ["DEF-shin.L", "DEF-shin.R", "shin.L", "shin.R"]
## Angular steps through each joint. 4 leaves at most a quarter of the bend for
## any one vertex to blend across, which is a ~3% collapse at a hard tuck.
const SEGMENTS := 4
## A vertex needs at least this much of BOTH bones before it is worth moving.
const MIN_BLEND := 0.02
## And the joint must hold at least this share of the vertex overall, or it is a
## stray cross-body influence rather than the joint this vertex belongs to.
const DOMINANT_SHARE := 0.5


## Returns an Array of [child_bone_idx, helper_bone_idx, fraction].
static func install(root: Node, skeleton: Skeleton3D) -> Array:
	if skeleton == null:
		return []
	var driven: Array = []
	# child bone -> [parent bone, [helper bones, inner first]]
	var joint_helpers := {}
	for child_name in JOINT_CHILDREN:
		var child := skeleton.find_bone(child_name)
		if child < 0:
			continue
		var parent := skeleton.get_bone_parent(child)
		if parent < 0:
			continue
		if skeleton.find_bone("HELPER1-" + child_name) >= 0:
			continue  # already installed
		var helpers: Array = []
		for step in range(1, SEGMENTS):
			var hname: String = "HELPER%d-%s" % [step, child_name]
			skeleton.add_bone(hname)
			var h := skeleton.find_bone(hname)
			skeleton.set_bone_parent(h, parent)
			# Same rest as the child, so "this fraction of the child's local
			# rotation" lands the helper at that fraction of the bend.
			skeleton.set_bone_rest(h, skeleton.get_bone_rest(child))
			skeleton.reset_bone_pose(h)
			helpers.append(h)
			driven.append([child, h, float(step) / float(SEGMENTS)])
		joint_helpers[child] = [parent, helpers]
	if driven.is_empty():
		return []

	var chain_l := _leg_chain(skeleton, ".L")
	var chain_r := _leg_chain(skeleton, ".R")
	var limb_radius := 0.09
	if chain_l.size() > 0 and chain_r.size() > 0:
		limb_radius = maxf(absf(chain_l[0].x - chain_r[0].x) * 0.5, 0.02)
	for mi in root.find_children("*", "MeshInstance3D", true, false):
		if mi.mesh == null or mi.skin == null:
			continue
		if mi.mesh.get_blend_shape_count() > 0:
			continue  # rebuilding would drop the blend shapes
		_reweight(mi, skeleton, joint_helpers, chain_l, chain_r, limb_radius)
	return driven


## Rest-pose joints down one leg, as a polyline for limb-membership tests.
static func _leg_chain(skeleton: Skeleton3D, suffix: String) -> PackedVector3Array:
	var out := PackedVector3Array()
	for stem in ["DEF-thigh", "DEF-shin", "DEF-foot", "DEF-toe"]:
		var i := skeleton.find_bone(stem + suffix)
		if i >= 0:
			out.append(skeleton.get_bone_global_rest(i).origin)
	return out


static func _dist_to_chain(p: Vector3, chain: PackedVector3Array) -> float:
	var best := INF
	for i in range(chain.size() - 1):
		var a: Vector3 = chain[i]
		var ab: Vector3 = chain[i + 1] - a
		var l2: float = ab.length_squared()
		var t: float = 0.0 if l2 < 0.000001 else clampf((p - a).dot(ab) / l2, 0.0, 1.0)
		best = minf(best, p.distance_to(a + ab * t))
	return best


## Drive the helpers. MUST run inside the skeleton's modification pass.
static func update(skeleton: Skeleton3D, driven: Array) -> void:
	for d in driven:
		skeleton.set_bone_pose_rotation(d[1], Quaternion.IDENTITY.slerp(
			skeleton.get_bone_pose_rotation(d[0]), d[2]))


static func _reweight(mi: MeshInstance3D, skeleton: Skeleton3D,
		joint_helpers: Dictionary, chain_l: PackedVector3Array,
		chain_r: PackedVector3Array, limb_radius: float) -> void:
	var safe_bounds := mi.mesh.get_aabb()
	var skin: Skin = mi.skin
	var bone_of := {}
	for b in skin.get_bind_count():
		var bi := skin.get_bind_bone(b)
		if bi < 0:
			bi = skeleton.find_bone(skin.get_bind_name(b))
		bone_of[b] = bi
	var bind_of_bone := {}
	for b in skin.get_bind_count():
		bind_of_bone[bone_of[b]] = b

	# Per joint: the ordered bind chain parent -> helpers -> child. Helpers share
	# the child's bind pose because their global rests are identical.
	var jobs: Array = []
	for child in joint_helpers:
		var parent: int = joint_helpers[child][0]
		var helpers: Array = joint_helpers[child][1]
		if not bind_of_bone.has(child) or not bind_of_bone.has(parent):
			continue
		var child_bind: int = bind_of_bone[child]
		var chain: Array = [bind_of_bone[parent]]
		for h in helpers:
			skin.add_named_bind(skeleton.get_bone_name(h),
				skin.get_bind_pose(child_bind))
			chain.append(skin.get_bind_count() - 1)
		chain.append(child_bind)
		jobs.append(chain)
	if jobs.is_empty():
		return

	var surfaces: Array = []
	var touched := 0
	for s in range(mi.mesh.get_surface_count()):
		var arrays: Array = mi.mesh.surface_get_arrays(s)
		touched += _reweight_surface(arrays, jobs, chain_l, chain_r, limb_radius)
		surfaces.append({
			"arrays": arrays,
			"material": mi.mesh.surface_get_material(s),
			"name": mi.mesh.surface_get_name(s),
		})
	if touched == 0:
		return
	var rebuilt := ArrayMesh.new()
	for i in surfaces.size():
		var e: Dictionary = surfaces[i]
		rebuilt.add_surface_from_arrays(Mesh.PRIMITIVE_TRIANGLES, e["arrays"])
		rebuilt.surface_set_material(i, e["material"])
		if e["name"] != "":
			rebuilt.surface_set_name(i, e["name"])
	# Preserve an explicit conservative bound across this second mesh rebuild.
	# See SkinLegRepair: sparse runtime Skin binds can otherwise produce more
	# bone AABBs than the renderer skeleton contains.
	var padding := Vector3.ONE * maxf(safe_bounds.size.length() * 0.12, 0.25)
	safe_bounds.position -= padding
	safe_bounds.size += padding * 2.0
	rebuilt.custom_aabb = safe_bounds
	mi.mesh = rebuilt


static func _reweight_surface(arrays: Array, jobs: Array,
		chain_l: PackedVector3Array, chain_r: PackedVector3Array,
		limb_radius: float) -> int:
	var verts: PackedVector3Array = arrays[Mesh.ARRAY_VERTEX]
	var bones: PackedInt32Array = arrays[Mesh.ARRAY_BONES]
	var weights: PackedFloat32Array = arrays[Mesh.ARRAY_WEIGHTS]
	if bones.is_empty() or verts.is_empty():
		return 0
	var per: int = bones.size() / verts.size()
	var touched := 0

	for v in verts.size():
		# Cloth is not a limb. The gate is generous — the outside of a knee sits
		# further from the bone axis than the limb radius and DOES need helping —
		# but without it the skirt and the outer boot flare get bound to a knee
		# helper and swing with the shin.
		if minf(_dist_to_chain(verts[v], chain_l),
				_dist_to_chain(verts[v], chain_r)) > limb_radius * 1.6:
			continue
		# Pick the joint this vertex most belongs to — the chain holding the most
		# of its weight — NOT merely the first chain that matches.
		#
		# Taking the first match bound right-thigh vertices that still carried a
		# little residual left-leg weight to the LEFT knee's helpers, which then
		# dragged them across the body: measured worst 0.46, and no amount of
		# extra subdivision touched it because the vertex was on the wrong
		# joint entirely.
		var best_chain: Array = []
		var best_kp := -1
		var best_kc := -1
		var best_total := 0.0
		for chain in jobs:
			var parent_bind: int = chain[0]
			var child_bind: int = chain[chain.size() - 1]
			var k_parent := -1
			var k_child := -1
			for k in per:
				var b: int = bones[v * per + k]
				if b == parent_bind:
					k_parent = k
				elif b == child_bind:
					k_child = k
			if k_parent < 0 or k_child < 0:
				continue
			var w_parent: float = weights[v * per + k_parent]
			var w_child: float = weights[v * per + k_child]
			if minf(w_parent, w_child) < MIN_BLEND:
				continue
			var total: float = w_parent + w_child
			if total > best_total:
				best_total = total
				best_chain = chain
				best_kp = k_parent
				best_kc = k_child
		# The joint must genuinely own this vertex. Below this it is a stray
		# cross-body influence, and re-binding it would pull the mesh across.
		if best_chain.is_empty() or best_total < DOMINANT_SHARE:
			continue
		# How far through the joint the vertex sits, then hand it to the two
		# ADJACENT bones either side of that point. Its widest blend is now one
		# segment instead of the whole joint.
		var w_c: float = weights[v * per + best_kc]
		var pos: float = (w_c / best_total) * float(SEGMENTS)
		var lo: int = clampi(int(floor(pos)), 0, SEGMENTS - 1)
		var frac: float = clampf(pos - float(lo), 0.0, 1.0)
		bones[v * per + best_kp] = best_chain[lo]
		weights[v * per + best_kp] = best_total * (1.0 - frac)
		bones[v * per + best_kc] = best_chain[lo + 1]
		weights[v * per + best_kc] = best_total * frac
		touched += 1
	if touched > 0:
		arrays[Mesh.ARRAY_BONES] = bones
		arrays[Mesh.ARRAY_WEIGHTS] = weights
	return touched