# Garment pipeline: clothing mesh (any source) -> game-ready outfit-part GLBs
# on the shared 65-bone Quaternius skeleton. See README.md for design rationale
# (GLM-reviewed: fit-before-decimate, skirt proxy + cone weights, inner-shell
# clearance, layer stacking).
#
# Staged; every stage loads the previous checkpoint .blend from work/ and saves
# its own, so any stage can be re-run alone while iterating:
#
#   blender --background --python garment_pipeline.py -- \
#       --config configs/dress.json --stage prepare|fit|reduce|bake|skin|export
#
# Naming inside the .blend (stable contract between stages):
#   GARM_<Part>   garment part mesh (high-poly until `reduce`, low-poly after)
#   HI_<Part>     hidden high-poly copy kept for baking (created by `reduce`)
#   BODY          the target body mesh (from the game GLB, with vertex groups)
#   BODY_SHELL    inflated clearance copy of BODY (fit target)
#   RIG           the Quaternius armature imported with BODY
#   CONE_<Part>   skirt weight-transfer proxy (created by `skin`, kept for QA)

import bpy, bmesh, json, math, os, sys, argparse, random
from mathutils import Vector, Matrix

random.seed(7)

ARGS = None
CFG = None
HERE = os.path.dirname(os.path.abspath(__file__))


# ── plumbing ─────────────────────────────────────────────────────────────────

def log(msg):
    print(f"[cloth] {msg}", flush=True)


def work_path(name):
    d = os.path.join(HERE, "work", CFG["name"])
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, name)


def save_checkpoint(stage):
    p = work_path(f"{stage}.blend")
    bpy.ops.wm.save_as_mainfile(filepath=p, compress=True)
    log(f"checkpoint saved: {p}")


def load_checkpoint(stage):
    p = work_path(f"{stage}.blend")
    if not os.path.exists(p):
        raise SystemExit(f"missing checkpoint {p} — run stage '{stage}' first")
    bpy.ops.wm.open_mainfile(filepath=p)
    log(f"checkpoint loaded: {p}")


def obj(name):
    o = bpy.data.objects.get(name)
    if o is None:
        raise SystemExit(f"expected object '{name}' not in scene")
    return o


def objs_prefixed(prefix):
    return [o for o in bpy.data.objects if o.name.startswith(prefix)]


def select_only(objects, active=None):
    bpy.ops.object.select_all(action="DESELECT")
    for o in objects:
        o.select_set(True)
    bpy.context.view_layer.objects.active = active or objects[0]


def apply_all_transforms(objects):
    select_only(objects)
    bpy.ops.object.parent_clear(type="CLEAR_KEEP_TRANSFORM")
    for o in objects:
        select_only([o])
        bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)


def apply_modifier(o, mod_name):
    select_only([o])
    bpy.ops.object.modifier_apply(modifier=mod_name)


def tris_of(o):
    o.data.calc_loop_triangles()
    return len(o.data.loop_triangles)


def world_verts(o):
    return [o.matrix_world @ v.co for v in o.data.vertices]


def zspan(objects):
    vs = [v for o in objects for v in world_verts(o)]
    zs = [v.z for v in vs]
    return min(zs), max(zs)


def qa_render(tag, focus_objects=None, ortho_pad=1.25, color_type="TEXTURE"):
    """Front + 3/4 workbench renders of the current scene into work/.

    Defaults to TEXTURE shading so a QA render shows the fabric print, not just
    the placeholder colour — census passes MATERIAL to keep its island coding.
    """
    scene = bpy.context.scene
    scene.render.engine = "BLENDER_WORKBENCH"
    scene.display.shading.light = "STUDIO"
    scene.display.shading.color_type = color_type
    scene.render.resolution_x, scene.render.resolution_y = 900, 1200
    vs = [v for o in (focus_objects or [o for o in bpy.data.objects if o.type == "MESH" and o.visible_get()])
          for v in world_verts(o)]
    if not vs:
        return
    zmin, zmax = min(v.z for v in vs), max(v.z for v in vs)
    cz, h = (zmin + zmax) / 2, (zmax - zmin)
    for view, loc, rot in (("front", (0, -4, cz), (math.pi / 2, 0, 0)),
                           ("quarter", (2.7, -2.9, cz + 0.15), (math.radians(80), 0, math.radians(42)))):
        cam_name = f"_qa_{view}"
        cam = bpy.data.objects.get(cam_name)
        if cam is None:
            cam = bpy.data.objects.new(cam_name, bpy.data.cameras.new(cam_name))
            bpy.context.scene.collection.objects.link(cam)
        cam.data.type = "ORTHO"
        cam.data.ortho_scale = h * ortho_pad
        cam.location, cam.rotation_euler = loc, rot
        scene.camera = cam
        scene.render.filepath = work_path(f"qa_{tag}_{view}.png")
        bpy.ops.render.render(write_still=True)
        log(f"QA render: {scene.render.filepath}")


def bone_head_z(rig, bone):
    return (rig.matrix_world @ rig.data.bones[bone].head_local).z


# ── garment import (shared by census + prepare) ─────────────────────────────
# MD exports hide the real structure: fabric may live on ONE material while
# millions of faces are topstitch threads. So: drop configured stitch/trash
# materials, weld, split by LOOSE PARTS, and identify pieces as "islands"
# (indexed by vertex count, deterministic). Parts are then configured as island
# index lists (see census renders).

