#!/usr/bin/env python3
"""
Headless Blender animation merger for Papaya-Shooter.

Merges every animation clip from the shared animation library onto a rigged
character (Mixamo-compatible bone names) and exports one game-ready GLB whose
animations use the game's canonical clip names (Idle, Walk, Run, Jump, ...).

Animation library — two supported layouts in <animations_dir>:

  A. `_library.glb` and optional `_library_v2.glb` — multi-clip libraries
     (Quaternius Universal Animation Libraries 1 and 2, CC0, committed at
     assets/characters/animations/). Clip names map through the tables below.
     The two editions use the same rest skeleton and proportions but different
     bone names; V2_BONE_MAP describes that exact rename.

  B. Loose FBX/GLB files, ONE clip per file, all on the Mixamo skeleton
     (manual Mixamo downloads). Filenames map to canonical clip names:
       idle.fbx        -> Idle
       run.fbx         -> Run
       crouch_walk.fbx -> CrouchWalk

Usage:
  blender --background --python tools/merge_animations.py -- \
      <rigged_character.glb> <animations_dir> <output.glb> [--keep-root-motion]

Locomotion clips are exported in place (hips horizontal motion stripped)
unless --keep-root-motion is given, since the gameplay code moves the body.
"""
import bpy
import math
import sys
import os
from mathutils import Matrix, Quaternion, Vector

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rig_map import RigRoles, build_map
from retarget_pose import (
    authored_world_rotation, build_segment_pairs, heading_inverse
)

argv = sys.argv
argv = argv[argv.index("--") + 1:] if "--" in argv else []
if len(argv) < 3:
    print("Usage: blender --background --python tools/merge_animations.py -- "
          "<character.glb> <animations_dir> <output.glb> [--keep-root-motion]")
    sys.exit(1)

CHARACTER = argv[0]
ANIM_DIR = argv[1]
OUTPUT = argv[2]
STRIP_ROOT_MOTION = "--keep-root-motion" not in argv

MIX = "mixamorig:"
UP = Vector((0.0, 0.0, 1.0))

# Quaternius Universal Animation Library clip -> game canonical clip name.
# Unmapped clips are skipped to keep character GLBs small.
LIBRARY_CLIP_MAP = {
    "Idle_Loop": "Idle",
    "Walk_Loop": "Walk",
    "Jog_Fwd_Loop": "Run",
    "Sprint_Loop": "Sprint",
    "Jump_Start": "Jump",
    "Jump_Loop": "Fall",
    "Jump_Land": "Land",
    "Crouch_Idle_Loop": "CrouchIdle",
    "Crouch_Fwd_Loop": "CrouchWalk",
    "Roll": "Roll",
    "Death01": "Death",
    "Hit_Chest": "Hit",
    "Dance_Loop": "Dance",
    # Movement-shooter extras:
    "Pistol_Idle_Loop": "PistolIdle",   # armed idle (weapon actually held up)
    "Pistol_Shoot": "PistolShoot",
    "Pistol_Reload": "PistolReload",
    "Sword_Attack": "SwordAttack",
}

# Quaternius Universal Animation Library 2 Standard adds authored traversal
# clips that the first library does not contain. Keep entry/loop/exit clips
# distinct: runtime code can now play the motion the animator authored instead
# of manufacturing a pose by rotating a few bones.
LIBRARY_V2_CLIP_MAP = {
    "Slide_Start": "SlideStart",
    "Slide_Loop": "Slide",
    "Slide_Exit": "SlideExit",
    "NinjaJump_Start": "WallRunStart",
    "NinjaJump_Land": "WallRunExit",
    "Idle_Rail_Loop": "WallCling",
    "ClimbUp_1m": "WallClimb",
    "Shield_Dash": "Dash",
    "Sword_Dash": "SwordDash",
    "OverhandThrow": "Throw",
    "Hit_Knockback": "Knockback",
    "Chest_Open": "EmoteStretch",
    "Idle_Rail_Call": "EmoteCall",
    "Yes": "EmoteYes",
    "Idle_No_Loop": "EmoteNo",
}

