"""
blender_segment_body.py
Usage:
  tooling/blender --background --python tooling/blender_segment_body.py -- \
    <input.gltf> <output_dir> [--scale sx sy sz]

Segments a Quaternius FullBody GLTF into 18 production GLBs:
  seg_head, seg_neck, seg_torso, seg_torso_upper,
  seg_arm_upper_l/r, seg_arm_lower_l/r, seg_hand_l/r,
  seg_leg_upper_l/r, seg_leg_lower_l/r, seg_foot_l/r,
  seg_eyes, seg_eyebrows

Each segment contains the vertices primarily weighted to its bone group
plus 1-ring boundary overlap for seam-free deformation.

Optional --scale sx sy sz applies a vertex-level scale to the mesh before
segmentation (for fork body types: thin, heavy, child).

Outputs: <output_dir>/seg_{name}.glb (18 files total)

Design: D-160 (18 segments per body type), D-164 (Source .blends as starting point)
"""

import sys
import os
import bpy

# --- Segment → bone vertex group mappings ---
# Each segment selects vertices with weight > 0 for ANY listed bone.
# 1-ring expansion adds boundary overlap for seam-free deformation (D-160).

SEGMENT_BONES = {
    "seg_head":        ["Head"],
    "seg_neck":        ["neck_01"],
    "seg_torso":       ["pelvis", "spine_01", "spine_02", "spine_03"],
    "seg_torso_upper": ["spine_02", "spine_03"],
    "seg_arm_upper_l": ["clavicle_l", "upperarm_l"],
    "seg_arm_upper_r": ["clavicle_r", "upperarm_r"],
    "seg_arm_lower_l": ["lowerarm_l"],
    "seg_arm_lower_r": ["lowerarm_r"],
    "seg_hand_l": [
        "hand_l",
        "index_01_l",  "index_02_l",  "index_03_l",  "index_04_leaf_l",
        "middle_01_l", "middle_02_l", "middle_03_l", "middle_04_leaf_l",
        "pinky_01_l",  "pinky_02_l",  "pinky_03_l",  "pinky_04_leaf_l",
        "ring_01_l",   "ring_02_l",   "ring_03_l",   "ring_04_leaf_l",
        "thumb_01_l",  "thumb_02_l",  "thumb_03_l",  "thumb_04_leaf_l",
    ],
    "seg_hand_r": [
        "hand_r",
        "index_01_r",  "index_02_r",  "index_03_r",  "index_04_leaf_r",
        "middle_01_r", "middle_02_r", "middle_03_r", "middle_04_leaf_r",
        "pinky_01_r",  "pinky_02_r",  "pinky_03_r",  "pinky_04_leaf_r",
        "ring_01_r",   "ring_02_r",   "ring_03_r",   "ring_04_leaf_r",
        "thumb_01_r",  "thumb_02_r",  "thumb_03_r",  "thumb_04_leaf_r",
    ],
    "seg_leg_upper_l": ["thigh_l"],
    "seg_leg_upper_r": ["thigh_r"],
    "seg_leg_lower_l": ["calf_l"],
    "seg_leg_lower_r": ["calf_r"],
    "seg_foot_l": ["foot_l", "ball_l", "ball_leaf_l"],
    "seg_foot_r": ["foot_r", "ball_r", "ball_leaf_r"],
}

# Segments using a dedicated sub-object (not vertex-group-based segmentation)
OBJECT_SEGMENTS = {
    "seg_eyes":      "Eyes",
    "seg_eyebrows":  "Eyebrows",
}

# Ordered list for consistent output
SEGMENT_ORDER = [
    "seg_head", "seg_neck",
    "seg_torso", "seg_torso_upper",
    "seg_arm_upper_l", "seg_arm_upper_r",
    "seg_arm_lower_l", "seg_arm_lower_r",
    "seg_hand_l", "seg_hand_r",
    "seg_leg_upper_l", "seg_leg_upper_r",
    "seg_leg_lower_l", "seg_leg_lower_r",
    "seg_foot_l", "seg_foot_r",
    "seg_eyes", "seg_eyebrows",
]

WEIGHT_THRESHOLD = 0.01   # Minimum weight to count as "belonging" to a bone
EXPAND_RINGS     = 1      # 1-ring overlap at segment boundaries


def find_objects(scene):
    """Identify the main body mesh, eyes, eyebrows, and armature."""
    armature  = None
    body_mesh = None
    special   = {}

    all_meshes = [o for o in scene.objects if o.type == 'MESH']

    for obj in scene.objects:
        if obj.type == 'ARMATURE':
            armature = obj
        elif obj.type == 'MESH':
            name_upper = obj.name.upper()
            if 'EYES' in name_upper and 'BROW' not in name_upper:
                special['Eyes'] = obj
            elif 'BROW' in name_upper:
                special['Eyebrows'] = obj

    # Body mesh = largest mesh not in special set
    special_objs = set(special.values())
    candidates = [o for o in all_meshes if o not in special_objs]
    if candidates:
        body_mesh = max(candidates, key=lambda o: len(o.data.vertices))

    return body_mesh, special, armature