def import_garment_islands():
    pre = set(bpy.data.objects)
    src = CFG["source"]
    ext = os.path.splitext(src)[1].lower()
    if ext == ".fbx":
        bpy.ops.import_scene.fbx(filepath=src)
    elif ext in (".glb", ".gltf"):
        bpy.ops.import_scene.gltf(filepath=src)
    elif ext == ".obj":
        bpy.ops.wm.obj_import(filepath=src)
    else:
        raise SystemExit(f"unsupported garment format {ext}")
    new = [o for o in bpy.data.objects if o not in pre]
    meshes = [o for o in new if o.type == "MESH"]
    for o in [o for o in new if o.type != "MESH"]:
        bpy.data.objects.remove(o, do_unlink=True)
    for g in meshes:
        for m in list(g.modifiers):
            g.modifiers.remove(m)
        g.vertex_groups.clear()
    apply_all_transforms(meshes)

    drop = set(CFG.get("drop_materials", []))
    for g in meshes:
        if not g.data.materials:
            continue
        drop_idx = {i for i, m in enumerate(g.data.materials) if m and m.name in drop}
        if not drop_idx:
            continue
        bm = bmesh.new()
        bm.from_mesh(g.data)
        doomed = [f for f in bm.faces if f.material_index in drop_idx]
        bmesh.ops.delete(bm, geom=doomed, context="FACES")
        bm.to_mesh(g.data)
        bm.free()
    meshes = [g for g in meshes if len(g.data.vertices) > 0]
    log(f"after drop_materials: {sum(len(g.data.vertices) for g in meshes)} verts")

    # weld panels, then split to loose islands
    select_only(meshes)
    bpy.ops.object.join()
    fabric = bpy.context.view_layer.objects.active
    select_only([fabric])
    bpy.ops.object.mode_set(mode="EDIT")
    bpy.ops.mesh.select_all(action="SELECT")
    bpy.ops.mesh.remove_doubles(threshold=CFG.get("weld_threshold", 0.0006))
    bpy.ops.mesh.normals_make_consistent(inside=False)
    bpy.ops.mesh.separate(type="LOOSE")
    bpy.ops.object.mode_set(mode="OBJECT")
    islands = [o for o in bpy.data.objects if o.type == "MESH" and o.name not in ("BODY", "BODY_SHELL")]
    min_verts = CFG.get("min_island_verts", 40)
    kept = []
    for o in islands:
        if len(o.data.vertices) < min_verts:
            bpy.data.objects.remove(o, do_unlink=True)
        else:
            kept.append(o)
    # deterministic index: vertex count desc, then centroid x/z
    def key(o):
        vs = world_verts(o)
        c = sum(vs, Vector()) / len(vs)
        return (-len(o.data.vertices), round(c.x, 4), round(c.z, 4))
    kept.sort(key=key)
    for i, o in enumerate(kept):
        o.name = f"ISL_{i:03d}"
    log(f"islands kept: {len(kept)} (dropped tiny <{min_verts} verts)")
    return kept


CENSUS_COLORS = [
    (0.90, 0.10, 0.10), (0.10, 0.55, 0.95), (0.10, 0.80, 0.20), (0.95, 0.75, 0.10),
    (0.75, 0.15, 0.85), (0.05, 0.85, 0.80), (0.95, 0.45, 0.05), (0.55, 0.35, 0.15),
    (0.95, 0.55, 0.75), (0.45, 0.95, 0.45), (0.25, 0.25, 0.95), (0.80, 0.80, 0.30),
    (0.50, 0.05, 0.30), (0.05, 0.45, 0.35), (0.65, 0.65, 0.95), (0.95, 0.85, 0.65),
    (0.30, 0.10, 0.60), (0.10, 0.30, 0.10), (0.60, 0.30, 0.30), (0.30, 0.60, 0.60),
]


def stage_census():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    islands = import_garment_islands()
    report = []
    for i, o in enumerate(islands):
        vs = world_verts(o)
        zs = [v.z for v in vs]; xs = [v.x for v in vs]; ys = [v.y for v in vs]
        c = sum(vs, Vector()) / len(vs)
        color = CENSUS_COLORS[i % len(CENSUS_COLORS)] if i < len(CENSUS_COLORS) else (0.5, 0.5, 0.5)
        mat = bpy.data.materials.new(f"census_{i:03d}")
        mat.diffuse_color = (*color, 1.0)
        o.data.materials.clear()
        o.data.materials.append(mat)
        report.append({"island": i, "verts": len(o.data.vertices), "tris": tris_of(o),
                       "z": [round(min(zs), 3), round(max(zs), 3)],
                       "x": [round(min(xs), 3), round(max(xs), 3)],
                       "centroid": [round(c.x, 3), round(c.y, 3), round(c.z, 3)],
                       "color": [round(v, 2) for v in color]})
        log(f"ISL_{i:03d} verts={len(o.data.vertices):7d} z=[{min(zs):.3f},{max(zs):.3f}] "
            f"x=[{min(xs):+.3f},{max(xs):+.3f}] color={tuple(round(v,2) for v in color)}")
    with open(work_path("census.json"), "w", encoding="utf-8") as f:
        json.dump(report, f, indent=1)
    qa_render("00_census", focus_objects=islands, color_type="MATERIAL")
    save_checkpoint("00_census")


# ── fabric texture ───────────────────────────────────────────────────────────
# Marvelous Designer bakes no texture into the export, but it DOES export the UVs
# it draped the fabric with, so re-pointing the same PNG at those UVs restores the
# print exactly where it was authored — no offset needed. The UVs are centred on
# zero and measured in texture repeats (a skirt whose v spans -0.5..0.5 wears the
# tile exactly once); negative coordinates wrap, so they sample correctly as-is.
#
# Verify a scale/offset change in the GAME, not a Blender QA render: Workbench
# shading ignores shader Mapping nodes, so a wrong transform renders correct here
# and wrong in Godot. That is why the transform is baked into the UV data below.

