# Re-skin a Meshy rigged GLB, keeping its skeleton and clips.
# blender -b -P reskin.py -- in.glb out.glb [voxel]
# env: ZONES=zones.json  KEEP=1 (keep Meshy weights, zones only)  RIGID=1 (one bone per vertex)
import bpy, sys, os, json, time, collections, bmesh, mathutils

a = sys.argv[sys.argv.index("--") + 1:]
src, dst = a[0], a[1]
voxel = float(a[2]) if len(a) > 2 else 0.012
KEEP = os.environ.get("KEEP") == "1"
RIGID = os.environ.get("RIGID") == "1"
T0 = time.time()


def log(*x):
    print("STEP %.1fs" % (time.time() - T0), *x, flush=True)


bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=src)
for o in list(bpy.context.scene.objects):
    if o.type == "MESH" and o.name.startswith("Icosphere"):
        bpy.data.objects.remove(o)
arm = [o for o in bpy.context.scene.objects if o.type == "ARMATURE"][0]
me = [o for o in bpy.context.scene.objects if o.type == "MESH"][0]
arm.data.pose_position = "REST"
bpy.context.view_layer.update()


def recenter_joints(targets):
    """Move listed joints to where the mesh actually is (world x/y, z kept); orientation unchanged,
    so Meshy clips (local rotations + near-zero translation deltas) pivot at the right place."""
    if not targets:
        return
    Ainv = arm.matrix_world.inverted().to_3x3()
    bpy.ops.object.select_all(action="DESELECT")
    bpy.context.view_layer.objects.active = arm; arm.select_set(True)
    bpy.ops.object.mode_set(mode="EDIT")
    for name, (tx, ty) in targets.items():
        eb = arm.data.edit_bones[name]
        h = arm.matrix_world @ eb.head
        d = Ainv @ mathutils.Vector((tx - h.x, ty - h.y, 0.0))
        eb.head += d; eb.tail += d
        log("recentred %-13s by (%.2f, %.2f)" % (name, tx - h.x, ty - h.y))
    bpy.ops.object.mode_set(mode="OBJECT")


def select_only(obj):
    bpy.ops.object.select_all(action="DESELECT")
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)


def weight_ops(*ops):
    select_only(me)
    bpy.ops.object.mode_set(mode="WEIGHT_PAINT")
    for op, kw in ops:
        op(**kw)
    bpy.ops.object.mode_set(mode="OBJECT")


Z = json.load(open(os.environ["ZONES"])) if os.environ.get("ZONES") else {}
recenter_joints(Z.get("recenter", {}))
bpy.context.view_layer.update()

if not KEEP:
    # glTF import splits vertices at UV seams; weld them so both sides share weights.
    bm = bmesh.new(); bm.from_mesh(me.data); n0 = len(bm.verts)
    bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=0.01)
    bm.to_mesh(me.data); bm.free()
    log("welded", n0, "->", len(me.data.vertices))
    # Watertight voxel proxy in world units (the Meshy mesh is in cm under a 0.01 armature).
    proxy = me.copy(); proxy.data = me.data.copy()
    bpy.context.scene.collection.objects.link(proxy)
    proxy.parent = None; proxy.matrix_world = me.matrix_world.copy()
    for m in list(proxy.modifiers):
        proxy.modifiers.remove(m)
    proxy.vertex_groups.clear()
    select_only(proxy)
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    r = proxy.modifiers.new("vox", "REMESH"); r.mode = "VOXEL"; r.voxel_size = voxel
    bpy.ops.object.modifier_apply(modifier="vox")
    dec = proxy.modifiers.new("dec", "DECIMATE"); dec.ratio = min(1.0, 25000 / max(1, len(proxy.data.polygons)))
    bpy.ops.object.modifier_apply(modifier="dec")
    proxy.select_set(True); arm.select_set(True); bpy.context.view_layer.objects.active = arm
    bpy.ops.object.parent_set(type="ARMATURE_AUTO")
    log("heat done on proxy", len(proxy.data.vertices))
    names = [b.name for b in arm.data.bones if b.use_deform]
    me.vertex_groups.clear()
    for n in names:
        me.vertex_groups.new(name=n)
    dt = me.modifiers.new("dt", "DATA_TRANSFER"); dt.object = proxy; 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"; dt.use_object_transform = True
    select_only(me)
    bpy.ops.object.modifier_move_to_index(modifier="dt", index=0)
    bpy.ops.object.modifier_apply(modifier="dt")
    bpy.data.objects.remove(proxy)
    weight_ops((bpy.ops.object.vertex_group_smooth, dict(group_select_mode="ALL", factor=0.5, repeat=2)),
               (bpy.ops.object.vertex_group_clean, dict(group_select_mode="ALL", limit=0.02)),
               (bpy.ops.object.vertex_group_limit_total, dict(group_select_mode="ALL", limit=4)),
               (bpy.ops.object.vertex_group_normalize_all, dict(lock_active=False)))
