# Kevin Iglesias FBX → Quaternius-skeleton GLB batch retargeter (Blender 5.x, headless).
#
# THE one-time migration path (Ozan 2026-07-07 "I want to do this migration once"):
# originals live in ariki-assets/animations/kevin-iglesias (Unity FBX, T-pose "B-*" rig);
# this script bakes them onto the game's Quaternius UE-named skeleton so the runtime loads
# them exactly like the UAL GLBs (PlayerController.LoadUalLibrary) — no runtime retarget,
# no deprecated Godot BoneMap import path.
#
# Method: world-rotation-delta retarget. Both rigs rest in a T-POSE (verified 2026-07-07:
# Quaternius upperarm_l dir=(1,0,0), Kevin B-upperArm.L dir=(1,0,0)), so for every mapped
# bone the world rotation delta from rest transfers 1:1:
#     q_target_world(f) = (q_src_world(f) · q_src_rest⁻¹) · q_target_rest
# Pelvis additionally carries scaled root translation. Unmapped target bones (spine_03,
# leaf/ball bones, missing fingers) ride their parent rigidly at rest offset.
#
# Usage:
#   blender --background --python tools/kevin_retarget.py -- \
#       --src "<dir-or-fbx>[,<dir-or-fbx>…]" --out <out.glb> [--limit N]
#
# Each FBX becomes one glTF animation named from the file: "HumanM@Eat01_R - Loop" →
# "Eat01_R_Loop". Output GLB = armature + all animations (no mesh).

import bpy, sys, os, re
from mathutils import Matrix, Quaternion, Vector

TARGET_GLTF = "/Users/behcetozanbozkurt/Documents/tinqs-ltd/ariki-game/assets/quaternius/base-characters/Universal Base Characters[Standard]/Base Characters/Godot - UE/Superhero_Male_FullBody.gltf"

# Kevin "B-*" → Quaternius UE names. Fingers map 1:1 (both rigs have 3-segment fingers).
BONE_MAP = {
    "B-hips": "pelvis",
    "B-spine": "spine_01",
    "B-chest": "spine_02",
    "B-neck": "neck_01",
    "B-head": "Head",
}
for S, T in (("L", "l"), ("R", "r")):
    BONE_MAP[f"B-shoulder.{S}"] = f"clavicle_{T}"
    BONE_MAP[f"B-upperArm.{S}"] = f"upperarm_{T}"
    BONE_MAP[f"B-forearm.{S}"] = f"lowerarm_{T}"
    BONE_MAP[f"B-hand.{S}"] = f"hand_{T}"
    BONE_MAP[f"B-thigh.{S}"] = f"thigh_{T}"
    BONE_MAP[f"B-shin.{S}"] = f"calf_{T}"
    BONE_MAP[f"B-foot.{S}"] = f"foot_{T}"
    BONE_MAP[f"B-toe.{S}"] = f"ball_{T}"
    for i in (1, 2, 3):
        BONE_MAP[f"B-indexFinger0{i}.{S}"] = f"index_0{i}_{T}"
        BONE_MAP[f"B-middleFinger0{i}.{S}"] = f"middle_0{i}_{T}"
        BONE_MAP[f"B-ringFinger0{i}.{S}"] = f"ring_0{i}_{T}"
        BONE_MAP[f"B-pinky0{i}.{S}"] = f"pinky_0{i}_{T}"
        BONE_MAP[f"B-thumb0{i}.{S}"] = f"thumb_0{i}_{T}"
INV_MAP = {v: k for k, v in BONE_MAP.items()}


def apply_object_transform(obj):
    """Bake the importer's object-level rotation/scale into the armature data. Makes Blender
    world == data space (meters, Z-up) for BOTH rigs, so exported glTF matches what Godot's
    importer produced for the original Quaternius file — same units, same axes, same rests."""
    bpy.ops.object.select_all(action="DESELECT")
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    bpy.context.view_layer.update()


def clip_name(path):
    base = os.path.splitext(os.path.basename(path))[0]
    base = base.split("@", 1)[-1]                 # drop HumanM@ / HumanF@ / Villager@ …
    base = base.replace(" - ", "_").replace(" ", "")
    return re.sub(r"[^A-Za-z0-9_\-]", "", base)


def hierarchy_order(arm):
    out, stack = [], [b for b in arm.data.bones if b.parent is None]
    while stack:
        b = stack.pop(0)
        out.append(b)
        stack.extend(b.children)
    return out


