"""
blender_batch_fit_skinned.py  (T-1089 gap G1 — the biggest wardrobe gap)

Fit a reference clothing GLB (authored on average_m) to every body type and
export ANIMATABLE skinned variants. This supersedes
tooling/blender_surface_deform_batch.py, whose output used export_skins=False
(:142) — those variants cannot animate on the shared skeleton, which is how the
runtime loads clothing (character_visual.gd:582-594).

It merges three proven codepaths that had never been combined:
  * the 11-body loop + headless temp_override Surface-Deform bind + Shrinkwrap
    fallback + pipeline_log.json   — from blender_surface_deform_batch.py
  * the weight flow                — from
      spikes/quaternius-aesthetic/scripts/blender/fit_outfits_to_bodies.py:160-257
      Surface Deform bind -> apply -> Data Transfer VGROUP_WEIGHTS
      (POLYINTERP_NEAREST) from the fitted body -> normalize -> retarget the
      armature modifier to the body's armature -> export_skins=True
  * optional Solidify              — from tooling/convert_outfit.py (skipped by
      default; offset-shell references are already solidified)

Per body type:
  average_m (REFERENCE_BODY): direct copy of the reference GLB (already correct).
  others: SD-bind the reference garment to the target body surface, bake the
  deformed rest shape, transfer + normalise weights from that body, retarget the
  armature, export a skinned GLB.

Run:
  tooling/blender --background --python \
    tooling/garment-fit/blender_batch_fit_skinned.py -- \
    <reference_glb> <bodies_dir> <output_dir> \
    [--only average_m,average_f,...] [--solidify 0.0] [--self-check]

Output:
  <output_dir>/<body_type>.glb      one skinned variant per fitted body type
  <output_dir>/pipeline_log.json    per-variant method (surface_deform/shrinkwrap/
                                     copy) + status, for the review gate

Decisions: D-162 (clothing pre-fitted per body type), D-251 (in-house wardrobe).
"""

import sys
import os
import json
import glob
import shutil
import bpy

BODY_TYPES = [
    "thin_m", "thin_f", "average_m", "average_f", "muscular_m", "muscular_f",
    "teen_m", "teen_f", "heavy_m", "heavy_f", "child",
]
REFERENCE_BODY = "average_m"
SHRINKWRAP_OFFSET = 0.002
SD_FALLOFF = 4.0   # generous — clothing sits proud of the body (fit_outfits :212)


def log(msg):
    print(f"[batch-fit] {msg}")


# --------------------------------------------------------------------------
# Scene helpers
# --------------------------------------------------------------------------

def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    bpy.ops.outliner.orphans_purge(do_recursive=True)


def import_glb(path):
    before = set(bpy.context.scene.objects)
    bpy.ops.import_scene.gltf(filepath=path)
    return [o for o in bpy.context.scene.objects if o not in before]


def is_body_mesh(obj):
    """Skinned body-segment mesh — excludes Icosphere debris that rides in GLBs."""
    return (
        obj.type == 'MESH'
        and not obj.name.startswith("Icosphere")
        and len(obj.vertex_groups) > 0
        and len(obj.data.vertices) >= 50
    )


def deselect_all():
    bpy.ops.object.select_all(action='DESELECT')


def set_active(obj):
    bpy.context.view_layer.objects.active = obj


# --------------------------------------------------------------------------
# Body surface (join of skinned segments) + one target armature
# --------------------------------------------------------------------------

