#!/usr/bin/env python3
"""Lane C — retarget a SKEL-24 BVH onto the character_kit mixamorig rig + QC render.

Headless Blender only:
  /Applications/Blender.app/Contents/MacOS/Blender --background --python retarget.py -- \
      --bvh out/<name>/motion.bvh [--rig <female_game.glb>] [--name <n>] [--travel] [--no-qc]

Pipeline: import mixamorig rig (female_game.glb) + the SKEL BVH → Copy-Rotation-map SKEL bones
onto mixamorig (spec/skel_to_mixamorig.json), world-space → bake to keyframes → trim near-rest
bookends → export anim-only FBX (Mixamo "Without Skin", merge_anims.py eats these) → Workbench
turntable QC mp4. Root is in-place by default (HSMR root is camera-space; --travel copies it).

# ponytail: constraint-bake retarget (world Copy Rotation). If the QC shows arm-roll twist
# (SKEL vs mixamorig rest-roll differ), upgrade the ONE function `map_and_bake` to a rest-delta
# transfer (Rt_f = Rs_f · Rs_rest⁻¹ · Rt_rest); everything else stays.
"""
import bpy, json, sys, os
from pathlib import Path
from mathutils import Matrix

REPO = Path(__file__).resolve().parent
DEF_RIG = Path.home() / "Documents" / "character_kit" / "female" / "female_game.glb"
MAP = REPO / "spec" / "skel_to_mixamorig.json"


def args_after_ddash():
    a = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--bvh", required=True)
    p.add_argument("--rig", default=str(DEF_RIG))
    p.add_argument("--name", default=None)
    p.add_argument("--out", default=str(REPO / "clips_out"))
    p.add_argument("--travel", action="store_true", help="copy root translation (default: in-place)")
    p.add_argument("--foot-fix", action="store_true", help="clamp planted feet to the floor (heuristic)")
    p.add_argument("--no-qc", action="store_true")
    return p.parse_args(a)


def clean():
    bpy.ops.wm.read_factory_settings(use_empty=True)


def imported(op, **kw):
    before = set(bpy.data.objects)
    op(**kw)
    return [o for o in bpy.data.objects if o not in before]


def find_bone(arm, name):
    """Tolerate mixamorig: colon vs _ underscore variants (the PLAN's warned trap)."""
    pbs = arm.pose.bones
    if name in pbs:
        return pbs[name]
    alt = name.replace(":", "_") if ":" in name else name.replace("mixamorig_", "mixamorig:")
    return pbs.get(alt)


def action_fcurves(action):
    """Blender 5.0 removed Action.fcurves for the layered ('Baklava') system — collect fcurves
    from layers/strips/channelbags, with a fallback to the legacy flat list."""
    if hasattr(action, "fcurves"):
        return list(action.fcurves)
    fcs = []
    for layer in getattr(action, "layers", []):
        for strip in layer.strips:
            for cb in getattr(strip, "channelbags", []):
                fcs.extend(cb.fcurves)
    return fcs


def map_and_bake(target, source, name_map, f0, f1, travel):
    """Add world-space Copy Rotation constraints target←source per the map, then bake to keys."""
    bpy.context.view_layer.objects.active = target
    n = 0
    for skel_bone, mixamo in name_map.items():
        if not mixamo:                                   # null = twist/heel, folded elsewhere
            continue
        pb = find_bone(target, mixamo)
        if pb is None or skel_bone not in source.pose.bones:
            print(f"  skip {skel_bone}->{mixamo} (missing)")
            continue
        c = pb.constraints.new("COPY_ROTATION")
        c.target, c.subtarget = source, skel_bone
        c.target_space = c.owner_space = "WORLD"
        c.mix_mode = "REPLACE"
        n += 1
    if travel:                                           # root position (units differ; scale it)
        hips = find_bone(target, name_map["pelvis"])
        c = hips.constraints.new("COPY_LOCATION")
        c.target, c.subtarget = source, "pelvis"
        c.target_space = c.owner_space = "WORLD"
        c.use_offset = True
    print(f"mapped {n} bones; baking frames {f0}..{f1}")

    bpy.ops.object.mode_set(mode="POSE")
    bpy.ops.pose.select_all(action="SELECT")
    bpy.ops.nla.bake(frame_start=f0, frame_end=f1, only_selected=True, visual_keying=True,
                     clear_constraints=True, use_current_action=True, bake_types={"POSE"})
    bpy.ops.object.mode_set(mode="OBJECT")