# Four authored combat-strafe cycles from the Godot TPS Demo. These are kept as
# a tiny animation-only GLB; see _directional_LICENSE.md. They give the runtime
# a real backward and lateral stride instead of yawing a forward run's hips.
DIRECTIONAL_CLIP_MAP = {
    "strafe_front-cycle": "StrafeWalkForward",
    "strafe_back-cycle": "StrafeWalkBackward",
}

MESH2MOTION_CLIP_MAP = {
    "Strafe_left": [("StrafeWalkLeft", 0.0), ("RunLeft", 0.0)],
    "Strafe_right": [("StrafeWalkRight", 0.0), ("RunRight", 0.0)],
    "Run_Anime": [("RunForward", 0.0)],
    # Authored held flight/reach with only subtle foot motion.
    "Flying Forward Super": [("Grapple", 0.0)],
}

WALLRUN_CLIP_MAP = {
    "WallRunLeft": "WallRunLeft",
    "WallRunRight": "WallRunRight",
}

CMU_CLIP_MAP = {
    "RunBackward": "RunBackward",
}

# UAL2 renamed the original Quaternius/Rigify bones to Unreal-style names, but
# their rest transforms are byte-for-byte equivalent (verified in Blender).
# The map is target/game bone -> UAL2 source bone.
V2_BONE_MAP = {
    "root": "root",
    "DEF-hips": "pelvis",
    "DEF-spine.001": "spine_01",
    "DEF-spine.002": "spine_02",
    "DEF-spine.003": "spine_03",
    "DEF-neck": "neck_01",
    "DEF-head": "Head",
    "DEF-shoulder.L": "clavicle_l",
    "DEF-upper_arm.L": "upperarm_l",
    "DEF-forearm.L": "lowerarm_l",
    "DEF-hand.L": "hand_l",
    "DEF-shoulder.R": "clavicle_r",
    "DEF-upper_arm.R": "upperarm_r",
    "DEF-forearm.R": "lowerarm_r",
    "DEF-hand.R": "hand_r",
    "DEF-thigh.L": "thigh_l",
    "DEF-shin.L": "calf_l",
    "DEF-foot.L": "foot_l",
    "DEF-toe.L": "ball_l",
    "DEF-thigh.R": "thigh_r",
    "DEF-shin.R": "calf_r",
    "DEF-foot.R": "foot_r",
    "DEF-toe.R": "ball_r",
}
for _side, _suffix in (("L", "l"), ("R", "r")):
    for _game_digit, _source_digit in (
        ("f_index", "index"),
        ("f_middle", "middle"),
        ("f_ring", "ring"),
        ("f_pinky", "pinky"),
        ("thumb", "thumb"),
    ):
        for _joint in range(1, 4):
            V2_BONE_MAP[f"DEF-{_game_digit}.{_joint:02d}.{_side}"] = (
                f"{_source_digit}_{_joint:02d}_{_suffix}"
            )

# The TPS robot's lower-body rig is humanoid and its leg rest axes are within a
# few degrees of the Quaternius rig. Only the authored lower-body motion is
# transferred; the project's weapon layer keeps ownership of the upper body.
DIRECTIONAL_BONE_MAP = {
    "root": "root",
    "DEF-hips": "hips",
    "DEF-spine.001": "spine1",
    "DEF-spine.002": "spine2",
    "DEF-spine.003": "chest",
    "DEF-thigh.L": "thigh.L",
    "DEF-shin.L": "shin.L",
    "DEF-foot.L": "foot.L",
    "DEF-toe.L": "toe.L",
    "DEF-thigh.R": "thigh.R",
    "DEF-shin.R": "shin.R",
    "DEF-foot.R": "foot.R",
    "DEF-toe.R": "toe.R",
}


def clear_scene():
    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()


def action_fcurves(action):
    """Return an action's F-curves across Blender versions.

    Blender 4.4+/5.x replaced `Action.fcurves` with the slotted-action system
    (layers -> strips -> channelbags -> fcurves). This yields the curves either
    way so the merge logic doesn't care which Blender it runs under.
    """
    legacy = getattr(action, "fcurves", None)
    if legacy is not None:
        return list(legacy)
    out = []
    for layer in getattr(action, "layers", []):
        for strip in layer.strips:
            cbags = getattr(strip, "channelbags", None)
            if cbags is not None:
                for cbag in cbags:
                    out.extend(cbag.fcurves)
            else:
                for slot in getattr(action, "slots", []):
                    cbag = strip.channelbag(slot)
                    if cbag:
                        out.extend(cbag.fcurves)
    return out