def build_body(bodies_dir, body_type):
    """Import all seg_*.glb for a body; return (body_surface_mesh, armature).

    The surface is a join of the skinned segment meshes (keeps merged vertex
    groups so it can serve as both the Surface-Deform target and the weight
    source). One armature is kept as the retarget destination; extras dropped.
    """
    body_dir = os.path.join(bodies_dir, body_type)
    if not os.path.isdir(body_dir):
        return None, None

    seg_paths = sorted(glob.glob(os.path.join(body_dir, "seg_*.glb")))
    meshes = []
    armature = None
    for p in seg_paths:
        for o in import_glb(p):
            if o.type == 'ARMATURE' and armature is None:
                armature = o
            elif o.type == 'ARMATURE':
                bpy.data.objects.remove(o, do_unlink=True)
            elif is_body_mesh(o):
                meshes.append(o)
            elif o.type == 'MESH':
                bpy.data.objects.remove(o, do_unlink=True)

    if not meshes or armature is None:
        return None, None

    deselect_all()
    for m in meshes:
        m.select_set(True)
    set_active(meshes[0])
    bpy.ops.object.join()
    surface = bpy.context.active_object
    surface.name = f"body_surface_{body_type}"
    # Detach from armature so it is pure geometry for SD/weight source.
    for mod in list(surface.modifiers):
        if mod.type == 'ARMATURE':
            surface.modifiers.remove(mod)
    return surface, armature


# --------------------------------------------------------------------------
# Reference garment
# --------------------------------------------------------------------------

def import_reference_garment(reference_glb):
    """Import the reference garment; return its single skinned mesh + armature."""
    objs = import_glb(reference_glb)
    meshes = [o for o in objs if is_body_mesh(o)]
    if not meshes:
        # offset-shell garments always have vgroups; guard anyway
        meshes = [o for o in objs if o.type == 'MESH' and len(o.data.vertices) >= 20]
    armature = next((o for o in objs if o.type == 'ARMATURE'), None)
    if len(meshes) > 1:
        deselect_all()
        for m in meshes:
            m.select_set(True)
        set_active(meshes[0])
        bpy.ops.object.join()
        return bpy.context.active_object, armature
    return (meshes[0] if meshes else None), armature


# --------------------------------------------------------------------------
# Fitting (Surface Deform + Shrinkwrap fallback)
# --------------------------------------------------------------------------

def try_surface_deform(clothing, body_surface):
    """Bind + apply Surface Deform via a headless context override. Returns bool."""
    clothing.select_set(True)
    set_active(clothing)
    mod = clothing.modifiers.new("SurfaceDeformFit", 'SURFACE_DEFORM')
    mod.target = body_surface
    mod.falloff = SD_FALLOFF
    mod_name = mod.name
    try:
        with bpy.context.temp_override(
            active_object=clothing, object=clothing, selected_objects=[clothing]
        ):
            bpy.ops.object.surfacedeform_bind(modifier=mod_name)
    except Exception as exc:
        log(f"  SD bind EXCEPTION: {exc}")
        clothing.modifiers.remove(mod)
        return False
    if not mod.is_bound:
        log("  SD bind did not complete (is_bound=False)")
        clothing.modifiers.remove(mod)
        return False
    deselect_all()
    clothing.select_set(True)
    set_active(clothing)
    bpy.ops.object.modifier_apply(modifier=mod_name)
    return True


def shrinkwrap_fallback(clothing, body_surface):
    log("  Shrinkwrap fallback (SD bind failed)")
    sw = clothing.modifiers.new("ShrinkwrapFit", 'SHRINKWRAP')
    sw.target = body_surface
    sw.wrap_method = 'NEAREST_SURFACEPOINT'
    sw.wrap_mode = 'ON_SURFACE'
    sw.offset = SHRINKWRAP_OFFSET
    deselect_all()
    clothing.select_set(True)
    set_active(clothing)
    bpy.ops.object.modifier_apply(modifier=sw.name)


def transfer_weights(clothing, body_surface):
    """Data Transfer VGROUP_WEIGHTS from the fitted body (fit_outfits :160-186)."""
    deselect_all()
    set_active(clothing)
    clothing.select_set(True)
    dt = clothing.modifiers.new("WeightTransfer", 'DATA_TRANSFER')
    dt.object = body_surface
    dt.use_vert_data = True
    dt.data_types_verts = {'VGROUP_WEIGHTS'}
    dt.vert_mapping = 'POLYINTERP_NEAREST'
    dt.layers_vgroup_select_src = 'ALL'
    dt.layers_vgroup_select_dst = 'NAME'
    bpy.ops.object.datalayout_transfer(modifier=dt.name)
    bpy.ops.object.modifier_apply(modifier=dt.name)