def trim_bookends(arm, f0, f1, thresh_deg=6.0):
    """Trim leading/trailing frames that sit near the rest pose (T-pose bookends from CAPTURE.md).
    Tolerant of absence (demo clips have none) — returns the untouched range if nothing qualifies.
    # ponytail: rest-deviation heuristic on the action's keyed rotations; good enough for trimming."""
    import math
    act = arm.animation_data.action if arm.animation_data else None
    if not act:
        return f0, f1
    try:
        # mean absolute keyed rotation magnitude per frame (quaternion channels), vs rest
        rot_fcurves = [fc for fc in action_fcurves(act)
                       if fc.data_path.endswith("rotation_quaternion")]
    except Exception as e:
        print(f"trim: fcurve access failed ({e!r}); keeping full range")
        return f0, f1
    if not rot_fcurves:
        return f0, f1

    def near_rest(fr):
        # quaternion identity is (1,0,0,0); deviation ~ how far w drifts from 1 across bones
        dev = 0.0
        for fc in rot_fcurves:
            if fc.array_index == 0:                       # w component
                dev += abs(1.0 - fc.evaluate(fr))
        return math.degrees(dev) < thresh_deg

    a, b = f0, f1
    while a < b and near_rest(a):
        a += 1
    while b > a and near_rest(b):
        b -= 1
    if (a, b) != (f0, f1):
        print(f"trimmed bookends: {f0}..{f1} -> {a}..{b}")
    return a, b


def foot_fix(arm, f0, f1):
    """# ponytail: minimal stub — real IK foot-pinning is Lane-C-plus. Logs planted frames so the
    upgrade has a target; does not yet re-solve. Enable heavy version when sliding annoys."""
    print("foot-fix: heuristic stub (no IK yet) — feet may slide; upgrade path noted in PLAN")


def export_fbx(arm, path):
    bpy.ops.object.select_all(action="DESELECT")
    arm.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.export_scene.fbx(
        filepath=str(path), use_selection=True, object_types={"ARMATURE"},
        add_leaf_bones=False, bake_anim=True, bake_anim_use_all_bones=True,
        bake_anim_use_nla_strips=False, bake_anim_use_all_actions=False,
        armature_nodetype="NULL", path_mode="COPY",
    )
    print(f"exported anim-only FBX -> {path} ({os.path.getsize(path)//1024}KB)")


def qc_render(target_arm, mesh_objs, f0, f1, path, fps):
    """Workbench turntable: camera orbits the character while the clip plays. Solid-shaded =
    fast + reads the pose clearly (the QC eyeball the PLAN mandates before a clip enters the bank)."""
    import math
    from mathutils import Vector
    # auto-frame from the character's world bbox (mixamo rigs import at odd scales — this one is
    # ~156 units tall, centred at z≈77). Scale-agnostic so any rig frames correctly.
    cs = []
    for m in mesh_objs:
        for c in m.bound_box:
            cs.append(m.matrix_world @ Vector(c))
    mn = Vector((min(c.x for c in cs), min(c.y for c in cs), min(c.z for c in cs)))
    mx = Vector((max(c.x for c in cs), max(c.y for c in cs), max(c.z for c in cs)))
    center = (mn + mx) / 2
    size = max((mx - mn).x, (mx - mn).y, (mx - mn).z) or 1.0
    dist = size * 1.9

    piv = bpy.data.objects.new("qc_pivot", None)
    piv.location = center
    bpy.context.scene.collection.objects.link(piv)

    cam_data = bpy.data.cameras.new("qc_cam")
    cam_data.clip_start, cam_data.clip_end = size * 0.01, dist * 4
    cam = bpy.data.objects.new("qc_cam", cam_data)
    bpy.context.scene.collection.objects.link(cam)
    cam.parent = piv
    cam.matrix_parent_inverse = Matrix()
    cam.location = (0, -dist, 0)                           # local to pivot → orbits at eye level
    tc = cam.constraints.new("TRACK_TO")
    tc.target = piv
    bpy.context.scene.camera = cam

    piv.rotation_euler = (0, 0, 0)
    piv.keyframe_insert("rotation_euler", frame=f0)
    piv.rotation_euler = (0, 0, math.radians(360))
    piv.keyframe_insert("rotation_euler", frame=f1)
    for fc in action_fcurves(piv.animation_data.action):  # linear spin, no ease
        for kp in fc.keyframe_points:
            kp.interpolation = "LINEAR"

    import subprocess, shutil, tempfile
    sc = bpy.context.scene
    sc.render.engine = "BLENDER_WORKBENCH"
    sc.frame_start, sc.frame_end = f0, f1
    sc.render.resolution_x = sc.render.resolution_y = 640
    # Blender 5.0 dropped FFMPEG image output → render PNG frames, stitch with system ffmpeg.
    frames = Path(tempfile.mkdtemp())
    sc.render.image_settings.file_format = "PNG"
    sc.render.filepath = str(frames / "f_")
    bpy.ops.render.render(animation=True)
    ff = shutil.which("ffmpeg") or "/opt/homebrew/bin/ffmpeg"
    r = subprocess.run([ff, "-y", "-framerate", f"{fps:.3f}", "-start_number", str(f0),
                        "-i", str(frames / "f_%04d.png"), "-c:v", "libx264", "-pix_fmt", "yuv420p",
                        "-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", str(path)],
                       capture_output=True, text=True)
    if r.returncode or not Path(path).exists():
        print("QC ffmpeg stitch FAILED:", r.stderr[-400:])
    else:
        print(f"QC render -> {path} ({os.path.getsize(path)//1024}KB)")