def apply_scale(mesh_obj, sx, sy, sz):
    """Scale mesh vertices in-place (mesh-local space)."""
    if sx == 1.0 and sy == 1.0 and sz == 1.0:
        return
    print(f"  Applying mesh scale: ({sx:.3f}, {sy:.3f}, {sz:.3f})")
    for v in mesh_obj.data.vertices:
        v.co.x *= sx
        v.co.y *= sy
        v.co.z *= sz
    mesh_obj.data.update()


def export_glb(objects, output_path):
    """Select the given objects and export as GLB."""
    bpy.ops.object.select_all(action='DESELECT')
    for obj in objects:
        obj.select_set(True)
    if objects:
        bpy.context.view_layer.objects.active = objects[0]

    bpy.ops.export_scene.gltf(
        filepath=output_path,
        use_selection=True,
        export_format='GLB',
        export_animations=False,
        export_yup=True,
        export_image_format='AUTO',
        export_texcoords=True,
        export_normals=True,
        export_skins=True,
        export_materials='EXPORT',
    )


def segment_by_bones(body_mesh, armature, bone_names, output_path):
    """
    Extract a segment from body_mesh based on bone weights.
    Uses BMesh API directly (reliable in headless mode without context issues).
    Includes 1-ring boundary overlap for seam-free deformation.
    Exports segment + armature.
    """
    import bmesh

    # Duplicate the body mesh
    bpy.ops.object.select_all(action='DESELECT')
    body_mesh.select_set(True)
    bpy.context.view_layer.objects.active = body_mesh
    bpy.ops.object.duplicate(linked=False)
    dup = bpy.context.active_object

    # Ensure we're in object mode
    if dup.mode != 'OBJECT':
        bpy.ops.object.mode_set(mode='OBJECT')

    # Get vertex group indices for target bones
    vg_indices = set()
    for name in bone_names:
        vg = dup.vertex_groups.get(name)
        if vg:
            vg_indices.add(vg.index)

    if not vg_indices:
        print(f"    WARNING: No vertex groups found for bones {bone_names}")

    # Build BMesh from the duplicated mesh data
    bm = bmesh.new()
    bm.from_mesh(dup.data)
    bm.verts.ensure_lookup_table()
    bm.edges.ensure_lookup_table()

    deform_layer = bm.verts.layers.deform.verify()

    # Identify vertices that belong to this segment (primary set)
    primary_set = set()
    for v in bm.verts:
        weights = v[deform_layer]
        for idx in vg_indices:
            if idx in weights and weights[idx] > WEIGHT_THRESHOLD:
                primary_set.add(v)
                break

    print(f"    Primary vertices: {len(primary_set)} / {len(bm.verts)}")

    # Expand by EXPAND_RINGS to get boundary overlap
    keep_set = set(primary_set)
    for _ in range(EXPAND_RINGS):
        boundary = set()
        for v in keep_set:
            for edge in v.link_edges:
                for other_v in edge.verts:
                    if other_v not in keep_set:
                        boundary.add(other_v)
        keep_set.update(boundary)

    print(f"    Keep set after {EXPAND_RINGS}-ring expansion: {len(keep_set)}")

    # Delete vertices NOT in keep_set
    verts_to_delete = [v for v in bm.verts if v not in keep_set]
    bmesh.ops.delete(bm, geom=verts_to_delete, context='VERTS')

    # Write back to mesh
    bm.to_mesh(dup.data)
    bm.free()
    dup.data.update()

    remaining = len(dup.data.vertices)
    print(f"    Segment vertices after trim: {remaining}")

    # Export segment + armature
    bpy.ops.object.select_all(action='DESELECT')
    dup.select_set(True)
    if armature:
        armature.select_set(True)
    bpy.context.view_layer.objects.active = dup

    export_glb([dup] + ([armature] if armature else []), output_path)

    size = os.path.getsize(output_path)
    print(f"    Exported: {os.path.basename(output_path)} ({size:,} bytes)")

    # Clean up duplicate
    bpy.ops.object.select_all(action='DESELECT')
    dup.select_set(True)
    bpy.ops.object.delete()


def segment_special_object(special_obj, armature, output_path):
    """Export a special sub-object (Eyes / Eyebrows) as its own GLB segment."""
    export_objs = [special_obj]
    if armature:
        export_objs.append(armature)
    export_glb(export_objs, output_path)
    size = os.path.getsize(output_path)
    print(f"    Exported: {os.path.basename(output_path)} ({size:,} bytes)")