def normalize_weights(clothing):
    deselect_all()
    set_active(clothing)
    clothing.select_set(True)
    bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
    bpy.ops.object.vertex_group_normalize_all(lock_active=False)
    bpy.ops.object.mode_set(mode='OBJECT')


def retarget_armature(clothing, new_armature):
    has_arm = False
    for mod in clothing.modifiers:
        if mod.type == 'ARMATURE':
            mod.object = new_armature
            has_arm = True
    if not has_arm:
        mod = clothing.modifiers.new("Armature", 'ARMATURE')
        mod.object = new_armature
    clothing.parent = new_armature
    clothing.matrix_parent_inverse = new_armature.matrix_world.inverted()


def apply_solidify(clothing, thickness):
    if thickness <= 0.0:
        return
    deselect_all()
    clothing.select_set(True)
    set_active(clothing)
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.normals_make_consistent(inside=False)
    bpy.ops.object.mode_set(mode='OBJECT')
    sol = clothing.modifiers.new("Solidify", 'SOLIDIFY')
    sol.thickness = thickness
    sol.offset = 1.0
    sol.use_rim = True
    bpy.ops.object.modifier_apply(modifier=sol.name)


def export_variant(clothing, armature, out_path):
    deselect_all()
    clothing.select_set(True)
    armature.select_set(True)
    set_active(armature)
    bpy.ops.export_scene.gltf(
        filepath=out_path,
        export_format='GLB',
        use_selection=True,
        export_apply=False,      # keep armature modifier for skinning
        export_animations=False,
        export_skins=True,       # <-- the whole point of G1
        export_yup=True,
        export_texcoords=True,
        export_normals=True,
        export_materials='EXPORT',
        export_image_format='AUTO',
    )


# --------------------------------------------------------------------------
# Per-body processing
# --------------------------------------------------------------------------

def process_body(body_type, reference_glb, bodies_dir, output_dir, solidify_mm):
    out_path = os.path.join(output_dir, f"{body_type}.glb")
    log(f"=== {body_type} ===")

    if body_type == REFERENCE_BODY:
        shutil.copy2(reference_glb, out_path)
        log("  reference body — direct copy")
        return {"body_type": body_type, "status": "ok", "method": "copy"}

    clear_scene()
    clothing, _ref_arm = import_reference_garment(reference_glb)
    if clothing is None:
        return {"body_type": body_type, "status": "error", "error": "no garment mesh"}

    body_surface, armature = build_body(bodies_dir, body_type)
    if body_surface is None:
        return {"body_type": body_type, "status": "error", "error": "no body surface"}

    if try_surface_deform(clothing, body_surface):
        method = "surface_deform"
    else:
        shrinkwrap_fallback(clothing, body_surface)
        method = "shrinkwrap"

    transfer_weights(clothing, body_surface)
    normalize_weights(clothing)
    retarget_armature(clothing, armature)
    apply_solidify(clothing, solidify_mm)

    # Drop the body surface so only garment + armature export.
    bpy.data.objects.remove(body_surface, do_unlink=True)
    export_variant(clothing, armature, out_path)
    log(f"  exported [{method}] -> {os.path.basename(out_path)}")
    return {"body_type": body_type, "status": "ok", "method": method}


# --------------------------------------------------------------------------
# Self-check (headless smoke test): refit peasant_tunic to average_f
# --------------------------------------------------------------------------

