"""
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":       ["spine_01", "spine_02"],
    "seg_torso_upper": ["spine_03", "clavicle_l", "clavicle_r"],
    "seg_hips":        ["pelvis"],
    "seg_arm_upper_l": ["upperarm_l"],
    "seg_arm_upper_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_upper", "seg_torso", "seg_hips",
    "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     = 0      # No overlap — clean segment boundaries for hiding/amputation


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 face-based assignment: each face belongs to the segment whose bones
    have the highest total weight across the face's vertices. No vertices are
    deleted — only faces that don't belong to this segment are removed.
    This keeps all boundary vertices intact (shared with neighbors) so there
    are no gaps, no holes, and no need for caps.
    """
    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

    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
    bm = bmesh.new()
    bm.from_mesh(dup.data)
    bm.verts.ensure_lookup_table()
    bm.faces.ensure_lookup_table()

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

    # Each face belongs to exactly ONE segment — the one whose bones have
    # the highest total weight across the face's vertices. This prevents
    # any face from appearing in two segments.
    #
    # We compute per-face the sum of weights for EVERY segment's bone set,
    # then assign the face to the segment with the highest sum. We only
    # keep faces assigned to THIS segment.

    # Build a map of ALL segments' bone group indices for comparison.
    # Exclude swappable variants (torso_upper) — they are subsets of their
    # parent segment and should not compete in exclusive face assignment.
    # torso_upper gets the same faces as torso, filtered to its bone subset.
    VARIANT_SEGMENTS = set()  # no variants — all segments are independent

    all_segment_vg_indices = {}
    for seg_name_key, seg_bones in SEGMENT_BONES.items():
        if seg_name_key in VARIANT_SEGMENTS:
            continue  # skip variants in competition
        seg_vg = set()
        for bname in seg_bones:
            vg = dup.vertex_groups.get(bname)
            if vg:
                seg_vg.add(vg.index)
        all_segment_vg_indices[seg_name_key] = seg_vg

    # For the current segment, use the key from SEGMENT_BONES that matches our bone_names
    current_seg_key = None
    for seg_name_key, seg_bones in SEGMENT_BONES.items():
        if set(seg_bones) == set(bone_names):
            current_seg_key = seg_name_key
            break

    if current_seg_key is None:
        # Fallback: match by vg_indices
        for seg_name_key, seg_vg in all_segment_vg_indices.items():
            if seg_vg == vg_indices:
                current_seg_key = seg_name_key
                break

    is_variant = current_seg_key in VARIANT_SEGMENTS

    keep_faces = set()
    if is_variant:
        # Variant segments (e.g. torso_upper) are subsets of a parent.
        # Keep only faces where the dominant bone (highest weight vertex)
        # is exclusively in this variant's bone set, not the parent's
        # extra bones. For torso_upper (spine_02, spine_03): keep faces
        # where spine_02/spine_03 outweigh spine_01.
        for face in bm.faces:
            variant_w = 0.0
            total_w = 0.0
            for v in face.verts:
                weights = v[deform_layer]
                for idx, w in weights.items():
                    total_w += w
                    if idx in vg_indices:
                        variant_w += w
            # Face belongs to variant if variant bones are dominant
            if total_w > 0 and variant_w / total_w > 0.5:
                keep_faces.add(face)
    else:
        # Primary segments: exclusive assignment via competition
        for face in bm.faces:
            best_seg = None
            best_weight = -1.0
            for seg_name_key, seg_vg in all_segment_vg_indices.items():
                total_w = 0.0
                for v in face.verts:
                    weights = v[deform_layer]
                    for idx in seg_vg:
                        if idx in weights:
                            total_w += weights[idx]
                if total_w > best_weight:
                    best_weight = total_w
                    best_seg = seg_name_key
            if best_seg == current_seg_key:
                keep_faces.add(face)

    print(f"    Faces to keep: {len(keep_faces)} / {len(bm.faces)}")

    # Delete faces NOT in keep_faces
    faces_to_delete = [f for f in bm.faces if f not in keep_faces]
    bmesh.ops.delete(bm, geom=faces_to_delete, context='FACES')

    # Clean up: remove vertices that have no faces left
    bm.verts.ensure_lookup_table()
    orphan_verts = [v for v in bm.verts if not v.link_faces]
    if orphan_verts:
        bmesh.ops.delete(bm, geom=orphan_verts, context='VERTS')

    # Write back
    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}")