else:
    log("keeping Meshy weights")

# ---- zone rules ----
A = arm.matrix_world; seg = {}
for b in arm.data.bones:
    h = A @ b.head_local
    ch = list(b.children)
    if ch:  # towards the farthest child (Head -> head_end rather than headfront)
        t = max((A @ c.head_local for c in ch), key=lambda q: (q - h).length)
    elif b.parent:  # glTF leaf tails are arbitrary (tens of metres away): extend along the parent
        ph = A @ b.parent.head_local; d = h - ph
        t = h + d.normalized() * max(0.05, 0.6 * d.length) if d.length > 1e-6 else h
    else:
        t = h
    seg[b.name] = (h, t)


def dist(p, s0, s1):
    ab = s1 - s0; l = ab.length_squared
    u = 0 if l < 1e-12 else max(0, min(1, (p - s0).dot(ab) / l))
    return (s0 + ab * u - p).length


gi = {g.name: g.index for g in me.vertex_groups}; ig = {v: k for k, v in gi.items()}


def rigid_bone(p):
    for r in Z.get("rigid", []):
        if all(eval(c, {"x": p.x, "y": p.y, "z": p.z, "abs": abs}) for c in r["where"]):
            return r["bone"]
    return None


def reassert_rigid():
    """Smoothing blurs zone borders inward; put rigid zones back to a single bone."""
    n = 0
    for v in me.data.vertices:
        b = rigid_bone(MW @ v.co)
        if b:
            for g in list(v.groups):
                me.vertex_groups[g.group].remove([v.index])
            me.vertex_groups[b].add([v.index], 1.0, "REPLACE"); n += 1
    log("rigid zones re-asserted", n)
MW = me.matrix_world; nrig = nlim = 0
limits = Z.get("limit", {})
for v in me.data.vertices:
    p = MW @ v.co
    w = {ig[g.group]: g.weight for g in v.groups if g.weight > 0}
    rig = rigid_bone(p)
    near = None
    for r in Z.get("nearest", []):
        if all(eval(c, {"x": p.x, "y": p.y, "z": p.z, "abs": abs}) for c in r["where"]):
            near = r["bones"]; break
    if near and not rig:
        rig = min(near, key=lambda b: dist(p, *seg[b]))
    if rig:
        w = {rig: 1.0}; nrig += 1
    else:
        for bn, lim in limits.items():
            if bn in w and dist(p, *seg[bn]) > lim:
                w.pop(bn); nlim += 1
        if sum(w.values()) < 1e-4:
            cand = [b for b in seg if b in gi and b not in limits]
            w = {min(cand, key=lambda b: dist(p, *seg[b])): 1.0}
        if RIGID:
            w = {max(w, key=w.get): 1.0}
        tot = sum(w.values()); w = {k: x / tot for k, x in w.items()}
    for g in list(v.groups):
        me.vertex_groups[g.group].remove([v.index])
    for k, x in w.items():
        me.vertex_groups[k].add([v.index], x, "REPLACE")
log("zones: rigid", nrig, "limited", nlim, "rigid-mode" if RIGID else "")
if RIGID:  # blend only along block borders; block interiors keep a single bone
    weight_ops((bpy.ops.object.vertex_group_smooth, dict(group_select_mode="ALL", factor=0.5, repeat=int(os.environ.get("SEAM", "3")))),
               (bpy.ops.object.vertex_group_normalize_all, dict(lock_active=False)))
elif Z and not KEEP:
    weight_ops((bpy.ops.object.vertex_group_smooth, dict(group_select_mode="ALL", factor=0.3, repeat=1)),
               (bpy.ops.object.vertex_group_normalize_all, dict(lock_active=False)))
if Z.get("rigid") and (RIGID or not KEEP):
    reassert_rigid()

cnt = collections.Counter()
for v in me.data.vertices:
    gs = [g for g in v.groups if g.weight > 0.01]
    if gs:
        cnt[me.vertex_groups[max(gs, key=lambda g: g.weight).group].name] += 1
log("DOMINANT", cnt.most_common(8))
arm.data.pose_position = "POSE"
bpy.ops.object.select_all(action="SELECT")
bpy.ops.export_scene.gltf(filepath=dst, export_format="GLB", use_selection=True, export_animations=True)
log("exported", dst)