def apply_fabric_texture(mat, part_cfg, mesh):
    tex = part_cfg.get("texture")
    if not tex:
        return
    path = tex if os.path.isabs(tex) else os.path.join(HERE, tex)
    if not os.path.exists(path):
        log(f"WARNING: texture not found, keeping flat colour: {path}")
        return
    if not mesh.uv_layers:
        log(f"WARNING: no UV map on {mesh.name}, keeping flat colour")
        return

    scale = part_cfg.get("texture_scale", [1.0, 1.0])
    offset = part_cfg.get("texture_offset", [0.0, 0.0])

    if scale != [1.0, 1.0] or offset != [0.0, 0.0]:
        for loop_uv in mesh.uv_layers[0].data:
            loop_uv.uv[0] = loop_uv.uv[0] * scale[0] + offset[0]
            loop_uv.uv[1] = loop_uv.uv[1] * scale[1] + offset[1]

    nodes, links = mat.node_tree.nodes, mat.node_tree.links
    img_node = nodes.new("ShaderNodeTexImage")
    img_node.image = bpy.data.images.load(path, check_existing=True)
    bsdf = nodes.get("Principled BSDF")
    if bsdf:
        links.new(img_node.outputs["Color"], bsdf.inputs["Base Color"])
    log(f"texture {os.path.basename(path)} → {mesh.name} "
        f"(uv={mesh.uv_layers[0].name}, scale={scale}, offset={offset})")