def to_pascal(stem):
    return "".join(part.capitalize() for part in stem.replace("-", "_").split("_"))


def find_armature(objects):
    for o in objects:
        if o.type == "ARMATURE":
            return o
    return None


def normalize_prefix(name):
    """mixamorig1:Hips / mixamorig_Hips / Hips -> mixamorig:Hips"""
    base = name
    for i in range(10):
        base = base.replace(f"mixamorig{i}:", MIX)
    base = base.replace("mixamorig_", MIX)
    if ":" not in base and base in CORE_BONES:
        base = MIX + base
    return base


CORE_BONES = {
    "Hips", "Spine", "Spine1", "Spine2", "Neck", "Head",
    "LeftShoulder", "LeftArm", "LeftForeArm", "LeftHand",
    "RightShoulder", "RightArm", "RightForeArm", "RightHand",
    "LeftUpLeg", "LeftLeg", "LeftFoot", "LeftToeBase",
    "RightUpLeg", "RightLeg", "RightFoot", "RightToeBase",
}


def normalize_action_paths(action):
    for fc in action_fcurves(action):
        if 'pose.bones["' in fc.data_path:
            start = fc.data_path.index('"') + 1
            end = fc.data_path.index('"', start)
            bone = fc.data_path[start:end]
            fixed = normalize_prefix(bone)
            if fixed != bone:
                fc.data_path = fc.data_path[:start] + fixed + fc.data_path[end:]


def hips_height(arm):
    for b in arm.data.bones:
        if b.name.lower().endswith("hips") or b.name.lower() == "pelvis":
            return (arm.matrix_world @ b.head_local).z
    return 1.0


def scale_location_curves(action, ratio):
    if abs(ratio - 1.0) < 0.01:
        return
    for fc in action_fcurves(action):
        if fc.data_path.endswith(".location"):
            for kp in fc.keyframe_points:
                kp.co.y *= ratio
                kp.handle_left.y *= ratio
                kp.handle_right.y *= ratio


def strip_hips_horizontal(action):
    """Zero horizontal root motion so clips play in place (gameplay code
    moves the body). Applies to hips AND any dedicated root bone; keeps the
    vertical channel so bob/land weight survives.
    """
    for fc in action_fcurves(action):
        if not fc.data_path.endswith(".location"):
            continue
        path_lower = fc.data_path.lower()
        if ("hips" in path_lower or '"root"' in path_lower) and fc.array_index in (0, 2):
            for kp in fc.keyframe_points:
                kp.co.y = 0.0
                kp.handle_left.y = 0.0
                kp.handle_right.y = 0.0