def main():
    a = args_after_ddash()
    bvh = Path(a.bvh).resolve()
    name = a.name or bvh.parent.name
    outdir = Path(a.out); outdir.mkdir(parents=True, exist_ok=True)
    name_map = json.loads(MAP.read_text())["skel_to_mixamorig"]

    clean()
    rig_objs = imported(bpy.ops.import_scene.gltf, filepath=a.rig)
    target = next(o for o in rig_objs if o.type == "ARMATURE")
    meshes = [o for o in rig_objs if o.type == "MESH"]
    print(f"rig: armature '{target.name}', {len(target.data.bones)} bones, {len(meshes)} mesh(es)")

    src_objs = imported(bpy.ops.import_anim.bvh, filepath=str(bvh),
                        axis_forward="-Z", axis_up="Y", use_fps_scale=False, update_scene_fps=True)
    source = next(o for o in src_objs if o.type == "ARMATURE")
    act = source.animation_data.action
    f0, f1 = (int(round(x)) for x in act.frame_range)
    fps = bpy.context.scene.render.fps / max(bpy.context.scene.render.fps_base, 1e-9)
    print(f"bvh: source '{source.name}', frames {f0}..{f1} @ ~{fps:.2f}fps")

    map_and_bake(target, source, name_map, f0, f1, a.travel)
    bpy.data.objects.remove(source, do_unlink=True)        # drop the SKEL rig, keep baked mixamorig
    if a.foot_fix:
        foot_fix(target, f0, f1)
    f0, f1 = trim_bookends(target, f0, f1)

    fbx = outdir / f"{name}.fbx"
    # the FBX take inherits the action name — un-named it exports as 'Armature|Scene' for EVERY clip,
    # colliding in any take-addressed bank (Unreal/3GOD). merge_anims renames from the filename, but
    # the take itself must be right (PROGRESS.md bug, fixed 2026-07-17).
    target.animation_data.action.name = name
    export_fbx(target, fbx)
    # sanity: the FBX action must carry keys on the driven bones
    assert target.animation_data and target.animation_data.action, "no baked action on target"
    nkeys = len(action_fcurves(target.animation_data.action))
    print(f"baked action '{target.animation_data.action.name}': {nkeys} fcurves")
    assert nkeys > 20, f"too few baked channels ({nkeys}) — retarget likely failed"

    if not a.no_qc:
        qc_render(target, meshes, f0, f1, outdir / f"{name}_qc.mp4", fps)
    print(f"LANE-C DONE: {name} -> {fbx.name}" + ("" if a.no_qc else f" + {name}_qc.mp4"))


if __name__ == "__main__":
    main()