backend/blender/process_model.py

# backend/blender/process_model.py
# S4: headless Blender pass for the OMEGA model pipeline.
# Runs INSIDE Blender:  blender --background --factory-startup --python process_model.py -- <args>
# Two modes:
#   --input model.glb          → import, clean, normalize part names, apply template, export
#   --blueprint blueprint.json → procedural build from a Forge blueprint (Builder flow)
# Output is always a GLB (no draco — the WebAR GLTFLoader ships without a decoder).
#
# Kept to bpy APIs stable across 4.2 LTS → 5.x: import_scene.gltf, export_scene.gltf,
# object.transform_apply, mesh primitive ops.

import json
import math
import re
import sys

import bpy


def parse_args():
    argv = sys.argv
    argv = argv[argv.index("--") + 1:] if "--" in argv else []
    args = {"input": "", "output": "out.glb", "template": "", "blueprint": "",
            "target_size": "1.0", "name": "model"}
    key = None
    for token in argv:
        if token.startswith("--"):
            key = token[2:].replace("-", "_")
            if key not in args:
                args[key] = ""
        elif key:
            args[key] = token
            key = None
    return args


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


def slug(value):
    return re.sub(r"^_+|_+$", "", re.sub(r"[^a-z0-9]+", "_", str(value or "").lower()))


def normalize_part_names():
    """Lowercase, underscore-separated object names so Forge template matching
    (client findMatchingBlueprintTemplate / server match_template) hits reliably."""
    renamed = {}
    for obj in bpy.data.objects:
        clean = slug(obj.name)
        if clean and clean != obj.name:
            renamed[obj.name] = clean
            obj.name = clean
    return renamed


def cleanup_meshes():
    """Conservative cleanup: recalc outside normals + drop loose geometry.
    Deliberately NO merge-by-distance / decimate — destructive on UV seams."""
    for obj in bpy.data.objects:
        if obj.type != "MESH":
            continue
        bpy.context.view_layer.objects.active = obj
        obj.select_set(True)
        bpy.ops.object.mode_set(mode="EDIT")
        bpy.ops.mesh.select_all(action="SELECT")
        bpy.ops.mesh.normals_make_consistent(inside=False)
        bpy.ops.mesh.delete_loose()
        bpy.ops.object.mode_set(mode="OBJECT")
        obj.select_set(False)


def apply_template(template):
    """Stamp Forge interactivity metadata as glTF extras (custom properties).
    The WebAR client's discoverPartsFromObject reads userData.partType/partActions."""
    if not template:
        return 0
    aliases = {slug(k): v for k, v in (template.get("partAliases") or {}).items()}
    part_meta = {}
    for part in template.get("parts") or []:
        part_meta[slug(part.get("name"))] = part
    stamped = 0
    for obj in bpy.data.objects:
        key = slug(obj.name)
        if key in aliases:
            obj.name = slug(aliases[key])
            key = obj.name
        meta = part_meta.get(key)
        if meta:
            obj["partType"] = str(meta.get("type") or "detail")
            obj["partActions"] = json.dumps(meta.get("actions") or [])
            stamped += 1
    return stamped


def center_and_scale(target_size):
    """Center at origin, feet on floor (y=0 in glTF +Y-up terms), uniform scale so
    the largest dimension equals target_size — honest scale for AR placement."""
    import mathutils
    meshes = [o for o in bpy.data.objects if o.type == "MESH"]
    if not meshes:
        return
    mins = mathutils.Vector((math.inf,) * 3)
    maxs = mathutils.Vector((-math.inf,) * 3)
    for obj in meshes:
        for corner in obj.bound_box:
            world = obj.matrix_world @ mathutils.Vector(corner)
            mins = mathutils.Vector(map(min, mins, world))
            maxs = mathutils.Vector(map(max, maxs, world))
    size = max(maxs - mins)
    if size <= 0:
        return
    scale = float(target_size) / size
    center = (mins + maxs) / 2
    for obj in bpy.data.objects:
        if obj.parent is None:
            obj.location = (obj.location - center) * scale
            obj.location.z += (center.z - mins.z) * scale  # rest on floor
            obj.scale = obj.scale * scale


GEOMETRY_OPS = {
    "box": lambda: bpy.ops.mesh.primitive_cube_add(size=0.25),
    "cube": lambda: bpy.ops.mesh.primitive_cube_add(size=0.25),
    "sphere": lambda: bpy.ops.mesh.primitive_uv_sphere_add(radius=0.15),
    "cylinder": lambda: bpy.ops.mesh.primitive_cylinder_add(radius=0.1, depth=0.28),
    "cone": lambda: bpy.ops.mesh.primitive_cone_add(radius1=0.12, depth=0.26),
    "torus": lambda: bpy.ops.mesh.primitive_torus_add(major_radius=0.14, minor_radius=0.02),
    "capsule": lambda: bpy.ops.mesh.primitive_uv_sphere_add(radius=0.08),
    "octahedron": lambda: bpy.ops.mesh.primitive_ico_sphere_add(subdivisions=1, radius=0.12),
    "icosahedron": lambda: bpy.ops.mesh.primitive_ico_sphere_add(subdivisions=2, radius=0.12),
}