def self_check(bodies_dir):
    """Prove the SD bind + weight flow + skinned export path runs headless.

    Builds average_f's body surface, binds a trivial one-quad plane to it, and
    confirms the export produces JOINTS_0/WEIGHTS_0 accessors. Non-fatal probe.
    """
    import struct
    clear_scene()
    bpy.ops.mesh.primitive_plane_add(size=0.3, location=(0, 0.1, 1.2))
    plane = bpy.context.active_object
    body_surface, armature = build_body(bodies_dir, "average_f")
    if body_surface is None:
        log("self-check: no average_f body — SKIP")
        return
    ok = try_surface_deform(plane, body_surface)
    log(f"self-check: SD bind {'ok' if ok else 'fell back'}")
    if not ok:
        shrinkwrap_fallback(plane, body_surface)
    transfer_weights(plane, body_surface)
    normalize_weights(plane)
    retarget_armature(plane, armature)
    bpy.data.objects.remove(body_surface, do_unlink=True)
    out = os.path.join(bpy.app.tempdir, "selfcheck.glb")
    export_variant(plane, armature, out)
    with open(out, 'rb') as f:
        f.read(12)
        clen = struct.unpack('<I', f.read(4))[0]
        f.read(4)
        j = json.loads(f.read(clen))
    attrs = set()
    for m in j.get("meshes", []):
        for pr in m["primitives"]:
            attrs |= set(pr["attributes"].keys())
    has_skin = "JOINTS_0" in attrs and "WEIGHTS_0" in attrs
    log(f"self-check: exported attrs={sorted(attrs)} skinned={has_skin}")
    log(f"self-check: {'PASS' if has_skin else 'FAIL'}")


# --------------------------------------------------------------------------
# Entry
# --------------------------------------------------------------------------

def main():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []

    if "--self-check" in argv:
        # <bodies_dir> is the first positional in self-check mode
        pos = [a for a in argv if not a.startswith("--")]
        bodies_dir = pos[0] if pos else "client/assets/characters/bodies"
        self_check(bodies_dir)
        return

    if len(argv) < 3:
        print("Usage: -- <reference_glb> <bodies_dir> <output_dir> "
              "[--only a,b,c] [--solidify MM]")
        sys.exit(1)
    reference_glb, bodies_dir, output_dir = argv[0], argv[1], argv[2]

    only = None
    if "--only" in argv:
        only = [s.strip() for s in argv[argv.index("--only") + 1].split(",")]
    solidify_mm = 0.0
    if "--solidify" in argv:
        solidify_mm = float(argv[argv.index("--solidify") + 1])

    if not os.path.isfile(reference_glb):
        print(f"ERROR: reference not found: {reference_glb}")
        sys.exit(1)
    os.makedirs(output_dir, exist_ok=True)

    types = only if only else BODY_TYPES
    log(f"reference={reference_glb}")
    log(f"bodies={bodies_dir}")
    log(f"output={output_dir}")
    log(f"types={types} solidify={solidify_mm*1000:.0f}mm")

    results = []
    for bt in types:
        if bt != REFERENCE_BODY and not os.path.isdir(os.path.join(bodies_dir, bt)):
            results.append({"body_type": bt, "status": "skipped",
                            "error": "body dir missing"})
            continue
        try:
            results.append(process_body(bt, reference_glb, bodies_dir,
                                        output_dir, solidify_mm))
        except Exception as exc:
            log(f"  ERROR {bt}: {exc}")
            results.append({"body_type": bt, "status": "error", "error": str(exc)})

    ok = [r for r in results if r["status"] == "ok"]
    sd = len([r for r in ok if r.get("method") == "surface_deform"])
    sw = len([r for r in ok if r.get("method") == "shrinkwrap"])
    cp = len([r for r in ok if r.get("method") == "copy"])
    log("=" * 50)
    for r in results:
        log(f"  {r['body_type']:12s} {r['status'].upper():8s} "
            f"{r.get('method', r.get('error', ''))}")
    log(f"OK={len(ok)} (surface_deform={sd} shrinkwrap={sw} copy={cp})")
    if sw:
        log(f"WARNING: {sw} variant(s) used Shrinkwrap — verify extremities at zoom")

    with open(os.path.join(output_dir, "pipeline_log.json"), 'w') as f:
        json.dump({
            "reference": os.path.basename(reference_glb),
            "bodies_dir": bodies_dir,
            "variants": {r["body_type"]: {
                "status": r["status"],
                "method": r.get("method"),
                "error": r.get("error"),
            } for r in results},
        }, f, indent=2)
    log("wrote pipeline_log.json")

    if any(r["status"] == "error" for r in results):
        sys.exit(1)


if __name__ == "__main__":
    main()