# MPFB character export for the Blender -> Unreal pipeline (0.1.0 acceptance).
#
# Generates a default human, adds the standard deform rig with weights,
# materializes the v2 skin, bakes its diffuse to a texture, and exports
# FBX with the named export preset "ue_character_v1".
#
#   blender -b --factory-startup --python Tools/blender/mpfb_character_export.py -- \
#     --out /tmp/gigline/character.fbx --pose idle \
#     --save-blend Art/blender/character/character-<date>.blend

# --pose applies one of the named standing poses to the standard rig and bakes
# it into the rest pose before export, so the caller-controlled vertex
# weights still deform the exported bind correctly (pose | none | idle).

import argparse
import os
import struct
import sys
import zlib

import bpy
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import fbx_export_preset as fbx_preset  # noqa: E402

import numpy

from mathutils import Euler


def write_png(path, img):
    # Write image pixels to PNG. `img.save()` writes black for an in-memory
    # image with no backing file even though the bake updates img.pixels, and a
    # per-pixel stdlib loop is far too slow (256^2 took 25 s). numpy makes the
    # whole encode ~0.3 s for 1024^2; Blender ships numpy.
    w, h = img.size
    arr = numpy.asarray(img.pixels[:], dtype=numpy.float32).reshape(h, w, 4)
    rgba = (numpy.clip(arr, 0.0, 1.0) * 255.0 + 0.5).astype(numpy.uint8)[::-1]
    byt = rgba.tobytes()
    raw = bytearray()
    for y in range(h):
        raw.append(0)          # filter type 0
        raw += byt[y * w * 4:(y + 1) * w * 4]

    def chunk(tag, data):
        c = struct.pack(">I", len(data)) + tag + data
        c += struct.pack(">I", zlib.crc32(tag + data) & 0xffffffff)
        return c

    ihdr = struct.pack(">IIBBBBB", w, h, 8, 6, 0, 0, 0)
    png = b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr)
    png += chunk(b"IDAT", zlib.compress(bytes(raw))) + chunk(b"IEND", b"")
    with open(path, "wb") as f:
        f.write(png)

IDLE_POSE = {
    # Brings the A/T bind arms down into a relaxed standing position.
    # Values are local-pose euler deltas; verified 2026-11-07.
    "shoulder01.L": (0.0, -0.15, -0.03),
    "shoulder01.R": (0.0, 0.15, 0.03),
    "upperarm01.L": (0.0, 0.0, -1.55),
    "upperarm01.R": (0.0, 0.0, 1.55),
    "lowerarm01.L": (0.0, 0.0, -0.10),
    "lowerarm01.R": (0.0, 0.0, 0.10),
}

REPORT = "/tmp/mpfb-export-report.txt"
lines = []


def log(msg):
    lines.append(str(msg))


def enable_mpfb():
    bpy.ops.preferences.addon_enable(module="bl_ext.blender_org.mpfb")
    assert "bl_ext.blender_org.mpfb" in bpy.context.preferences.addons, "MPFB did not enable"


def apply_pose(arm, name="idle"):
    if name == "none":
        return
    table = IDLE_POSE
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.mode_set(mode="POSE")
    for bone_name, eul in table.items():
        bone = arm.pose.bones[bone_name]
        bone.rotation_mode = "XYZ"
        bone.rotation_euler = Euler(eul)
    bpy.ops.pose.armature_apply()
    bpy.ops.object.mode_set(mode="OBJECT")
    log(f"pose: applied '{name}' to rest pose")


def build_character(pose="idle"):
    for obj in list(bpy.data.objects):
        bpy.data.objects.remove(obj, do_unlink=True)

    bpy.ops.mpfb.create_human("EXEC_DEFAULT")
    human = bpy.context.active_object
    log(f"create_human: {human.name}, {len(human.data.vertices)} verts, dims={[round(d, 3) for d in human.dimensions]}")

    bpy.ops.mpfb.add_standard_rig("EXEC_DEFAULT")
    armatures = [o for o in bpy.data.objects if o.type == "ARMATURE"]
    log(f"add_standard_rig: {len(armatures)} armature(s), bones={[len(a.data.bones) for a in armatures]}")
    apply_pose(armatures[0], pose)

    bpy.ops.object.select_all(action="DESELECT")
    human.select_set(True)
    bpy.context.view_layer.objects.active = human
    bpy.ops.mpfb.create_v2_skin("EXEC_DEFAULT")
    mats = [m.name for m in human.data.materials]
    log(f"create_v2_skin: materials={mats}")
    return human


def make_diffuse(human, size=1024):
    # A Cycles bake of MPFB's v2-skin node group returns black and is slow at
    # 1024^2 headless. We ship a flat export material anyway (the node tree
    # does not map to FBX), so synthesize a warm skin-tone diffuse directly.
    img = bpy.data.images.new("T_Character_Diffuse", size, size)
    body = (0.72, 0.52, 0.38, 1.0)   # warm skin
    cloth = (0.30, 0.34, 0.40, 1.0)  # grey-blue, bottom band to suggest clothing
    band = int(size * 0.25)
    pixels = []
    for y in range(size):
        tone = cloth if y < band else body
        pixels += list(tone) * size
    img.pixels[:] = pixels
    log(f"make_diffuse: synthesized {size}x{size} body/shorts-suggestive diffuse")

    write_png("/tmp/gigline/T_Character_Diffuse.png", img)
    log("make_diffuse: wrote diffuse via stdlib PNG writer -> /tmp/gigline/T_Character_Diffuse.png")

    # MPFB's v2 skin is a node-group tree the FBX exporter cannot map to a
    # texture. Build a flat export material carrying the diffuse so the FBX
    # (and therefore Unreal) gets a real texture.
    export_mat = bpy.data.materials.new("M_Character_Export")
    export_mat.use_nodes = True
    nt = export_mat.node_tree
    nt.nodes.clear()
    out = nt.nodes.new("ShaderNodeOutputMaterial")
    bsdf = nt.nodes.new("ShaderNodeBsdfPrincipled")
    tex = nt.nodes.new("ShaderNodeTexImage")
    tex.image = img
    nt.links.new(tex.outputs["Color"], bsdf.inputs["Base Color"])
    nt.links.new(bsdf.outputs["BSDF"], out.inputs["Surface"])
    human.data.materials.clear()
    human.data.materials.append(export_mat)
    log("export material: M_Character_Export with diffuse assigned")


def export_fbx(human, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    # Named export preset "ue_character_v1" — change only via a new preset name.
    fbx_preset.export_character(path)
    log(f"export: FBX preset ue_character_v1 -> {path}")


def main():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", default="/tmp/gigline/character.fbx")
    parser.add_argument("--pose", choices=("none", "idle"), default="idle")
    parser.add_argument("--save-blend", default=None)
    args = parser.parse_args(argv)

    enable_mpfb()
    human = build_character(pose=args.pose)
    make_diffuse(human)
    if args.save_blend:
        os.makedirs(os.path.dirname(args.save_blend), exist_ok=True)
        bpy.ops.wm.save_as_mainfile(filepath=args.save_blend)
        log(f"save: {args.save_blend}")
    export_fbx(human, args.out)

    with open(REPORT, "w") as f:
        f.write("\n".join(lines) + "\n")


main()