MATERIAL_PRESETS = {
    # (base_color RGBA, metallic, roughness, emission_strength)
    "matte": ((0.48, 0.74, 1.0, 1.0), 0.1, 0.55, 0.0),
    "metal": ((0.56, 0.78, 1.0, 1.0), 0.92, 0.16, 0.0),
    "glass": ((0.81, 0.94, 1.0, 0.36), 0.0, 0.05, 0.0),
    "glow": ((0.0, 0.78, 1.0, 1.0), 0.0, 0.4, 2.0),
    "wood": ((0.55, 0.42, 0.31, 1.0), 0.02, 0.82, 0.0),
}


def make_material(kind):
    name = f"forge_{kind}"
    mat = bpy.data.materials.get(name)
    if mat:
        return mat
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes.get("Principled BSDF")
    color, metallic, roughness, emission = MATERIAL_PRESETS.get(kind, MATERIAL_PRESETS["matte"])
    if bsdf:
        bsdf.inputs["Base Color"].default_value = color
        bsdf.inputs["Metallic"].default_value = metallic
        bsdf.inputs["Roughness"].default_value = roughness
        if emission and "Emission Strength" in bsdf.inputs:
            bsdf.inputs["Emission Strength"].default_value = emission
            bsdf.inputs["Emission Color"].default_value = color
        if color[3] < 1.0:
            mat.blend_method = "BLEND"
            bsdf.inputs["Alpha"].default_value = color[3]
    return mat


def build_from_blueprint(blueprint):
    """Procedural Builder mode — mirrors the client's model-forge part vocabulary
    so a Gemini blueprint renders identically whether built client-side or here."""
    built = 0
    for part in blueprint.get("parts") or []:
        geometry = str(part.get("geometry") or "box").lower()
        op = GEOMETRY_OPS.get(geometry, GEOMETRY_OPS["box"])
        op()
        obj = bpy.context.active_object
        obj.name = slug(part.get("name") or f"part_{built}")
        pos = part.get("position") or [0, 0, 0]
        rot = part.get("rotation") or [0, 0, 0]
        scale = part.get("scale") or [1, 1, 1]
        # Blueprint space is glTF (+Y up); Blender is +Z up → swap Y/Z.
        obj.location = (float(pos[0]), -float(pos[2]), float(pos[1]))
        obj.rotation_euler = (float(rot[0]), -float(rot[2]), float(rot[1]))
        obj.scale = (float(scale[0]), float(scale[2]), float(scale[1]))
        obj["partType"] = str(part.get("type") or "detail")
        obj["partActions"] = json.dumps(part.get("actions") or [])
        obj.data.name = obj.name  # keep mesh data-block names aligned — GLB scanners see both
        obj.data.materials.append(make_material(str(part.get("material") or "matte").lower()))
        built += 1
    return built


def export_glb(path):
    bpy.ops.export_scene.gltf(
        filepath=path,
        export_format="GLB",
        export_yup=True,
        export_extras=True,      # partType/partActions custom props → glTF extras
        export_apply=True,
        export_draco_mesh_compression_enable=False,
    )


def main():
    args = parse_args()
    clean_scene()
    report = {"mode": "", "parts": 0, "renamed": 0, "stamped": 0, "output": args["output"]}

    template = None
    if args["template"]:
        with open(args["template"], encoding="utf-8") as fh:
            template = json.load(fh)

    if args["blueprint"]:
        with open(args["blueprint"], encoding="utf-8") as fh:
            blueprint = json.load(fh)
        report["mode"] = "builder-procedural"
        report["parts"] = build_from_blueprint(blueprint)
    elif args["input"]:
        report["mode"] = "forge-process"
        bpy.ops.import_scene.gltf(filepath=args["input"])
        report["renamed"] = len(normalize_part_names())
        cleanup_meshes()
        report["parts"] = len([o for o in bpy.data.objects if o.type == "MESH"])
    else:
        print("OMEGA_BLENDER_RESULT " + json.dumps({"error": "no --input or --blueprint"}))
        sys.exit(2)

    report["stamped"] = apply_template(template)
    center_and_scale(args["target_size"])
    export_glb(args["output"])
    print("OMEGA_BLENDER_RESULT " + json.dumps(report))


if __name__ == "__main__":
    main()