# ── stage: prepare ───────────────────────────────────────────────────────────
# Group census islands into configured parts (each part lists island indices),
# import body+rig, align garment onto body.

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

    # body + rig first (defines target space)
    body_path = CFG["body"]
    bpy.ops.import_scene.gltf(filepath=body_path)
    rig = next(o for o in bpy.data.objects if o.type == "ARMATURE")
    rig.name = "RIG"
    body = max((o for o in bpy.data.objects if o.type == "MESH"),
               key=lambda o: len(o.data.vertices))
    # drop face-part extras (eyes etc.) that ride along in the body GLB
    for o in [o for o in bpy.data.objects if o.type == "MESH" and o is not body]:
        bpy.data.objects.remove(o, do_unlink=True)
    body.name = "BODY"
    apply_all_transforms([rig, body])
    log(f"body: {len(body.data.vertices)} verts, groups={len(body.vertex_groups)}, "
        f"bones={len(rig.data.bones)}")

    # garment islands (same deterministic indexing the census printed)
    islands = import_garment_islands()
    by_idx = {int(o.name.split("_")[1]): o for o in islands}
    claimed = set()
    for part_name, part_cfg in CFG["parts"].items():
        ids = part_cfg.get("islands", [])
        members = [by_idx[i] for i in ids if i in by_idx]
        if not members:
            log(f"WARNING: part {part_name} matched no islands ({ids})")
            continue
        claimed.update(ids)
        select_only(members, active=members[0])
        bpy.ops.object.join()
        part = bpy.context.view_layer.objects.active
        part.name = f"GARM_{part_name}"
        mat = bpy.data.materials.new(f"{part_name}_mat")
        col = part_cfg.get("color", [0.7, 0.7, 0.7])
        mat.diffuse_color = (*col, 1.0)
        if mat.use_nodes is False:
            mat.use_nodes = True
        bsdf = mat.node_tree.nodes.get("Principled BSDF")
        if bsdf:
            bsdf.inputs["Base Color"].default_value = (*col, 1.0)
            bsdf.inputs["Roughness"].default_value = 0.85
        apply_fabric_texture(mat, part_cfg, part.data)
        part.data.materials.clear()
        part.data.materials.append(mat)
        log(f"part {part.name}: islands={ids} verts={len(part.data.vertices)} tris={tris_of(part)}")
    for i, o in list(by_idx.items()):
        if i not in claimed and o.name.startswith("ISL_"):
            bpy.data.objects.remove(o, do_unlink=True)

    # bodyshell parts: fitted clothing built FROM the body (Islander technique):
    # copy the body, keep the configured z-band + torso-dominant vertices, push
    # outward along normals. Fit and weights are correct by construction.
    for part_name, part_cfg in CFG["parts"].items():
        if part_cfg.get("type") != "bodyshell":
            continue
        shellp = body.copy()
        shellp.data = body.data.copy()
        shellp.name = f"GARM_{part_name}"
        bpy.context.scene.collection.objects.link(shellp)
        z0, z1 = part_cfg["z0"], part_cfg["z1"]
        allowed = set(part_cfg.get("bones", ["spine_01", "spine_02", "spine_03", "pelvis"]))
        gnames = {i: g.name for i, g in enumerate(shellp.vertex_groups)}
        keep = set()
        for v in shellp.data.vertices:
            z = (shellp.matrix_world @ v.co).z
            if not (z0 <= z <= z1):
                continue
            best, bw = None, 0.0
            for g in v.groups:
                if g.weight > bw:
                    best, bw = gnames.get(g.group), g.weight
            if best in allowed:
                keep.add(v.index)
        bm = bmesh.new()
        bm.from_mesh(shellp.data)
        bm.verts.ensure_lookup_table()
        doomed = [v for v in bm.verts if v.index not in keep]
        bmesh.ops.delete(bm, geom=doomed, context="VERTS")
        bm.to_mesh(shellp.data)
        bm.free()
        d = shellp.modifiers.new("Offset", "DISPLACE")
        d.strength = part_cfg.get("offset_mm", 6) / 1000.0
        d.mid_level = 0.0
        apply_modifier(shellp, d.name)
        mat = bpy.data.materials.new(f"{part_name}_mat")
        col = part_cfg.get("color", [0.7, 0.7, 0.7])
        mat.use_nodes = True
        bsdf = mat.node_tree.nodes.get("Principled BSDF")
        if bsdf:
            bsdf.inputs["Base Color"].default_value = (*col, 1.0)
            bsdf.inputs["Roughness"].default_value = 0.85
        mat.diffuse_color = (*col, 1.0)
        shellp.data.materials.clear()
        shellp.data.materials.append(mat)
        log(f"bodyshell {shellp.name}: verts={len(shellp.data.vertices)} tris={tris_of(shellp)}")

    parts = [o for o in objs_prefixed("GARM_")
             if CFG["parts"][o.name[5:]].get("type") != "bodyshell"]
    if not parts:
        raise SystemExit("no parts matched any islands — check config vs census.json")

    # sleeve pose-warp FIRST, in garment-local space (unambiguous: sleeves are
    # the only geometry beyond x_beyond). Rigid rotation around the garment's
    # own shoulder point, bringing each cuff to horizontal (T-pose).
    for part_name, part_cfg in CFG["parts"].items():
        wcfg = part_cfg.get("sleeve_warp")
        if not wcfg:
            continue
        p = bpy.data.objects.get(f"GARM_{part_name}")
        if p is None:
            continue
        x_beyond = wcfg.get("x_beyond", 0.28)
        sx, sz = wcfg.get("shoulder", [0.19, 1.43])
        for sign in (1.0, -1.0):
            pivot = Vector((sign * sx, 0.0, sz))
            sleeve = [v for v in p.data.vertices if sign * v.co.x > x_beyond]
            if len(sleeve) < 10:
                log(f"sleeve_warp {part_name} side {sign:+.0f}: none found")
                continue
            tip = max(sleeve, key=lambda v: sign * v.co.x)
            cur = math.atan2(tip.co.z - pivot.z, sign * (tip.co.x - pivot.x))
            rot = Matrix.Rotation(-sign * cur, 4, "Y")
            M = Matrix.Translation(pivot) @ rot @ Matrix.Translation(-pivot)
            for v in sleeve:
                v.co = M @ v.co
            log(f"sleeve_warp {part_name} side {sign:+.0f}: {math.degrees(cur):+.1f} deg, "
                f"{len(sleeve)} verts")
    a = CFG.get("align", {})
    sxy = a.get("scale_xy", a.get("scale", 1.0)) or 1.0
    sz = a.get("scale_z", a.get("scale", 1.0)) or 1.0
    M = Matrix.Diagonal(Vector((sxy, sxy, sz))).to_4x4()
    for p in parts:
        p.data.transform(M)
    gz0, gz1 = zspan(parts)
    b_top = bone_head_z(obj("RIG"), a.get("top_bone", "neck_01"))
    dz = (b_top - gz1) + a.get("z_nudge", 0.0)
    dxy = Vector(a.get("xy_nudge", (0, 0)))
    T = Matrix.Translation(Vector((dxy.x, dxy.y, dz)))
    for p in parts:
        p.data.transform(T)
    log(f"align: scale_xy={sxy:.3f} scale_z={sz:.3f} dz={dz:+.4f}")

    # torso boost: heroic bodies are wider than MD mannequins — radially inflate
    # the above-waist region about the body axis so bodice/collar fabric starts
    # OUTSIDE the body (direction-preserving; the OUTSIDE shrinkwrap then only
    # catches stragglers instead of shredding buried cloth to random sides).
    tb = CFG.get("torso_boost")
    if tb:
        z0, z1, boost = tb.get("z0", 1.00), tb.get("z1", 1.08), tb.get("xy", 1.18)
        for p in parts:
            for v in p.data.vertices:
                z = v.co.z
                if z <= z0:
                    continue
                t = 1.0 if z >= z1 else (z - z0) / max(1e-9, z1 - z0)
                f = 1.0 + (boost - 1.0) * t
                v.co.x *= f
                v.co.y *= f
        log(f"torso_boost: xy x{boost} above z={z0}..{z1}")

    # post-align cut: drop configured fabric above a body-space height (v1 uses
    # this to remove the slim MD bodice+sleeves replaced by the bodyshell part)
    for part_name, part_cfg in CFG["parts"].items():
        cut = part_cfg.get("cut_above_z")
        p = bpy.data.objects.get(f"GARM_{part_name}")
        if cut is None or p is None or part_cfg.get("type") == "bodyshell":
            continue
        bm = bmesh.new()
        bm.from_mesh(p.data)
        ox = part_cfg.get("cut_outboard_x")
        oz = part_cfg.get("cut_outboard_z_above", 0.0)
        def out(v):
            w = p.matrix_world @ v.co
            return w.z > cut or (ox is not None and w.z > oz and abs(w.x) > ox)
        doomed = [v for v in bm.verts if out(v)]
        bmesh.ops.delete(bm, geom=doomed, context="VERTS")
        bm.to_mesh(p.data)
        bm.free()
        log(f"cut {p.name} above z={cut}: now {len(p.data.vertices)} verts")
        # prune disconnected scraps left by the cuts (e.g. sleeve remnants):
        # separate loose, keep the biggest piece + anything near the center axis
        px = part_cfg.get("prune_scraps_x")
        if px is not None:
            pname = p.name
            select_only([p])
            bpy.ops.object.mode_set(mode="EDIT")
            bpy.ops.mesh.select_all(action="SELECT")
            bpy.ops.mesh.separate(type="LOOSE")
            bpy.ops.object.mode_set(mode="OBJECT")
            pieces = [o for o in bpy.data.objects if o.name.startswith(pname)]
            main = max(pieces, key=lambda o: len(o.data.vertices))
            keepers = [main]
            for o in pieces:
                if o is main:
                    continue
                vs = world_verts(o)
                c = sum(vs, Vector()) / len(vs)
                if abs(c.x) <= px:
                    keepers.append(o)
                else:
                    bpy.data.objects.remove(o, do_unlink=True)
            select_only(keepers, active=main)
            bpy.ops.object.join()
            joined = bpy.context.view_layer.objects.active
            joined.name = pname
            log(f"prune {joined.name}: kept {len(keepers)} pieces, "
                f"{len(joined.data.vertices)} verts")

    qa_render("10_prepare")
    save_checkpoint("10_prepare")


