# Blender headless: render a rigged character + Mixamo anim clips to sprite frames.
#
# Usage:
#   /Applications/Blender.app/Contents/MacOS/Blender -b --python pipeline/render_moves.py -- \
#       --char ~/Documents/foitin_assets/woman_athleisure_01.fbx \
#       --clips "clipdir_or_file1.fbx,file2.fbx" \
#       --out renders/hero [--size 512] [--rotz 90] [--samples 16] [--step 1]
#
# One output folder per clip: renders/<id>/<clip_name_lowercase>/####.png
# (that layout feeds straight into pipeline/pack_character.py).
#
# Frames render with transparent background, fixed orthographic side camera,
# ground plane at z=0 mapped to a consistent pixel row -> stable ground pivot
# across every move (the BKB .geo lesson: bake a consistent pivot).
import argparse
import math
import sys
from pathlib import Path

import bpy

GROUND_ROW_FRAC = 0.96   # ground line at 96% of image height (bottom margin)


def parse_args():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    ap = argparse.ArgumentParser()
    ap.add_argument("--char", required=True)
    ap.add_argument("--clips", required=True, help="dir or comma-separated fbx list")
    ap.add_argument("--out", required=True)
    ap.add_argument("--size", type=int, default=512)
    ap.add_argument("--rotz", type=float, default=90.0,
                    help="extra Z rotation (deg) so the character faces screen-right")
    ap.add_argument("--height", type=float, default=1.75,
                    help="normalised character height in metres")
    ap.add_argument("--cam-height", type=float, default=2.6,
                    help="world metres covered by the image height")
    ap.add_argument("--samples", type=int, default=16)
    ap.add_argument("--step", type=int, default=1, help="render every Nth frame")
    ap.add_argument("--max-frames", type=int, default=0, help="cap source frames per clip (0 = all)")
    return ap.parse_args(argv)


def import_character(path: str):
    before = set(bpy.data.objects)
    if path.lower().endswith((".glb", ".gltf")):
        bpy.ops.import_scene.gltf(filepath=path)
    else:
        bpy.ops.import_scene.fbx(filepath=path)
    new = [o for o in bpy.data.objects if o not in before]
    arms = [o for o in new if o.type == "ARMATURE"]
    if not arms:
        raise RuntimeError("no armature in " + path)
    # drop helper meshes that aren't skinned to the rig (stray spheres etc.)
    kept = []
    for o in list(new):
        if o.type == "MESH" and not any(m.type == "ARMATURE" for m in o.modifiers):
            bpy.data.objects.remove(o, do_unlink=True)
            new.remove(o)
        else:
            kept.append(o)
    return arms[0], kept


def world_bbox(objs):
    import mathutils
    pts = []
    for o in objs:
        if o.type == "MESH":
            for c in o.bound_box:
                pts.append(o.matrix_world @ mathutils.Vector(c))
    lo = [min(p[i] for p in pts) for i in range(3)]
    hi = [max(p[i] for p in pts) for i in range(3)]
    return lo, hi


## Mixamo FBX often lands at centimetre scale (100x). Normalise the rig to
## target height with feet on the ground at the origin.
def normalize_character(arm, objs, target_h: float):
    bpy.context.view_layer.update()
    lo, hi = world_bbox(objs)
    h = hi[2] - lo[2]
    if h <= 0:
        return
    s = target_h / h
    arm.scale = (arm.scale[0] * s, arm.scale[1] * s, arm.scale[2] * s)
    bpy.context.view_layer.update()
    lo, hi = world_bbox(objs)
    arm.location.x -= (lo[0] + hi[0]) / 2
    arm.location.y -= (lo[1] + hi[1]) / 2
    arm.location.z -= lo[2]
    bpy.context.view_layer.update()


## Blender 4.4+ uses layered/slotted actions; older versions expose .fcurves.
def action_fcurves(act):
    if hasattr(act, "fcurves"):
        return list(act.fcurves)
    out = []
    for layer in act.layers:
        for strip in layer.strips:
            for cb in strip.channelbags:
                out.extend(cb.fcurves)
    return out


def assign_action(arm, act):
    if arm.animation_data is None:
        arm.animation_data_create()
    arm.animation_data.action = act
    if hasattr(act, "slots") and len(act.slots):
        arm.animation_data.action_slot = act.slots[0]


def import_clip_action(path: str, target_arm):
    """Import an anim-only FBX, steal its action, delete its skeleton.
    Mixamo namespaces bones per download (mixamorig:, mixamorig6:, ...);
    remap the action's fcurve paths onto the target rig's namespace."""
    before_obj = set(bpy.data.objects)
    before_act = set(bpy.data.actions)
    bpy.ops.import_scene.fbx(filepath=path, use_anim=True)
    new_obj = [o for o in bpy.data.objects if o not in before_obj]
    new_act = [a for a in bpy.data.actions if a not in before_act]
    for o in new_obj:
        bpy.data.objects.remove(o, do_unlink=True)
    if not new_act:
        raise RuntimeError("no action in " + path)
    act = new_act[0]

    target_ns = ""
    for b in target_arm.data.bones:
        if ":" in b.name:
            target_ns = b.name.split(":")[0] + ":"
        break
    for fc in action_fcurves(act):
        if '"' not in fc.data_path:
            continue
        bone = fc.data_path.split('"')[1]
        base = bone.split(":")[-1]
        new_bone = target_ns + base
        if new_bone != bone and target_arm.data.bones.get(new_bone):
            fc.data_path = fc.data_path.replace('"%s"' % bone, '"%s"' % new_bone)
    return act