def char_frame(rest, hips, head, upper_arm_l):
    """Orthonormal character frame from the rest pose: columns = (left-arm axis, hips→head up,
    their cross). Lets us align two rigs whatever world orientation their importer left them in
    — the old pipeline's limb bug was exactly an unhandled world-frame mismatch."""
    u = (rest[head].translation - rest[hips].translation).normalized()
    l = (rest[upper_arm_l].to_quaternion() @ Vector((0, 1, 0))).normalized()  # bone Y = along bone
    f = l.cross(u).normalized()
    l = u.cross(f).normalized()
    m = Matrix((l, u, f)).transposed()
    return m.to_quaternion().normalized()


def retarget_one(src_arm, tgt_arm, action, name, scn, unit):
    f0, f1 = int(action.frame_range[0]), int(action.frame_range[1])
    if src_arm.animation_data is None:
        src_arm.animation_data_create()
    src_arm.animation_data.action = action
    # Blender 5 slotted actions: bind the first slot so the action actually evaluates.
    if getattr(action, "slots", None) and src_arm.animation_data.action_slot is None:
        src_arm.animation_data.action_slot = action.slots[0]

    tgt_rest = {b.name: (tgt_arm.matrix_world @ b.matrix_local) for b in tgt_arm.data.bones}
    src_rest = {b.name: (src_arm.matrix_world @ b.matrix_local) for b in src_arm.data.bones}
    src_rest_q = {n: m.to_quaternion().normalized() for n, m in src_rest.items()}
    tgt_rest_q = {n: m.to_quaternion().normalized() for n, m in tgt_rest.items()}

    # World-frame alignment C: rotates source-world axes onto target-world axes (both rigs are
    # T-pose, so the character frames coincide semantically). All rotation deltas and root
    # translation deltas are conjugated/rotated through C before applying to the target.
    C = (char_frame(tgt_rest, "pelvis", "Head", "upperarm_l")
         @ char_frame(src_rest, "B-hips", "B-head", "B-upperArm.L").inverted()).normalized()
    Cinv = C.inverted()

    # Uniform size ratio from hips→head world distance (unit/scale agnostic).
    h_src = (src_rest["B-head"].translation - src_rest["B-hips"].translation).length
    h_tgt = (tgt_rest["Head"].translation - tgt_rest["pelvis"].translation).length
    tscale = h_tgt / h_src if h_src > 1e-5 else 1.0

    new_act = bpy.data.actions.new(name)
    if tgt_arm.animation_data is None:
        tgt_arm.animation_data_create()
    tgt_arm.animation_data.action = new_act

    order = hierarchy_order(tgt_arm)
    tw_inv = tgt_arm.matrix_world.inverted()
    for pb in tgt_arm.pose.bones:
        pb.rotation_mode = "QUATERNION"

    # ── Pelvis translation: the raw hips fcurve channel ──
    # Calibrated on Eat (static ≈0) + SleepGround Down (y 0→−86cm): Kevin's hips LOCAL location
    # is placement-free (the demo-scene offset lives on the unmapped B-root), cm-scaled, and
    # standing = 0 with −Y = down, X/Z = horizontal. Maps straight onto the target's Z-up world:
    # (x, z, y)·unit·tscale added to the pelvis rest. No world-frame algebra, no rest quirks.

    for f in range(f0, f1 + 1):
        scn.frame_set(f)
        dg = bpy.context.evaluated_depsgraph_get()
        se = src_arm.evaluated_get(dg)
        spose = {pb.name: (src_arm.matrix_world @ pb.matrix) for pb in se.pose.bones}

        tgt_world = {}
        for tb in order:
            nt = tb.name
            rest_t = tgt_rest[nt]
            ns = INV_MAP.get(nt)
            if ns is None or ns not in spose:
                if tb.parent is None:
                    tgt_world[nt] = rest_t
                else:
                    rel = tgt_rest[tb.parent.name].inverted() @ rest_t
                    tgt_world[nt] = tgt_world[tb.parent.name] @ rel
                continue
            q_delta = (spose[ns].to_quaternion().normalized() @ src_rest_q[ns].inverted()).normalized()
            q_delta = (C @ q_delta @ Cinv).normalized()   # express the delta in TARGET world axes
            q_world = (q_delta @ tgt_rest_q[nt]).normalized()
            if nt == "pelvis":
                loc = se.pose.bones["B-hips"].location
                t_world = rest_t.translation + Vector((loc.x, loc.z, loc.y)) * unit * tscale
            else:
                rel = tgt_rest[tb.parent.name].inverted() @ rest_t
                t_world = (tgt_world[tb.parent.name] @ rel).translation
            tgt_world[nt] = Matrix.LocRotScale(t_world, q_world, Vector((1, 1, 1)))

        for tb in order:
            pb = tgt_arm.pose.bones[tb.name]
            m_arm = tw_inv @ tgt_world[tb.name]
            if tb.parent:
                m_parent_arm = tw_inv @ tgt_world[tb.parent.name]
                rel_rest = tb.parent.matrix_local.inverted() @ tb.matrix_local
                basis = rel_rest.inverted() @ (m_parent_arm.inverted() @ m_arm)
            else:
                basis = tb.matrix_local.inverted() @ m_arm
            pb.matrix_basis = basis
            pb.keyframe_insert("rotation_quaternion", frame=f)
            if tb.name == "pelvis" or tb.parent is None:
                pb.keyframe_insert("location", frame=f)

    tgt_arm.animation_data.action = None
    return new_act