# ── stage: fit ───────────────────────────────────────────────────────────────
# Build the inflated BODY_SHELL, give each part a Z-gradient influence mask, and
# shrinkwrap high-poly parts in configured layer order (later layers can target
# earlier ones).

def make_shell():
    body = obj("BODY")
    shell = body.copy()
    shell.data = body.data.copy()
    shell.name = "BODY_SHELL"
    bpy.context.scene.collection.objects.link(shell)
    d = shell.modifiers.new("Inflate", "DISPLACE")
    d.strength = CFG.get("shell_mm", 4) / 1000.0
    d.mid_level = 0.0
    apply_modifier(shell, d.name)
    shell.hide_render = True
    return shell


def zgradient_group(o, name, z_full, z_zero):
    """Vertex group: weight 1 at z_full … 0 at z_zero, linear between.

    Works in either direction: z_full above z_zero grades upward (waistband
    masks), z_full below z_zero grades downward (hem clearance masks).
    """
    vg = o.vertex_groups.get(name) or o.vertex_groups.new(name=name)
    mw = o.matrix_world
    for v in o.data.vertices:
        z = (mw @ v.co).z
        t = (z - z_zero) / ((z_full - z_zero) or 1e-9)
        vg.add([v.index], min(1.0, max(0.0, t)), "REPLACE")
    return vg


def stage_fit():
    load_checkpoint("10_prepare")
    rig = obj("RIG")
    shell = make_shell()
    hip_z = bone_head_z(rig, "pelvis")

    for part_name, part_cfg in CFG["parts"].items():
        p = bpy.data.objects.get(f"GARM_{part_name}")
        if p is None:
            continue
        if part_cfg.get("type") == "bodyshell":
            log(f"fit GARM_{part_name}: skipped (bodyshell)")
            continue
        fit = part_cfg.get("fit", {})
        target_name = fit.get("target", "BODY_SHELL")
        target = obj(target_name if target_name != "BODY_SHELL" else "BODY_SHELL")
        mode = fit.get("mask", "full")  # full | above_hip | waistband | none
        if mode == "none":
            log(f"fit {p.name}: skipped (mask none)")
            continue
        if mode == "full":
            vg = None
        elif mode == "above_hip":
            vg = zgradient_group(p, "fit_mask", hip_z + 0.02, hip_z - 0.06)
        elif mode == "waistband":
            vg = zgradient_group(p, "fit_mask", hip_z + 0.10, hip_z - 0.02)
        sw = p.modifiers.new("Fit", "SHRINKWRAP")
        sw.target = target
        sw.wrap_method = "NEAREST_SURFACEPOINT"
        sw.wrap_mode = fit.get("wrap_mode", "OUTSIDE")
        sw.offset = fit.get("offset_mm", 2) / 1000.0
        if vg is not None:
            sw.vertex_group = vg.name
        apply_modifier(p, sw.name)
        log(f"fit {p.name}: target={target.name} mask={mode} offset={sw.offset if hasattr(sw,'offset') else '?'}")

        # Graded hem clearance: the uniform shell offset is right at the waist but
        # far too tight where the legs travel — swing clips straight through a hem
        # fitted 4mm off the bind pose. A second OUTSIDE shrinkwrap, masked 0 at the
        # hip → 1 at the hem, pushes only too-close fabric out to hem_mm beyond the
        # shell; fabric already hanging clear never moves.
        hem_mm = fit.get("hem_mm", 0)
        if hem_mm:
            z_top = fit.get("hem_z_top", hip_z)
            z_bot = min(v.z for v in world_verts(p))
            hem_vg = zgradient_group(p, "hem_mask", z_bot, z_top)
            hw = p.modifiers.new("HemClear", "SHRINKWRAP")
            hw.target = target
            hw.wrap_method = "NEAREST_SURFACEPOINT"
            hw.wrap_mode = "OUTSIDE"
            hw.offset = hem_mm / 1000.0
            hw.vertex_group = hem_vg.name
            apply_modifier(p, hw.name)
            log(f"hem clearance {p.name}: +{hem_mm}mm graded z={z_top:.3f}→{z_bot:.3f}")

    qa_render("20_fit")
    save_checkpoint("20_fit")


# ── stage: reduce ────────────────────────────────────────────────────────────
# Keep a hidden HI_<part> copy for baking. Skirt-style parts are replaced by a
# generated cylinder proxy shrinkwrapped to the high-poly; others get planar
# dissolve then collapse decimate down to budget.

def make_skirt_proxy(part, proxy_cfg):
    seg = proxy_cfg.get("segments", 30)
    rings = proxy_cfg.get("rings", 6)
    vs = world_verts(part)
    zs = sorted(v.z for v in vs)
    z0, z1 = zs[int(0.01 * len(zs))], zs[int(0.99 * len(zs))]
    cx = sum(v.x for v in vs) / len(vs)
    cy = sum(v.y for v in vs) / len(vs)

    def ring_radius(z):
        band = [v for v in vs if abs(v.z - z) < (z1 - z0) / (rings * 1.5)]
        if not band:
            return 0.2
        ds = sorted(math.hypot(v.x - cx, v.y - cy) for v in band)
        return ds[int(0.65 * len(ds))]

    me = bpy.data.meshes.new(f"{part.name}_proxy")
    bm = bmesh.new()
    ring_verts = []
    for r in range(rings + 1):
        z = z1 - (z1 - z0) * (r / rings)
        rad = ring_radius(z)
        ring = [bm.verts.new((cx + rad * math.cos(2 * math.pi * s / seg),
                              cy + rad * math.sin(2 * math.pi * s / seg), z))
                for s in range(seg)]
        ring_verts.append(ring)
    for r in range(rings):
        for s in range(seg):
            a, b = ring_verts[r][s], ring_verts[r][(s + 1) % seg]
            c, d = ring_verts[r + 1][(s + 1) % seg], ring_verts[r + 1][s]
            bm.faces.new((a, b, c, d))
    bm.to_mesh(me)
    bm.free()
    proxy = bpy.data.objects.new(f"{part.name}_proxytmp", me)
    bpy.context.scene.collection.objects.link(proxy)
    sw = proxy.modifiers.new("WrapToSkirt", "SHRINKWRAP")
    sw.target = part
    sw.wrap_method = "NEAREST_SURFACEPOINT"
    apply_modifier(proxy, sw.name)
    select_only([proxy])
    bpy.ops.object.shade_smooth()
    # carry the part's (single) material
    if part.data.materials:
        proxy.data.materials.append(part.data.materials[0])
    return proxy