def segment_body(gltf_path, output_dir, scale=(1.0, 1.0, 1.0)):
    """Main entry: load GLTF, apply optional scale, produce all 18 segment GLBs."""
    print(f"\n  Loading: {os.path.basename(gltf_path)}")
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=gltf_path)

    body_mesh, special, armature = find_objects(bpy.context.scene)

    if body_mesh is None:
        print("ERROR: Could not find main body mesh")
        sys.exit(1)

    print(f"  Body mesh: {body_mesh.name!r} ({len(body_mesh.data.vertices)} verts)")
    print(f"  Special:   {list(special.keys())}")
    print(f"  Armature:  {armature.name if armature else 'NONE'}")

    # Remove utility objects (Icospheres, rig widgets, empties) that are not
    # the body mesh, eyes, eyebrows, or armature. These can be children of the
    # armature and would appear in all exported GLBs otherwise.
    # Using bpy.data.objects.remove() (Python API) instead of ops — operators
    # have context issues in headless Blender.
    keepers = set(filter(None, [body_mesh, armature] + list(special.values())))
    utility_objs = [
        obj for obj in list(bpy.context.scene.objects)
        if obj not in keepers
    ]
    removed = 0
    for obj in utility_objs:
        # Clear parent relationship BEFORE removal so the armature stops
        # treating it as a hierarchy child during GLTF export
        if obj.parent is not None:
            obj.parent = None
        mesh_data = obj.data if obj.type == 'MESH' else None
        bpy.data.objects.remove(obj, do_unlink=True)
        if mesh_data and mesh_data.users == 0:
            bpy.data.meshes.remove(mesh_data)
        removed += 1
    if removed:
        remaining_names = [o.name for o in bpy.context.scene.objects]
        print(f"  Removed {removed} utility objects. Scene now: {remaining_names}")

    # Apply optional scale transform (for fork body types)
    sx, sy, sz = scale
    apply_scale(body_mesh, sx, sy, sz)
    # Also scale Eyes/Eyebrows for child (uniform scale applies to all meshes)
    if scale != (1.0, 1.0, 1.0):
        for obj in special.values():
            apply_scale(obj, sx, sy, sz)

    os.makedirs(output_dir, exist_ok=True)

    # Strip animation data
    for obj in bpy.context.scene.objects:
        if obj.animation_data:
            obj.animation_data_clear()
    for action in list(bpy.data.actions):
        bpy.data.actions.remove(action)

    exported = []
    skipped  = []

    for seg_name in SEGMENT_ORDER:
        output_path = os.path.join(output_dir, seg_name + ".glb")
        print(f"\n  [{seg_name}]")

        if seg_name in OBJECT_SEGMENTS:
            obj_key = OBJECT_SEGMENTS[seg_name]
            special_obj = special.get(obj_key)
            if special_obj is None:
                print(f"    SKIP: no {obj_key!r} object found in scene")
                skipped.append(seg_name)
                continue
            segment_special_object(special_obj, armature, output_path)
            exported.append(seg_name)

        elif seg_name in SEGMENT_BONES:
            bone_names = SEGMENT_BONES[seg_name]
            segment_by_bones(body_mesh, armature, bone_names, output_path)
            exported.append(seg_name)

        else:
            print(f"    SKIP: unknown segment {seg_name!r}")
            skipped.append(seg_name)

    return exported, skipped


if __name__ == "__main__":
    argv = sys.argv
    if "--" not in argv:
        print("Usage: tooling/blender --background --python tooling/blender_segment_body.py -- <input.gltf> <output_dir> [--scale sx sy sz]")
        sys.exit(1)

    args = argv[argv.index("--") + 1:]
    if len(args) < 2:
        print("ERROR: Provide <input.gltf> and <output_dir>")
        sys.exit(1)

    gltf_path  = args[0]
    output_dir = args[1]

    # Optional scale argument
    scale = (1.0, 1.0, 1.0)
    if "--scale" in args:
        idx = args.index("--scale")
        try:
            scale = (float(args[idx + 1]), float(args[idx + 2]), float(args[idx + 3]))
        except (IndexError, ValueError):
            print("ERROR: --scale requires three floats: sx sy sz")
            sys.exit(1)

    exported, skipped = segment_body(gltf_path, output_dir, scale)

    print(f"\n=== Segmentation complete: {len(exported)} segments, {len(skipped)} skipped ===")
    for seg in exported:
        path = os.path.join(output_dir, seg + ".glb")
        print(f"  {seg}.glb  ({os.path.getsize(path):,} bytes)")
    if skipped:
        print(f"  Skipped: {skipped}")