def main():
    clear_scene()
    print(f"Importing character {CHARACTER}...")
    bpy.ops.import_scene.gltf(filepath=CHARACTER)
    target_arm = find_armature(bpy.data.objects)
    if not target_arm:
        print("ERROR: character has no armature — run tools/autorig.py first")
        sys.exit(1)

    # Preserve the character's authored bone names. The rig sidecar is the
    # runtime contract for IK, aim and cloth, so silently renaming Hips/Head
    # here would invalidate otherwise-correct metadata on artist rigs. Role
    # mapping makes a shared naming convention unnecessary.

    target_hips = hips_height(target_arm)
    print(f"Character hips height: {target_hips:.3f}")
    target_roles = RigRoles(target_arm)
    missing = target_roles.missing_core()
    if missing:
        print(f"ERROR: could not identify target rig roles: {missing}")
        sys.exit(1)

    if not target_arm.animation_data:
        target_arm.animation_data_create()

    # Drop any pre-existing animation so only canonical clips ship.
    for track in list(target_arm.animation_data.nla_tracks):
        target_arm.animation_data.nla_tracks.remove(track)
    target_arm.animation_data.action = None
    for action in list(bpy.data.actions):
        if action.users == 0:
            bpy.data.actions.remove(action)

    library_specs = (
        ("_library.glb", LIBRARY_CLIP_MAP),
        ("_library_v2.glb", LIBRARY_V2_CLIP_MAP),
        ("_directional.glb", DIRECTIONAL_CLIP_MAP),
        ("_mesh2motion.glb", MESH2MOTION_CLIP_MAP),
        ("_cmu_locomotion.glb", CMU_CLIP_MAP),
        ("_wallrun.glb", WALLRUN_CLIP_MAP),
    )
    merged = 0
    found_library = False
    for filename, clip_map in library_specs:
        library_path = os.path.join(ANIM_DIR, filename)
        if not os.path.exists(library_path):
            continue
        found_library = True
        merged += merge_from_library(
            target_arm, target_roles, library_path, clip_map
        )
    if not found_library:
        merged = merge_from_files(target_arm, target_hips)

    if merged == 0:
        print("ERROR: no animations merged")
        sys.exit(1)

    print(f"Merged {merged} clips. Exporting {OUTPUT}...")
    os.makedirs(os.path.dirname(os.path.abspath(OUTPUT)), exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=OUTPUT,
        export_format="GLB",
        export_yup=True,
        export_animations=True,
        export_animation_mode="NLA_TRACKS",
        export_skins=True,
        # Only mapped body bones are keyed. Baking here would add frozen
        # rest-pose tracks to every hair/skirt bone and steal them from the
        # runtime spring solver.
        export_bake_animation=False,
        export_optimize_animation_size=True,
        export_optimize_animation_keep_anim_armature=False,
    )
    print("Done.")


def _add_clip_track(target_arm, action, clip_name):
    action.name = clip_name
    track = target_arm.animation_data.nla_tracks.new()
    track.name = clip_name
    strip = track.strips.new(clip_name, 0, action)
    strip.name = clip_name
    track.mute = True
    action.use_fake_user = True


def _assign_action(obj, action):
    """Assign an action to an object across Blender versions (slotted actions)."""
    if not obj.animation_data:
        obj.animation_data_create()
    obj.animation_data.action = action
    try:
        slots = action.slots
        if len(slots):
            obj.animation_data.action_slot = slots[0]
    except Exception:
        pass


def _world_rest(arm):
    return {
        bone.name: arm.matrix_world @ bone.matrix_local
        for bone in arm.data.bones
    }


def _bone_order(arm):
    out = []

    def walk(bone):
        out.append(bone.name)
        for child in bone.children:
            walk(child)

    for bone in arm.data.bones:
        if bone.parent is None:
            walk(bone)
    return out


def _rig_forward(arm, roles):
    acc = Vector((0.0, 0.0, 0.0))
    for side in ("L", "R"):
        for role in ("toe", "foot"):
            name = roles.limb.get((role, side))
            if not name:
                continue
            bone = arm.data.bones[name]
            direction = (
                (arm.matrix_world @ bone.tail_local)
                - (arm.matrix_world @ bone.matrix_local.translation)
            )
            direction.z = 0.0
            if direction.length > 1e-5:
                acc += direction.normalized()
            break
    return acc.normalized() if acc.length > 1e-5 else None


def _facing_correction(src_arm, src_roles, target_arm, target_roles):
    source = _rig_forward(src_arm, src_roles)
    target = _rig_forward(target_arm, target_roles)
    if source is None or target is None:
        return Quaternion()
    angle = Vector((source.x, source.y)).angle_signed(
        Vector((target.x, target.y)), 0.0
    )
    return Quaternion(UP, angle) if abs(angle) >= 1e-4 else Quaternion()


def _solve_pose(arm, order, rest_world, desired_rot, hips, hips_head):
    pose_world = {}
    basis = {}
    bones = arm.data.bones
    for name in order:
        bone = bones[name]
        rest = rest_world[name]
        if bone.parent is not None:
            parent = bone.parent.name
            matrix = (
                pose_world[parent]
                @ rest_world[parent].inverted()
                @ rest
            )
        else:
            matrix = rest
        rotation = Quaternion()
        if name in desired_rot:
            rotation = matrix.to_quaternion().inverted() @ desired_rot[name]
        location = Vector((0.0, 0.0, 0.0))
        if name == hips and hips_head is not None:
            location = matrix.inverted() @ hips_head
        basis[name] = (location, rotation)
        pose_world[name] = (
            matrix
            @ Matrix.Translation(location)
            @ rotation.to_matrix().to_4x4()
        )
    return basis


