# 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)")
    ap.add_argument("--torn", default="",
                    help="comma list of mesh-name substrings to tear (punch alpha "
                         "holes in their textures); output dirs get an @torn suffix")
    ap.add_argument("--torn-seed", type=int, default=7)
    ap.add_argument("--strip", default="",
                    help="comma list of mesh-name substrings to REMOVE (clothes "
                         "knocked off); output dirs get an @stripped suffix")
    ap.add_argument("--attach", default="",
                    help="gear to bone-attach, ';'-separated entries of "
                         "glb|Bone|scale|dx,dy,dz|rx,ry,rz (offsets in bone space, "
                         "rotations in degrees; '|' because mixamorig bones contain "
                         "':'). Output dirs get an @geared suffix")
    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]


## Punch jagged alpha holes into the base-color textures of matching meshes
## and switch their materials to alpha-clip — sprite-scale "shredded" look.
def tear_clothing(objs, name_filter: str, seed: int):
    import random

    import numpy as np
    rng = random.Random(seed)
    keys = [k.strip().lower() for k in name_filter.split(",") if k.strip()]
    done_images = set()
    for o in objs:
        if o.type != "MESH" or not any(k in o.name.lower() for k in keys):
            continue
        for slot in o.material_slots:
            mat = slot.material
            if mat is None or not mat.use_nodes:
                continue
            mat.blend_method = "CLIP"
            for node in mat.node_tree.nodes:
                if node.type != "TEX_IMAGE" or node.image is None:
                    continue
                img = node.image
                if img.name in done_images:
                    continue
                done_images.add(img.name)
                w, h = img.size
                if w == 0:
                    continue
                px = np.empty(w * h * 4, dtype=np.float32)
                img.pixels.foreach_get(px)
                px = px.reshape(h, w, 4)
                yy, xx = np.mgrid[0:h, 0:w]
                for _ in range(14):                    # jagged elliptical rips
                    cx, cy = rng.uniform(0, w), rng.uniform(0, h)
                    rx = rng.uniform(0.02, 0.07) * w
                    ry = rng.uniform(0.04, 0.12) * h
                    ang = rng.uniform(0, math.pi)
                    dx, dy = xx - cx, yy - cy
                    u = (dx * math.cos(ang) + dy * math.sin(ang)) / rx
                    v = (-dx * math.sin(ang) + dy * math.cos(ang)) / ry
                    d = u * u + v * v
                    noise = 0.55 + 0.45 * np.sin(u * 9.0) * np.cos(v * 7.0)
                    px[..., 3] = np.where(d < noise, 0.0, px[..., 3])
                img.pixels.foreach_set(px.reshape(-1))
                img.update()
                print("TORE texture", img.name, "on", o.name)


## Bone-parent rigid gear (helmets, gloves) to the armature. Entry format:
## path.glb|BoneName|scale|dx,dy,dz|rx,ry,rz — offsets in bone-local space.
def attach_gear(arm, spec: str):
    import mathutils
    for entry in spec.split(";"):
        parts = entry.split("|")
        glb, bone_name = parts[0], parts[1]
        scale = float(parts[2]) if len(parts) > 2 else 0.25
        off = [float(v) for v in parts[3].split(",")] if len(parts) > 3 else [0, 0, 0]
        rot = [float(v) for v in parts[4].split(",")] if len(parts) > 4 else [0, 0, 0]
        before = set(bpy.data.objects)
        bpy.ops.import_scene.gltf(filepath=str(Path(glb).expanduser()))
        new = [o for o in bpy.data.objects if o not in before and o.type == "MESH"]
        # bone-parented children inherit the armature's object scale (0.01 on
        # Mixamo FBX imports) — compensate so `scale`/`off` stay in metres
        arm_s = arm.scale[0] if arm.scale[0] != 0 else 1.0
        for o in new:
            o.parent = arm
            o.parent_type = "BONE"
            o.parent_bone = bone_name
            k = scale / arm_s
            o.scale = (k, k, k)
            o.location = mathutils.Vector([v / arm_s for v in off])
            o.rotation_euler = [math.radians(v) for v in rot]
            print("ATTACHED", Path(glb).stem, "->", bone_name, "world_scale", scale)


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)
    variant_suffix = ""
    if args.torn:
        tear_clothing(objs, args.torn, args.torn_seed)
        variant_suffix = "@torn"
    if args.strip:
        keys = [k.strip().lower() for k in args.strip.split(",") if k.strip()]
        for o in list(objs):
            if o.type == "MESH" and any(k in o.name.lower() for k in keys):
                print("STRIPPED", o.name)
                objs.remove(o)
                bpy.data.objects.remove(o, do_unlink=True)
        variant_suffix = "@stripped"
    if args.attach:
        attach_gear(arm, args.attach)
        variant_suffix = "@geared"

    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() + variant_suffix
        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()