def setup_scene(size: int, samples: int):
    sc = bpy.context.scene
    try:
        sc.render.engine = "BLENDER_EEVEE_NEXT"   # Blender 4.2+
    except TypeError:
        sc.render.engine = "BLENDER_EEVEE"
    try:
        sc.eevee.taa_render_samples = samples
    except AttributeError:
        pass
    sc.render.film_transparent = True
    sc.render.resolution_x = size
    sc.render.resolution_y = size
    sc.render.image_settings.file_format = "PNG"
    sc.render.image_settings.color_mode = "RGBA"

    # key + fill + rim, angled like a simple studio setup
    for name, rot, energy in (
        ("key", (math.radians(60), 0, math.radians(35)), 3.0),
        ("fill", (math.radians(70), 0, math.radians(-50)), 1.2),
        ("rim", (math.radians(-60), 0, math.radians(180)), 1.5),
    ):
        light_data = bpy.data.lights.new(name, type="SUN")
        light_data.energy = energy
        light = bpy.data.objects.new(name, light_data)
        light.rotation_euler = rot
        sc.collection.objects.link(light)


def setup_camera(cam_height: float, size: int):
    cam_data = bpy.data.cameras.new("cam")
    cam_data.type = "ORTHO"
    cam_data.ortho_scale = cam_height          # vertical extent (square image)
    cam = bpy.data.objects.new("cam", cam_data)
    bpy.context.scene.collection.objects.link(cam)
    bpy.context.scene.camera = cam
    # side view: camera on +X axis looking toward -X, up = +Z
    cam.location = (8.0, 0.0, 0.0)
    cam.rotation_euler = (math.radians(90), 0, math.radians(90))
    # place ground (z=0) at a fixed pixel row: centre the camera so that
    # z spans [c - h/2, c + h/2]; ground at row frac r -> c = h*(r - 0.5)
    cam.location.z = cam_height * (GROUND_ROW_FRAC - 0.5)
    return cam


def main():
    args = parse_args()
    bpy.ops.wm.read_factory_settings(use_empty=True)
    setup_scene(args.size, args.samples)
    setup_camera(args.cam_height, args.size)

    arm, objs = import_character(str(Path(args.char).expanduser()))
    # rotate about WORLD Z (local axes are already tilted by the FBX importer)
    import mathutils
    arm.matrix_world = mathutils.Matrix.Rotation(math.radians(args.rotz), 4, "Z") @ arm.matrix_world
    normalize_character(arm, objs, args.height)

    clips_arg = str(Path(args.clips).expanduser())
    p = Path(clips_arg)
    if p.is_dir():
        clip_paths = sorted(p.glob("*.fbx"))
    else:
        clip_paths = [Path(c).expanduser() for c in args.clips.split(",")]

    out_root = Path(args.out).expanduser()
    sc = bpy.context.scene
    for clip in clip_paths:
        action = import_clip_action(str(clip), arm)
        assign_action(arm, action)
        f0, f1 = (int(action.frame_range[0]), int(action.frame_range[1]))
        if args.max_frames > 0:
            f1 = min(f1, f0 + args.max_frames - 1)
        # Centre EVERY frame's hips over the origin (clips travel and turn);
        # the per-frame drift is exported as root-motion data instead — the
        # engine decides per move whether to consume it (BKB .geo, modernised).
        hips = next((pb for pb in arm.pose.bones if pb.name.endswith("Hips")), None)
        base_loc = (arm.location.x, arm.location.y)
        px_per_m = args.size / args.cam_height
        move = clip.stem.lower()
        move_dir = out_root / move
        move_dir.mkdir(parents=True, exist_ok=True)
        n = 0
        root_px = []       # forward travel (screen px) per rendered frame
        prev_y = None      # -cam-space y = character forward when rotz=155ish
        for f in range(f0, f1 + 1, args.step):
            sc.frame_set(f)
            bpy.context.view_layer.update()
            if hips is not None:
                arm.location.x = base_loc[0]
                arm.location.y = base_loc[1]
                bpy.context.view_layer.update()
                hw = arm.matrix_world @ hips.head
                arm.location.x -= hw.x
                arm.location.y -= hw.y
                bpy.context.view_layer.update()
                fwd = -hw.y    # camera looks -X; screen-right = world -Y
                root_px.append(0.0 if prev_y is None else (fwd - prev_y) * px_per_m)
                prev_y = fwd
            sc.render.filepath = str(move_dir / f"{n:04d}.png")
            bpy.ops.render.render(write_still=True)
            n += 1
        if root_px:
            import json
            (move_dir / "root_motion.json").write_text(json.dumps(
                [round(v, 2) for v in root_px]))
        print(f"RENDERED {move}: {n} frames ({f0}-{f1})")

    print("DONE")


main()