def retarget_action(target_arm, target_roles, src_arm, src_roles, mapping,
                    src_action, clip_name, yaw, scale,
                    lower_body_yaw_degrees=0.0):
    """Bake an authored motion onto any recognized humanoid rig.

    Major joint-to-joint directions come from the source performance while the
    target keeps its own bone roll. This prevents a crossed/T-pose modelling
    rest from surviving into every clip without twisting artist rigs.
    """
    src_rest = _world_rest(src_arm)
    target_rest = _world_rest(target_arm)
    target_order = _bone_order(target_arm)
    src_rest_rot = {
        name: matrix.to_quaternion() for name, matrix in src_rest.items()
    }
    target_rest_rot = {
        name: matrix.to_quaternion() for name, matrix in target_rest.items()
    }
    yaw_inverse = yaw.inverted()
    segment_pairs = build_segment_pairs(mapping, src_roles, target_roles)
    src_hips = src_roles.hips
    target_hips = target_roles.hips
    src_hips_rest = src_rest[src_hips].translation.copy()
    target_hips_rest = target_rest[target_hips].translation.copy()

    _assign_action(src_arm, src_action)
    frame_start, frame_end = (
        int(round(value)) for value in src_action.frame_range
    )
    source_heading = Quaternion()
    if clip_name == "RunBackward":
        bpy.context.scene.frame_set(frame_start)
        source_eval = src_arm.evaluated_get(
            bpy.context.evaluated_depsgraph_get()
        )
        source_heading = heading_inverse(
            (
                source_eval.matrix_world
                @ source_eval.pose.bones[src_hips].matrix
            ).to_quaternion(),
            src_rest_rot[src_hips],
            UP,
        )
    baked = bpy.data.actions.new(clip_name)
    _assign_action(target_arm, baked)
    for pose_bone in target_arm.pose.bones:
        pose_bone.rotation_mode = "QUATERNION"

    scene = bpy.context.scene
    for frame in range(frame_start, frame_end + 1):
        scene.frame_set(frame)
        depsgraph = bpy.context.evaluated_depsgraph_get()
        src_eval = src_arm.evaluated_get(depsgraph)
        source_world = src_eval.matrix_world
        desired = {}
        for target_name, source_name in mapping.items():
            if (
                source_name not in src_eval.pose.bones
                or target_name not in target_rest_rot
            ):
                continue
            target_rotation = authored_world_rotation(
                target_name,
                source_name,
                src_eval,
                source_world,
                src_rest_rot,
                target_rest,
                target_rest_rot,
                yaw,
                yaw_inverse,
                segment_pairs,
                source_heading,
            )
            desired[target_name] = target_rotation

        hips_head = None
        if src_hips in src_eval.pose.bones:
            moved = (
                source_world @ src_eval.pose.bones[src_hips].matrix
            ).translation
            displacement = yaw @ (source_heading @ (
                (moved - src_hips_rest) * scale
            ))
            if STRIP_ROOT_MOTION:
                displacement.x = 0.0
                displacement.y = 0.0
            hips_head = target_hips_rest + displacement

        basis = _solve_pose(
            target_arm,
            target_order,
            target_rest,
            desired,
            target_hips,
            hips_head,
        )
        for name in mapping:
            if name not in basis:
                continue
            pose_bone = target_arm.pose.bones[name]
            pose_bone.rotation_quaternion = basis[name][1]
            pose_bone.keyframe_insert("rotation_quaternion", frame=frame)
        if hips_head is not None:
            pose_bone = target_arm.pose.bones[target_hips]
            pose_bone.location = basis[target_hips][0]
            pose_bone.keyframe_insert("location", frame=frame)

    _assign_action(target_arm, None)
    return baked


