#!/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` — a single multi-clip library (e.g. Quaternius Universal
     Animation Library, CC0, committed at assets/characters/animations/).
     Clip names map through LIBRARY_CLIP_MAP below. The character must be
     rigged on this library's own skeleton (tools/autorig.py does this
     automatically when _library.glb is present).

  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 sys
import os

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:"

# 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": "Dash",
    "Death01": "Death",
    "Hit_Chest": "Hit",
    "Dance_Loop": "Dance",
    # Movement-shooter extras:
    "Swim_Fwd_Loop": "Grapple",     # superman reach reads as a swing pose
    "Pistol_Idle_Loop": "PistolIdle",   # armed idle (weapon actually held up)
    "Pistol_Shoot": "PistolShoot",
    "Pistol_Reload": "PistolReload",
    "Sword_Attack": "Throw",        # overhead swing reads as a grenade throw
}


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"):
            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)

    # Normalize character bone names too.
    for bone in target_arm.data.bones:
        fixed = normalize_prefix(bone.name)
        if fixed != bone.name:
            bone.name = fixed

    target_hips = hips_height(target_arm)
    print(f"Character hips height: {target_hips:.3f}")

    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

    library_path = os.path.join(ANIM_DIR, "_library.glb")
    if os.path.exists(library_path):
        merged = merge_from_library(target_arm, target_hips, library_path)
    else:
        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,
        export_bake_animation=True,
    )
    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 retarget_action(target_arm, src_arm, src_action, clip_name):
    """Retarget one source clip onto the character rig via world-space
    constraint baking.

    The character rig's arms were fitted to the mesh's A-pose (autorig.py), so
    its rest differs from the library's T-pose. Copying each bone's WORLD
    orientation (not its local channel) reproduces the library's actual motion
    regardless of that rest difference, then baking captures it as a clean
    action on the character rig.
    """
    _assign_action(src_arm, src_action)
    fr0 = int(src_action.frame_range[0])
    fr1 = int(src_action.frame_range[1])

    for pb in target_arm.pose.bones:
        if pb.name not in src_arm.pose.bones:
            continue
        cr = pb.constraints.new("COPY_ROTATION")
        cr.target = src_arm
        cr.subtarget = pb.name
        cr.target_space = "WORLD"
        cr.owner_space = "WORLD"
        if pb.name.lower().endswith("hips"):
            cl = pb.constraints.new("COPY_LOCATION")
            cl.target = src_arm
            cl.subtarget = pb.name
            cl.target_space = "WORLD"
            cl.owner_space = "WORLD"

    bpy.ops.object.select_all(action="DESELECT")
    target_arm.select_set(True)
    bpy.context.view_layer.objects.active = target_arm
    if target_arm.animation_data:
        target_arm.animation_data.action = None

    bpy.ops.nla.bake(frame_start=fr0, frame_end=fr1, step=1,
                     only_selected=False, visual_keying=True,
                     clear_constraints=True, clear_parents=False,
                     use_current_action=True, bake_types={"POSE"})

    baked = target_arm.animation_data.action
    if STRIP_ROOT_MOTION:
        strip_hips_horizontal(baked)
    return baked


def merge_from_library(target_arm, target_hips, library_path):
    """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)

    # Match the library rig's height to the character so hips bob / foot planting
    # translate correctly under the world-space location copy.
    src_hips = hips_height(src_arm)
    if src_hips > 1e-4:
        s = target_hips / src_hips
        src_arm.scale = (s, s, s)
        bpy.ops.object.select_all(action="DESELECT")
        src_arm.select_set(True)
        bpy.context.view_layer.objects.active = src_arm
        bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)

    merged = 0
    keep = []
    for action in new_actions:
        base = action.name.split(".")[0]  # tolerate .001 suffixes
        clip_name = LIBRARY_CLIP_MAP.get(base)
        if not clip_name:
            continue
        print(f"--- {base} -> '{clip_name}' (retarget bake)")
        baked = retarget_action(target_arm, src_arm, action, clip_name)
        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()