"""merge_hold.py — Pray01 v2: kneeling-prayer HOLD (loop-clean).

Lower body = Kneel held at its deepest-kneel frame (static, character stays kneeling).
Upper body = Praying looped FORWARD (rotation only) for gentle life — hands stay clasped.
Loops cleanly: lower is static, upper spans a whole number of Praying loop periods.
Matches Ozan's spec (kneel, hands together, hold ~2-3s) and the looping-pray pattern the
Kevin Reverence01_Loop stand-in used.
"""
import bpy, sys, re, argparse

def parse():
    argv = sys.argv[sys.argv.index("--")+1:] if "--" in sys.argv else []
    p = argparse.ArgumentParser()
    p.add_argument("--upper", required=True)      # Praying
    p.add_argument("--lower", required=True)      # Kneel
    p.add_argument("--hold", type=int, required=True)   # deepest-kneel frame in Kneel
    p.add_argument("--nframes", type=int, default=72)   # output length (72=2.4s=2 Praying loops)
    p.add_argument("--glb", required=True)
    p.add_argument("--name", default="PrayHold")
    return p.parse_args(argv)

def action_fcurves(a):
    if hasattr(a, "fcurves") and getattr(a, "fcurves", None) is not None:
        return list(a.fcurves)
    fcs = []
    for layer in getattr(a, "layers", []):
        for strip in layer.strips:
            for bag in strip.channelbags:
                fcs.extend(bag.fcurves)
    return fcs

UPPER_KEYS = ("Spine", "Neck", "Head", "Shoulder", "Arm", "Hand")
def bone_of(dp):
    m = re.search(r'pose\.bones\["([^"]+)"\]', dp); return m.group(1) if m else ""
def is_upper(b):
    b = b.split(":")[-1]; return any(k in b for k in UPPER_KEYS)

def main():
    a = parse()
    bpy.ops.wm.read_factory_settings(use_empty=True)
    scn = bpy.context.scene; scn.render.fps = 30

    bpy.ops.import_scene.fbx(filepath=a.lower)
    lower_arm = next(o for o in bpy.data.objects if o.type == 'ARMATURE')
    lower_act = lower_arm.animation_data.action
    if getattr(lower_act, "slots", None) and lower_arm.animation_data.action_slot is None:
        lower_arm.animation_data.action_slot = lower_act.slots[0]

    pre = set(bpy.data.objects)
    bpy.ops.import_scene.fbx(filepath=a.upper)
    upper_arm = next(o for o in (set(bpy.data.objects)-pre) if o.type == 'ARMATURE')
    upper_act = upper_arm.animation_data.action

    pf0 = int(upper_act.frame_range[0]); pf1 = int(upper_act.frame_range[1])
    period = pf1 - pf0 + 1
    N = a.nframes
    print(f"[hold] lower held @frame {a.hold} | upper forward-loop period {period} | out 0..{N}")

    umap = {(fc.data_path, fc.array_index): fc for fc in action_fcurves(upper_act)}

    up_ct = lo_ct = 0
    for fc in action_fcurves(lower_act):
        bone = bone_of(fc.data_path)
        is_rot = fc.data_path.endswith("rotation_quaternion") or fc.data_path.endswith("rotation_euler")
        upper = bone and is_upper(bone) and is_rot and (fc.data_path, fc.array_index) in umap
        if upper:
            src = umap[(fc.data_path, fc.array_index)]
            vals = [src.evaluate(pf0 + (f % period)) for f in range(N+1)]   # forward loop
            up_ct += 1
        else:
            hv = fc.evaluate(a.hold)                                        # static hold
            vals = [hv] * (N+1)
            lo_ct += 1
        need = N + 1; have = len(fc.keyframe_points)
        if have < need: fc.keyframe_points.add(need - have)
        elif have > need:
            for _ in range(have - need): fc.keyframe_points.remove(fc.keyframe_points[-1], fast=True)
        for k in range(need):
            kp = fc.keyframe_points[k]; kp.co = (k, vals[k]); kp.interpolation = 'LINEAR'
        fc.update()
    print(f"[hold] {up_ct} upper (Praying loop) + {lo_ct} lower/base (Kneel held) fcurves")

    lower_act.name = a.name
    bpy.data.objects.remove(upper_arm, do_unlink=True)
    for t in list(lower_arm.animation_data.nla_tracks):
        lower_arm.animation_data.nla_tracks.remove(t)
    track = lower_arm.animation_data.nla_tracks.new(); track.name = a.name
    track.strips.new(a.name, 0, lower_act).name = a.name
    lower_arm.animation_data.action = None
    scn.frame_start = 0; scn.frame_end = N
    bpy.ops.export_scene.gltf(filepath=a.glb, export_animation_mode="NLA_TRACKS",
                              export_force_sampling=True, export_optimize_animation_size=False)
    print(f"[hold] EXPORTED {a.glb} clip '{a.name}' ({N+1} frames)")

if __name__ == "__main__":
    main()