"""Blender headless: bake per-frame blendshape weights onto an avatar and export FBX.
usage: blender -b -P bake_fbx.py -- --avatar X.vrm --frames frames.json --out out.fbx
frames.json: {"fps":30,"names":[...],"frames":[[...],...]}  (weights 0..1)
"""
import bpy, json, sys, os, argparse
argv = sys.argv[sys.argv.index("--")+1:]
ap = argparse.ArgumentParser(); ap.add_argument("--avatar", required=True); ap.add_argument("--frames", required=True); ap.add_argument("--out", required=True)
a = ap.parse_args(argv)

bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=a.avatar, import_shading="NORMALS", bone_heuristic="TEMPERANCE")
data = json.load(open(a.frames))
fps = int(data.get("fps", 30)); names = data["names"]; frames = data["frames"]
scn = bpy.context.scene
scn.render.fps = fps; scn.frame_start = 0; scn.frame_end = len(frames)-1

# Drop divider "shape keys" (--Arkit--, ----EyeSize----, ...) — they carry no geometry and their
# names break three.js FBXLoader's animation-curve classification.
import re as _re
for obj in bpy.data.objects:
    if obj.type != "MESH" or not obj.data.shape_keys: continue
    for kb in list(obj.data.shape_keys.key_blocks):
        if _re.match(r"^-+.*-+$", kb.name):
            obj.shape_key_remove(kb)

# Strip custom properties (glTF/VRM extras) from objects, meshes, materials, armatures — the FBX
# exporter turns some of them into stray AnimationCurveNodes that break three.js FBXLoader.
for coll in (bpy.data.objects, bpy.data.meshes, bpy.data.materials, bpy.data.armatures, bpy.data.shape_keys):
    for idb in coll:
        for k in list(idb.keys()):
            try: del idb[k]
            except Exception: pass
        if getattr(idb, "animation_data", None) and idb.animation_data and coll is not bpy.data.shape_keys:
            idb.animation_data_clear()

# ── Materials: VRM/MToon imports as Emission+Transparent trees, which the FBX exporter cannot
# map, so textures would be dropped (grey model). Rebuild each material as Principled BSDF with
# the base-color image (+alpha), and write packed images to disk so they can be embedded.
tex_dir = os.path.join(os.path.dirname(os.path.abspath(a.out)), "_tex")
os.makedirs(tex_dir, exist_ok=True)
saved_imgs = {}
def _save_img(img):
    if img.name in saved_imgs: return saved_imgs[img.name]
    safe = "".join(ch if ch.isalnum() or ch in "-_." else "_" for ch in img.name.strip()) or "tex"
    path = os.path.join(tex_dir, f"{safe}.png")
    try:
        img.file_format = "PNG"; img.filepath_raw = path; img.save()
    except Exception as e:
        print("[bake] image save failed", img.name, e); path = None
    saved_imgs[img.name] = path
    return path
n_fixed = 0
for mat in bpy.data.materials:
    nt = mat.node_tree
    if not mat.use_nodes or not nt: continue
    tex_nodes = [n for n in nt.nodes if n.type == "TEX_IMAGE" and n.image]
    if not tex_nodes: continue
    # prefer the texture feeding an Emission/Base Color socket, else the first
    img = None
    for n in tex_nodes:
        for l in nt.links:
            if l.from_node is n and l.to_socket.name in ("Color", "Base Color", "Emission Color", "Emission"):
                img = n.image; break
        if img: break
    img = img or tex_nodes[0].image
    was_blend = mat.blend_method in ("BLEND", "HASHED", "CLIP")
    nt.nodes.clear()
    out = nt.nodes.new("ShaderNodeOutputMaterial")
    bsdf = nt.nodes.new("ShaderNodeBsdfPrincipled")
    bsdf.inputs["Roughness"].default_value = 1.0
    for nm in ("Specular IOR Level", "Specular"):
        if nm in bsdf.inputs: bsdf.inputs[nm].default_value = 0.0
    tex = nt.nodes.new("ShaderNodeTexImage"); tex.image = img
    nt.links.new(tex.outputs["Color"], bsdf.inputs["Base Color"])
    if was_blend: nt.links.new(tex.outputs["Alpha"], bsdf.inputs["Alpha"])
    nt.links.new(bsdf.outputs["BSDF"], out.inputs["Surface"])
    _save_img(img); n_fixed += 1
print(f"[bake] rebuilt {n_fixed} materials as Principled, {len([v for v in saved_imgs.values() if v])} textures saved")