def merge_from_library(target_arm, target_roles, library_path, clip_map):
    """Retarget every mapped clip from a single multi-clip library GLB onto the
    character rig (whose arm rest was fitted to the mesh)."""
    print(f"Using multi-clip library: {library_path}")
    before = set(bpy.data.objects)
    before_actions = set(bpy.data.actions)
    bpy.ops.import_scene.gltf(filepath=library_path)
    new_objects = [o for o in bpy.data.objects if o not in before]
    new_actions = [a for a in bpy.data.actions if a not in before_actions]
    src_arm = find_armature(new_objects)
    if not src_arm:
        print("ERROR: no armature in animation library")
        sys.exit(1)
    src_roles = RigRoles(src_arm)
    missing = src_roles.missing_core()
    if missing:
        print(f"ERROR: could not identify source rig roles: {missing}")
        sys.exit(1)
    mapping = build_map(src_roles, target_roles, {})
    if len(mapping) < 10:
        print(f"ERROR: only {len(mapping)} role-mapped bones in {library_path}")
        sys.exit(1)
    yaw = _facing_correction(src_arm, src_roles, target_arm, target_roles)
    src_height = (
        src_arm.matrix_world
        @ src_arm.data.bones[src_roles.hips].matrix_local
    ).translation.z
    target_height = (
        target_arm.matrix_world
        @ target_arm.data.bones[target_roles.hips].matrix_local
    ).translation.z
    scale = target_height / src_height if src_height > 1e-5 else 1.0
    print(
        f"Role-mapped {len(mapping)} target bones; "
        f"hips scale {scale:.3f}"
    )

    merged = 0
    keep = []
    for action in new_actions:
        base = action.name.split(".")[0]  # tolerate .001 suffixes
        mapped = clip_map.get(base)
        if not mapped:
            continue
        specs = [(mapped, 0.0)] if isinstance(mapped, str) else mapped
        for spec in specs:
            clip_name, travel_yaw_degrees = spec[:2]
            print(
                f"--- {base} -> '{clip_name}' (retarget bake, "
                f"lower-body heading {travel_yaw_degrees:+.0f} deg)"
            )
            baked = retarget_action(
                target_arm,
                target_roles,
                src_arm,
                src_roles,
                mapping,
                action,
                clip_name,
                yaw,
                scale,
                travel_yaw_degrees,
            )
            target_arm.animation_data.action = None
            _add_clip_track(target_arm, baked, clip_name)
            keep.append(baked)
            merged += 1

    for o in new_objects:
        bpy.data.objects.remove(o, do_unlink=True)
    for a in list(bpy.data.actions):
        if a not in keep and a in new_actions:
            bpy.data.actions.remove(a)
    return merged


def merge_from_files(target_arm, target_hips):
    """Merge clips from loose one-clip-per-file FBX/GLB files."""
    anim_files = sorted(
        f for f in os.listdir(ANIM_DIR)
        if f.lower().endswith((".fbx", ".glb", ".gltf")) and not f.startswith((".", "_"))
    )
    if not anim_files:
        print(f"ERROR: no animation files in {ANIM_DIR}")
        sys.exit(1)

    merged = 0
    for fname in anim_files:
        clip_name = to_pascal(os.path.splitext(fname)[0])
        path = os.path.join(ANIM_DIR, fname)
        print(f"--- {fname} -> '{clip_name}'")

        before = set(bpy.data.objects)
        before_actions = set(bpy.data.actions)
        try:
            if fname.lower().endswith(".fbx"):
                bpy.ops.import_scene.fbx(filepath=path, ignore_leaf_bones=True)
            else:
                bpy.ops.import_scene.gltf(filepath=path)
        except Exception as e:
            print(f"    SKIP: import failed: {e}")
            continue

        new_objects = [o for o in bpy.data.objects if o not in before]
        new_actions = [a for a in bpy.data.actions if a not in before_actions]
        src_arm = find_armature(new_objects)

        if not new_actions:
            print("    SKIP: no action found in file")
        else:
            action = max(new_actions, key=lambda a: len(action_fcurves(a)))
            normalize_action_paths(action)
            if src_arm:
                ratio = target_hips / max(hips_height(src_arm), 1e-6)
                scale_location_curves(action, ratio)
            if STRIP_ROOT_MOTION:
                strip_hips_horizontal(action)
            _add_clip_track(target_arm, action, clip_name)
            merged += 1

        # Remove the imported helper objects (keep the action).
        for o in new_objects:
            bpy.data.objects.remove(o, do_unlink=True)
        for a in new_actions:
            if a.name != clip_name:
                bpy.data.actions.remove(a)

    return merged


main()