def stage_reduce():
    load_checkpoint("20_fit")
    for part_name, part_cfg in CFG["parts"].items():
        p = bpy.data.objects.get(f"GARM_{part_name}")
        if p is None:
            continue
        # keep high-poly for baking
        hi = p.copy()
        hi.data = p.data.copy()
        hi.name = f"HI_{part_name}"
        bpy.context.scene.collection.objects.link(hi)
        hi.hide_render = hi.hide_viewport = True

        wb = part_cfg.get("weld_band")   # [z0, z1, dist]: fuse stacked layers
        if wb:
            z0, z1, dist = wb
            select_only([p])
            bm = bmesh.new()
            bm.from_mesh(p.data)
            mw = p.matrix_world
            band = [v for v in bm.verts if z0 <= (mw @ v.co).z <= z1]
            bmesh.ops.remove_doubles(bm, verts=band, dist=dist)
            bm.to_mesh(p.data)
            bm.free()
            log(f"weld_band {p.name}: z=[{z0},{z1}] dist={dist} -> {len(p.data.vertices)} verts")

        budget = part_cfg.get("tris", 1500)
        if part_cfg.get("proxy") == "cylinder":
            proxy = make_skirt_proxy(p, part_cfg.get("proxy_cfg", {}))
            name = p.name
            bpy.data.objects.remove(p, do_unlink=True)
            proxy.name = name
            p = proxy
        else:
            # planar dissolve first (MD panels are near-planar), then collapse
            dm = p.modifiers.new("Planar", "DECIMATE")
            dm.decimate_type = "DISSOLVE"
            dm.angle_limit = math.radians(part_cfg.get("planar_deg", 5))
            apply_modifier(p, dm.name)
            cur = tris_of(p)
            if cur > budget:
                dm = p.modifiers.new("Collapse", "DECIMATE")
                dm.decimate_type = "COLLAPSE"
                dm.ratio = budget / cur
                apply_modifier(p, dm.name)
        select_only([p])
        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")
        log(f"reduce {p.name}: tris={tris_of(p)} (budget {budget})")

    qa_render("30_reduce", focus_objects=objs_prefixed("GARM_"))
    save_checkpoint("30_reduce")


# ── stage: bake ──────────────────────────────────────────────────────────────
# Normal + AO from HI_<part> onto GARM_<part> via a second UV map. Slow (Cycles).

def stage_bake():
    load_checkpoint("30_reduce")
    scene = bpy.context.scene
    scene.render.engine = "CYCLES"
    scene.cycles.device = "CPU"
    scene.cycles.samples = 32
    scene.cycles.seed = 7
    for part_name, part_cfg in CFG["parts"].items():
        lo = bpy.data.objects.get(f"GARM_{part_name}")
        hi = bpy.data.objects.get(f"HI_{part_name}")
        if lo is None or hi is None:
            continue
        hi.hide_viewport = hi.hide_render = False
        size = part_cfg.get("bake_size", 512)
        # bake UVs
        select_only([lo])
        uv = lo.data.uv_layers.new(name="bake_uv")
        lo.data.uv_layers.active = uv
        bpy.ops.object.mode_set(mode="EDIT")
        bpy.ops.mesh.select_all(action="SELECT")
        bpy.ops.uv.smart_project(angle_limit=math.radians(66))
        bpy.ops.object.mode_set(mode="OBJECT")
        # target images + nodes on the part's material
        mat = lo.data.materials[0] if lo.data.materials else None
        if mat is None:
            mat = bpy.data.materials.new(f"{part_name}_mat")
            lo.data.materials.append(mat)
        mat.use_nodes = True
        nt = mat.node_tree
        imgs = {}
        for kind in ("NORMAL", "AO"):
            img = bpy.data.images.new(f"{part_name}_{kind.lower()}", size, size,
                                      alpha=False, float_buffer=(kind == "NORMAL"))
            node = nt.nodes.new("ShaderNodeTexImage")
            node.image = img
            nt.nodes.active = node
            select_only([hi, lo], active=lo)
            bpy.ops.object.bake(type=kind, use_selected_to_active=True,
                                cage_extrusion=0.004,
                                use_clear=True)
            out = work_path(f"{part_name}_{kind.lower()}.png")
            img.filepath_raw = out
            img.file_format = "PNG"
            img.save()
            imgs[kind] = img
            log(f"baked {kind} for {part_name} -> {out}")
        # wire into the material (normal map + AO multiply into base color)
        bsdf = next(n for n in nt.nodes if n.type == "BSDF_PRINCIPLED")
        nrm_tex = nt.nodes.new("ShaderNodeTexImage")
        nrm_tex.image = imgs["NORMAL"]
        nrm_tex.image.colorspace_settings.name = "Non-Color"
        nmap = nt.nodes.new("ShaderNodeNormalMap")
        nmap.uv_map = "bake_uv"
        nt.links.new(nmap.inputs["Color"], nrm_tex.outputs["Color"])
        nt.links.new(bsdf.inputs["Normal"], nmap.outputs["Normal"])
        hi.hide_viewport = hi.hide_render = True
    save_checkpoint("40_bake")


