"""merge_layer.py — layer upper body (Praying, reversed+looped) onto lower body (Kneel).

Both inputs are identical mixamorig skeletons. We take the upper-body bone tracks
(Spine chain, neck/head, both arms + hands/fingers) from Praying played BACKWARDS and
LOOPED, and keep the lower body (Hips, legs, feet) from Kneel. In-place fcurve rewrite
(same proven pattern as tools/reverse_clip.py). Then export a clean GLB and render a
color-coded stick-figure MP4 preview (upper = cyan, lower = orange).
"""
import bpy, sys, re, argparse
from mathutils import Vector

def parse():
    argv = sys.argv[sys.argv.index("--")+1:] if "--" in sys.argv else []
    p = argparse.ArgumentParser()
    p.add_argument("--upper", required=True)   # Praying (reversed+looped)
    p.add_argument("--lower", required=True)   # Kneel
    p.add_argument("--glb", required=True)
    p.add_argument("--mp4", required=True)
    p.add_argument("--name", default="PrayKneel")
    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")  # Arm covers ForeArm; Hand covers fingers
def bone_of(dp):
    m = re.search(r'pose\.bones\["([^"]+)"\]', dp)
    return m.group(1) if m else ""
def is_upper(bone):
    b = bone.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

    # import lower (Kneel) first -> base armature we keep
    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]

    # import upper (Praying)
    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

    kf0, kf1 = int(lower_act.frame_range[0]), int(lower_act.frame_range[1])
    kn = kf1 - kf0                       # lower length
    pf0, pf1 = int(upper_act.frame_range[0]), int(upper_act.frame_range[1])
    pn = pf1 - pf0                       # upper length
    period = pn + 1
    N = kn
    print(f"[merge] lower {kf0}..{kf1} (N={kn}) | upper {pf0}..{pf1} (period={period}) -> merged 0..{N}")

    # upper source fcurves keyed by (data_path, array_index)
    umap = {(fc.data_path, fc.array_index): fc for fc in action_fcurves(upper_act)}

    upper_ct = lower_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")
        # graft ONLY rotation from Praying — non-root bone location is a rest-length offset;
        # overwriting it detaches the torso from the hips. Location/scale stay from Kneel.
        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)]
            # reversed within each loop: local t=f%period, reversed = pn-t, sample pf0+reversed
            vals = [src.evaluate(pf0 + (pn - (f % period))) for f in range(N+1)]
            upper_ct += 1
        else:
            vals = [fc.evaluate(kf0 + min(f, kn)) for f in range(N+1)]
            lower_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"[merge] rewrote {upper_ct} upper (Praying) + {lower_ct} lower/base (Kneel) fcurves")

    lower_act.name = a.name

    # drop upper armature entirely
    for o in [upper_arm]:
        bpy.data.objects.remove(o, do_unlink=True)

    # single NLA track = clip name (pipeline convention)
    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"[merge] EXPORTED {a.glb} clip '{a.name}' ({N+1} frames)")

    # ---- stick-figure preview ----
    lower_arm.animation_data.action = lower_act  # re-attach for evaluation
    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]

    mat_up = bpy.data.materials.new("up"); mat_up.diffuse_color = (0.15, 0.75, 0.95, 1)
    mat_lo = bpy.data.materials.new("lo"); mat_lo.diffuse_color = (0.95, 0.5, 0.12, 1)

    # scale reference from bone lengths
    scn.frame_set(0)
    heads = [lower_arm.matrix_world @ pb.head for pb in lower_arm.pose.bones]
    size = max((max(v[i] for v in heads) - min(v[i] for v in heads)) for i in range(3)) or 1.0
    rad = size * 0.012

    for pb in lower_arm.pose.bones:
        if pb.bone.length < 1e-4:
            continue
        bpy.ops.mesh.primitive_cylinder_add(radius=rad, depth=1.0, vertices=8)
        cyl = bpy.context.active_object
        cyl.rotation_euler = (1.5708, 0, 0)  # cylinder long axis Z -> Y so STRETCH_TO (Y) aligns
        bpy.ops.object.transform_apply(rotation=True)
        up = is_upper(pb.name)
        cyl.data.materials.append(mat_up if up else mat_lo)
        cyl.color = (mat_up.diffuse_color if up else mat_lo.diffuse_color)
        c = cyl.constraints.new('COPY_LOCATION'); c.target = lower_arm; c.subtarget = pb.name; c.head_tail = 0.0
        s = cyl.constraints.new('STRETCH_TO'); s.target = lower_arm; s.subtarget = pb.name; s.head_tail = 1.0
        s.rest_length = 1.0; s.volume = 'NO_VOLUME'

    # camera framing over the whole motion
    mn = Vector(( 1e9, 1e9, 1e9)); mx = Vector((-1e9,-1e9,-1e9))
    for f in (0, N//4, N//2, 3*N//4, N):
        scn.frame_set(f)
        for pb in lower_arm.pose.bones:
            for p in (lower_arm.matrix_world @ pb.head, lower_arm.matrix_world @ pb.tail):
                for i in range(3):
                    mn[i] = min(mn[i], p[i]); mx[i] = max(mx[i], p[i])
    ctr = (mn + mx) * 0.5
    span = max((mx - mn)[i] for i in range(3))
    cam_data = bpy.data.cameras.new("cam"); cam = bpy.data.objects.new("cam", cam_data)
    scn.collection.objects.link(cam)
    d = span * 1.9
    cam.location = ctr + Vector((d*0.35, -d, d*0.10))
    cam.rotation_euler = (ctr - cam.location).to_track_quat('-Z', 'Y').to_euler()
    scn.camera = cam

    scn.render.engine = 'BLENDER_WORKBENCH'
    scn.display.shading.light = 'STUDIO'
    scn.display.shading.color_type = 'OBJECT'
    scn.render.resolution_x = 640; scn.render.resolution_y = 760
    scn.render.image_settings.file_format = 'PNG'
    scn.render.filepath = a.mp4  # here a.mp4 is a frame-dir prefix, e.g. .../frames/f_
    scn.frame_start = 0; scn.frame_end = N
    bpy.ops.render.render(animation=True)
    print(f"[merge] RENDERED frames to {a.mp4}####")

if __name__ == "__main__":
    main()