# ── Derived upper-face-only poses (same rule as the editor's addUpperFaceMorphs):
# mask = eyeMotion/(eyeMotion+mouthMotion) per vertex from the ARKit morphs, y-split fallback.
import numpy as _np
UPPER_SPLIT_POSES = ["Joy", "Angry", "Sorrow", "Fun", "Surprised", "Neutral"]
LOWER_RE = _re.compile(r"^(jaw|mouth|tongue|cheekPuff)"); UPPER_RE = _re.compile(r"^(eye|brow|cheekSquint|noseSneer)")
n_upper = 0
for obj in bpy.data.objects:
    if obj.type != "MESH" or not obj.data.shape_keys: continue
    kbs = obj.data.shape_keys.key_blocks
    poses = [n for n in UPPER_SPLIT_POSES if n in kbs and (n + "_upper") not in kbs]
    if not poses: continue
    basis = kbs[0]; N = len(basis.data)
    b = _np.empty(N * 3); basis.data.foreach_get("co", b); b = b.reshape(N, 3)
    dm = _np.zeros(N); de = _np.zeros(N)
    for kb in kbs:
        isL = bool(LOWER_RE.match(kb.name)); isU = bool(UPPER_RE.match(kb.name))
        if not (isL or isU): continue
        c = _np.empty(N * 3); kb.data.foreach_get("co", c); d = _np.linalg.norm(c.reshape(N, 3) - b, axis=1)
        if isL: dm = _np.maximum(dm, d)
        else: de = _np.maximum(de, d)
    mxE, mxM = de.max() if N else 0, dm.max() if N else 0
    selE, selM = de > 0.3 * mxE, dm > 0.3 * mxM
    ysplit = 0.5 * (b[selE, 1].mean() + b[selM, 1].mean()) if selE.any() and selM.any() else -_np.inf
    tot = de + dm
    mask = _np.where(tot > 1e-6, de / _np.where(tot > 1e-6, tot, 1), _np.where(b[:, 1] >= ysplit, 1.0, 0.0))
    mask[mask < 0.02] = 0.0
    for name in poses:
        c = _np.empty(N * 3); kbs[name].data.foreach_get("co", c); c = c.reshape(N, 3)
        nk = obj.shape_key_add(name=name + "_upper", from_mix=False)
        nk.data.foreach_set("co", (b + (c - b) * mask[:, None]).reshape(-1))
        n_upper += 1
# ── Brow-free variants of custom eye shapes (mirror of the editor's EYES_ONLY_POSES) ──
EYES_ONLY_POSES = ["eyesHappy"]
BROW_RE = _re.compile(r"^brow"); LID_RE = _re.compile(r"^(eyeBlink|eyeSquint|eyeWide|cheekSquint|eyeLook)")
n_eyes = 0
for obj in bpy.data.objects:
    if obj.type != "MESH" or not obj.data.shape_keys: continue
    kbs = obj.data.shape_keys.key_blocks
    poses = [n for n in EYES_ONLY_POSES if n in kbs and (n + "_eyes") not in kbs]
    if not poses: continue
    basis = kbs[0]; N = len(basis.data)
    b = _np.empty(N * 3); basis.data.foreach_get("co", b); b = b.reshape(N, 3)
    db = _np.zeros(N); dl = _np.zeros(N)
    for kb in kbs:
        isB = bool(BROW_RE.match(kb.name)); isL = bool(LID_RE.match(kb.name))
        if not (isB or isL): continue
        c = _np.empty(N * 3); kb.data.foreach_get("co", c); d = _np.linalg.norm(c.reshape(N, 3) - b, axis=1)
        if isB: db = _np.maximum(db, d)
        else: dl = _np.maximum(dl, d)
    tot = db + dl
    emask = _np.where(tot > 1e-6, dl / _np.where(tot > 1e-6, tot, 1), 1.0); emask[emask < 0.02] = 0.0
    for name in poses:
        c = _np.empty(N * 3); kbs[name].data.foreach_get("co", c); c = c.reshape(N, 3)
        nk = obj.shape_key_add(name=name + "_eyes", from_mix=False)
        nk.data.foreach_set("co", (b + (c - b) * emask[:, None]).reshape(-1)); n_eyes += 1
print(f"[bake] derived {n_upper} *_upper and {n_eyes} *_eyes shape keys")

targets = []  # (obj, key_block, column)
for obj in bpy.data.objects:
    if obj.type != "MESH" or not obj.data.shape_keys: continue
    kb = obj.data.shape_keys.key_blocks
    for col, n in enumerate(names):
        if n in kb: targets.append((obj, kb[n], col))
print(f"[bake] {len(frames)} frames @ {fps}fps, {len(targets)} shape-key bindings on {len({t[0].name for t in targets})} meshes")
if not targets: sys.exit("no matching shape keys")

# Write fcurves directly (much faster than keyframe_insert per frame)
by_key = {}
for obj, kb, col in targets:
    sk = obj.data.shape_keys
    if not sk.animation_data: sk.animation_data_create()
    if not sk.animation_data.action:
        sk.animation_data.action = bpy.data.actions.new(name=f"{obj.name}_face")
    act = sk.animation_data.action
    fc = act.fcurves.new(data_path=f'key_blocks["{kb.name}"].value')
    fc.keyframe_points.add(len(frames))
    flat = []
    for i, row in enumerate(frames):
        v = row[col] if col < len(row) else 0.0
        flat += [i, max(0.0, min(1.0, float(v)))]
    fc.keyframe_points.foreach_set("co", flat)
    for kp in fc.keyframe_points: kp.interpolation = "LINEAR"
    fc.update()

bpy.ops.object.select_all(action="SELECT")
bpy.ops.export_scene.fbx(filepath=a.out, use_selection=False, bake_anim=True,
    bake_anim_use_all_bones=True, bake_anim_use_nla_strips=False, bake_anim_use_all_actions=False,
    bake_anim_force_startend_keying=True, bake_anim_step=1.0, bake_anim_simplify_factor=0.0,
    add_leaf_bones=False, path_mode="COPY", embed_textures=True, apply_scale_options="FBX_SCALE_ALL",
    axis_forward="-Z", axis_up="Y", use_mesh_modifiers=False, mesh_smooth_type="OFF")
import shutil as _sh
_sh.rmtree(tex_dir, ignore_errors=True)
print("[bake] wrote", a.out, os.path.getsize(a.out))