# ── stage: skin ──────────────────────────────────────────────────────────────
# Weight transfer. Snug parts: Data Transfer (nearest face interpolated) from
# BODY. Skirt parts: cone proxy — data-transferred at the waist, pelvis-locked
# at the hem, then transferred onto the skirt.

def data_transfer_weights(dst, src):
    dst.vertex_groups.clear()
    for vg in src.vertex_groups:
        dst.vertex_groups.new(name=vg.name)
    dt = dst.modifiers.new("Weights", "DATA_TRANSFER")
    dt.object = src
    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"
    apply_modifier(dst, dt.name)


def smooth_and_normalize(o, iterations=3):
    select_only([o])
    bpy.ops.object.mode_set(mode="WEIGHT_PAINT")
    try:
        bpy.ops.object.vertex_group_smooth(group_select_mode="ALL",
                                           factor=0.5, repeat=iterations)
        bpy.ops.object.vertex_group_normalize_all(lock_active=False)
    finally:
        bpy.ops.object.mode_set(mode="OBJECT")


def strip_bone_groups(o, prefixes):
    """Remove vertex groups whose bone should never drive this garment (e.g.
    upperarm_* on a sleeveless top — nearest-face transfer grabs the arm surface
    at the armholes and arm swings then wrench the garment). The stripped verts
    are healed by the smooth+normalize pass that follows in stage_skin."""
    doomed = [vg for vg in o.vertex_groups
              if any(vg.name.startswith(pre) for pre in prefixes)]
    for vg in doomed:
        o.vertex_groups.remove(vg)
    return [vg.name for vg in doomed] if doomed else []


def pelvis_gradient(o, hip_z, hem_z, max_t=1.0):
    """Below the hip, blend all weights toward pelvis-only as z drops to hem.
    max_t < 1 leaves (1 - max_t) of the transferred leg weights alive at the
    hem so the cloth partially follows a raised thigh instead of letting it
    punch straight through a fully pelvis-locked skirt."""
    idx_pelvis = o.vertex_groups.find("pelvis")
    if idx_pelvis < 0:
        o.vertex_groups.new(name="pelvis")
        idx_pelvis = o.vertex_groups.find("pelvis")
    mw = o.matrix_world
    for v in o.data.vertices:
        z = (mw @ v.co).z
        if z >= hip_z:
            continue
        t = min(max_t, (hip_z - z) / max(1e-9, hip_z - hem_z))  # 0 at hip → max_t at hem
        for g in v.groups:
            g.weight *= (1.0 - t)
        o.vertex_groups[idx_pelvis].add([v.index], 0.0, "ADD")
        # set pelvis weight so the total stays 1 after normalize
        cur = sum(g.weight for g in v.groups)
        o.vertex_groups[idx_pelvis].add([v.index], max(0.0, 1.0 - cur), "ADD")


def skirt_bone_blend(o, rig, hip_z, ramp_m=0.20):
    """Blend transferred body weights toward the skirt_* bone ring below the hip
    (bodies built by ariki-game/tools/make_skirt_rig_body.py; the runtime
    SpringBoneSimulator3D in ariki-game SkirtSpringRig.cs owns those bones).

    Each vertex maps to the nearest TWO skirt bones by azimuth around the pelvis
    axis, angularly interpolated — adjacent verts land on near-identical blends,
    so there is no left/right seam to tear (the reason pelvis_gradient existed).
    Blend factor t ramps 0 at the hip → 1 at ramp_m below it: by mid-thigh the
    ring owns the cloth outright, which is what lets a thigh capsule push the
    skirt out of a lifted knee's way instead of averaging against leg weights.
    """
    mw_rig = rig.matrix_world
    pelvis = mw_rig @ rig.data.bones["pelvis"].head_local
    ring = []  # (azimuth, top-group idx, tip-group idx or None, z_mid)
    for b in rig.data.bones:
        if not b.name.startswith("skirt_") or b.name.endswith("_01"):
            continue
        head = mw_rig @ b.head_local
        az = math.atan2(head.y - pelvis.y, head.x - pelvis.x)
        if o.vertex_groups.find(b.name) < 0:
            o.vertex_groups.new(name=b.name)
        tip = rig.data.bones.get(b.name + "_01")
        tip_idx, z_mid = None, None
        if tip is not None:
            if o.vertex_groups.find(tip.name) < 0:
                o.vertex_groups.new(name=tip.name)
            tip_idx = o.vertex_groups.find(tip.name)
            z_mid = (mw_rig @ tip.head_local).z
        ring.append((az, o.vertex_groups.find(b.name), tip_idx, z_mid))
    if not ring:
        raise SystemExit("[skin] weights=skirt_bones but the body rig has no skirt_* "
                         "bones — point the config's `body` at a *_SkirtRig.glb copy")
    ring.sort()
    n = len(ring)
    skirt_idx = {i for _, top, tip, _ in ring for i in (top, tip) if i is not None}

    SEG_BAND = 0.05  # metres of top↔tip blend either side of the segment joint

    def strand_weights(entry, z, w):
        """Split strand weight w between its two segments by height."""
        _, top, tip, z_mid = entry
        if tip is None:
            return [(top, w)]
        fz = min(1.0, max(0.0, (z_mid + SEG_BAND - z) / (2 * SEG_BAND)))
        return [(top, w * (1.0 - fz)), (tip, w * fz)]

    tau = 2 * math.pi
    mw = o.matrix_world
    for v in o.data.vertices:
        co = mw @ v.co
        if co.z >= hip_z:
            continue
        t = min(1.0, (hip_z - co.z) / max(1e-9, ramp_m))
        az = math.atan2(co.y - pelvis.y, co.x - pelvis.x)
        # bracket az between ring[i-1] and ring[i] (cyclic)
        i = 0
        while i < n and ring[i][0] <= az:
            i += 1
        a, b_ = ring[(i - 1) % n], ring[i % n]
        span = (b_[0] - a[0]) % tau
        f = ((az - a[0]) % tau) / span if span > 1e-9 else 0.0
        # scale ALL existing (body) weights down, then add the two ring strands
        for g in v.groups:
            if g.group not in skirt_idx:
                g.weight *= (1.0 - t)
        for idx, w in strand_weights(a, co.z, t * (1.0 - f)) + strand_weights(b_, co.z, t * f):
            if w > 0.0:
                o.vertex_groups[idx].add([v.index], w, "ADD")
    log(f"skirt-bone blend {o.name}: {n} strands, ramp {ramp_m * 1000:.0f}mm below hip")