def main():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    args = dict(zip(argv[::2], argv[1::2]))
    srcs, out, limit = args["--src"].split(","), args["--out"], int(args.get("--limit", 0))

    fbx_files = []
    for s in srcs:
        s = s.strip()
        if os.path.isdir(s):
            for root, _, files in os.walk(s):
                fbx_files += [os.path.join(root, x) for x in sorted(files) if x.lower().endswith(".fbx")]
        else:
            fbx_files.append(s)
    if limit:
        fbx_files = fbx_files[:limit]
    print(f"[kevin] {len(fbx_files)} clips → {out}")

    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()
    scn = bpy.context.scene
    scn.render.fps = 30

    bpy.ops.import_scene.gltf(filepath=TARGET_GLTF)
    tgt_arm = next(o for o in bpy.data.objects if o.type == "ARMATURE")
    tgt_arm.name = "Armature"
    # Drop meshes — animations-only GLB (small, loads fast; the game binds clips to its own body).
    for o in [o for o in bpy.data.objects if o.type == "MESH"]:
        bpy.data.objects.remove(o, do_unlink=True)
    bpy.context.view_layer.update()
    apply_object_transform(tgt_arm)

    done = []
    for i, f in enumerate(fbx_files):
        name = clip_name(f)
        pre = set(bpy.data.objects) | set()
        pre_actions = set(bpy.data.actions)
        bpy.ops.import_scene.fbx(filepath=f)
        new_objs = [o for o in bpy.data.objects if o not in pre]
        src_arm = next((o for o in new_objs if o.type == "ARMATURE"), None)
        new_acts = [a for a in bpy.data.actions if a not in pre_actions]
        if src_arm is None or not new_acts:
            print(f"[kevin] SKIP (no armature/action): {f}")
        else:
            bpy.context.view_layer.update()
            # transform_apply keeps the ROTATION math in one consistent world frame (engine-
            # validated); the pelvis TRANSLATION reads the raw hips fcurve channel instead, so
            # capture the importer's unit scale (cm→m 0.01) before apply wipes it.
            unit = src_arm.scale.x
            apply_object_transform(src_arm)
            act = retarget_one(src_arm, tgt_arm, new_acts[0], name, scn, unit)
            done.append(act)
            print(f"[kevin] {i + 1}/{len(fbx_files)} {name} ({int(act.frame_range[1])}f)")
        for o in new_objs:
            bpy.data.objects.remove(o, do_unlink=True)
        for a in new_acts:
            bpy.data.actions.remove(a)

    # Stash every baked action on an NLA track so the glTF exporter emits it as an animation.
    for act in done:
        tr = tgt_arm.animation_data.nla_tracks.new()
        tr.name = act.name
        tr.strips.new(act.name, 1, act)
    bpy.ops.object.select_all(action="DESELECT")
    tgt_arm.select_set(True)
    bpy.context.view_layer.objects.active = tgt_arm
    bpy.ops.export_scene.gltf(
        filepath=out, use_selection=True,
        export_animations=True, export_animation_mode="NLA_TRACKS",
        export_force_sampling=True, export_optimize_animation_size=False)
    print(f"[kevin] EXPORTED {len(done)} clips → {out}")


main()