def stage_skin():
    src_stage = "40_bake" if os.path.exists(work_path("40_bake.blend")) else "30_reduce"
    load_checkpoint(src_stage)
    body, rig = obj("BODY"), obj("RIG")
    hip_z = bone_head_z(rig, "pelvis")

    for part_name, part_cfg in CFG["parts"].items():
        p = bpy.data.objects.get(f"GARM_{part_name}")
        if p is None:
            continue
        if part_cfg.get("type") == "bodyshell":
            pass  # weights inherited from the body copy — already correct
        elif part_cfg.get("weights") == "dress":
            # one-piece dress: body weights everywhere, then blend the loose
            # skirt region toward pelvis-only as z approaches the hem
            hem_z = zspan([p])[0]
            data_transfer_weights(p, body)
            pelvis_gradient(p, hip_z - 0.04, hem_z)
        elif part_cfg.get("weights") == "skirt_bones":
            # skirt_* ring bodies (make_skirt_rig_body.py): body weights at the
            # waistband, azimuth-blended to the spring-simulated ring below the
            # hip. The runtime capsules do what no static weighting can — move
            # the cloth out of a lifted knee's way.
            data_transfer_weights(p, body)
            skirt_bone_blend(p, rig, hip_z - 0.04,
                             ramp_m=part_cfg.get("skirt_ramp_m", 0.20))
        elif part_cfg.get("weights") == "skirt":
            hem_z = zspan([p])[0]
            cone = p.copy()  # skirt proxy IS already a leg-bridging cone shape
            cone.data = p.data.copy()
            cone.name = f"CONE_{part_name}"
            bpy.context.scene.collection.objects.link(cone)
            data_transfer_weights(cone, body)
            pelvis_gradient(cone, hip_z, hem_z)
            smooth_and_normalize(cone)
            data_transfer_weights(p, cone)
            cone.hide_viewport = cone.hide_render = True
        else:
            data_transfer_weights(p, body)
        smooth_and_normalize(p)
        # bind
        p.parent = rig
        p.matrix_parent_inverse = Matrix.Identity(4)
        am = p.modifiers.new("Armature", "ARMATURE")
        am.object = rig
        nonzero = sum(1 for v in p.data.vertices if v.groups)
        log(f"skin {p.name}: weighted {nonzero}/{len(p.data.vertices)} verts")

    qa_render("50_skin", focus_objects=objs_prefixed("GARM_"))
    save_checkpoint("50_skin")


# ── stage: export ────────────────────────────────────────────────────────────
# One GLB per outfit slot (multiple parts may share a slot). Armature + parts
# only; BODY/HI/CONE stay out.

def stage_export():
    load_checkpoint("50_skin")
    rig = obj("RIG")
    out_dir = CFG["export"]["out_dir"]
    os.makedirs(out_dir, exist_ok=True)
    slots = {}
    for part_name, part_cfg in CFG["parts"].items():
        p = bpy.data.objects.get(f"GARM_{part_name}")
        if p is not None:
            slots.setdefault(part_cfg["slot"], []).append(p)
    gender, set_name = CFG["export"]["gender"], CFG["export"]["set"]
    for slot, parts in slots.items():
        path = os.path.join(out_dir, f"{gender}_{set_name}_{slot}.gltf")
        select_only([rig] + parts, active=rig)
        bpy.ops.export_scene.gltf(filepath=path, export_format="GLTF_SEPARATE",
                                  use_selection=True, export_skins=True,
                                  export_animations=False, export_yup=True)
        log(f"EXPORTED {path} ({sum(tris_of(p) for p in parts)} tris, "
            f"{os.path.getsize(path)} bytes)")
    # determinism fingerprint
    import hashlib
    h = hashlib.sha1()
    for part in sorted(objs_prefixed("GARM_"), key=lambda o: o.name):
        for v in part.data.vertices:
            h.update(b"%d %d %d" % (int(v.co.x * 1e5), int(v.co.y * 1e5), int(v.co.z * 1e5)))
    log(f"vertex hash: {h.hexdigest()[:16]}")


# ── main ─────────────────────────────────────────────────────────────────────

STAGES = {"census": stage_census, "prepare": stage_prepare, "fit": stage_fit, "reduce": stage_reduce,
          "bake": stage_bake, "skin": stage_skin, "export": stage_export}


def main():
    global ARGS, CFG
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--stage", required=True, choices=sorted(STAGES))
    ARGS = ap.parse_args(argv)
    cfg_path = ARGS.config if os.path.isabs(ARGS.config) else os.path.join(HERE, ARGS.config)
    CFG = json.load(open(cfg_path, encoding="utf-8"))
    log(f"config={CFG['name']} stage={ARGS.stage}")
    STAGES[ARGS.stage]()


main()