"""Transfer TokenRig armature+weights from a rigged source onto an approved low mesh.

Blender:
  blender -b --factory-startup --python bl_transfer_skin.py -- \\
    <rigged_src.glb|fbx> <low_mesh.fbx|glb> <out.fbx>
"""
from __future__ import annotations

import argparse
import json
import math
import heapq
import sys
from pathlib import Path

import bpy
from mathutils import Matrix, Vector
from mathutils.bvhtree import BVHTree


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


def _import(path: Path):
    ext = path.suffix.lower()
    before = set(bpy.data.objects)
    if ext == ".fbx":
        bpy.ops.import_scene.fbx(filepath=str(path))
    elif ext in (".glb", ".gltf"):
        bpy.ops.import_scene.gltf(filepath=str(path))
    else:
        raise SystemExit(f"unsupported {ext}")
    return [o for o in bpy.data.objects if o not in before]


def _armatures(objs=None):
    pool = objs if objs is not None else bpy.data.objects
    return [o for o in pool if o.type == "ARMATURE"]


def _meshes(objs=None):
    pool = objs if objs is not None else bpy.data.objects
    # Skip TokenRig helper icosphere
    out = []
    for o in pool:
        if o.type != "MESH":
            continue
        if o.name.lower().startswith("icosphere"):
            continue
        if len(o.data.polygons) < 32:
            continue
        out.append(o)
    return out


def _largest_mesh(meshes):
    return max(meshes, key=lambda m: len(m.data.polygons))


def main():
    argv = sys.argv[sys.argv.index("--") + 1 :]
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("rigged_src", type=Path)
    parser.add_argument("low_mesh", type=Path)
    parser.add_argument("out", type=Path)
    parser.add_argument("--low-up", choices=("Y", "Z"), default="Y",
                        help="Project static FBX has Y-up vertex data; standard Blender/glTF imports use Z.")
    parser.add_argument("--align-bounds", action="store_true",
                        help="Align low bottom-center and height to its source high rig (rig/animations stay unchanged).")
    parser.add_argument('--weight-overrides',type=Path,help='Recipe with optional geodesic skinOverrides in approved mesh coordinates / height')
    args = parser.parse_args(argv)
    src_path, low_path, out_path = [p.resolve() for p in (args.rigged_src, args.low_mesh, args.out)]
    _clear()
    src_objs = _import(src_path)
    arms = _armatures(src_objs)
    src_meshes = _meshes(src_objs)
    if not arms or not src_meshes:
        raise SystemExit(f"source missing armature/mesh arms={len(arms)} meshes={len(src_meshes)}")
    arm = arms[0]
    donor = _largest_mesh(src_meshes)
    # Bind weights must be transferred in REST, never from a sampled action.
    arm.animation_data_clear()
    arm.data.pose_position = "REST"

    low_objs = _import(low_path)
    low_meshes = _meshes(low_objs)
    if not low_meshes:
        raise SystemExit("low mesh missing")
    low = _largest_mesh(low_meshes)

    # The static pipeline deliberately stores Y-up vertices (see bl_decimate).
    # A normal FBX re-import does NOT turn those vertices into Blender Z-up.
    # TokenRig/glTF is imported into Blender Z-up. Align BEFORE transferring weights.
    original_world = low.matrix_world.copy()
    original_points=[original_world @ v.co for v in low.data.vertices]
    axis = Matrix.Rotation(math.pi / 2, 4, "X") if args.low_up == "Y" else Matrix.Identity(4)
    low.parent = None
    low.matrix_world = Matrix.Identity(4)
    low.data.transform(axis @ original_world)
    low.data.update()
    bpy.context.view_layer.update()

    # Reject a different character/proxy unless it actually matches this low mesh.
    # Bone names/counts alone cannot establish compatible bind geometry.
    donor_points = [donor.matrix_world @ v.co for v in donor.data.vertices]
    alignment = Matrix.Identity(4)
    if args.align_bounds:
        def bottom_center(points):
            mn = Vector([min(p[i] for p in points) for i in range(3)])
            mx = Vector([max(p[i] for p in points) for i in range(3)])
            return Vector(((mn.x + mx.x)/2, (mn.y + mx.y)/2, mn.z)), mx.z - mn.z
        source_center, source_height = bottom_center(donor_points)
        low_center, low_height = bottom_center([v.co for v in low.data.vertices])
        scale = source_height / low_height
        alignment = Matrix.Translation(source_center) @ Matrix.Scale(scale, 4) @ Matrix.Translation(-low_center)
        low.data.transform(alignment)
        low.data.update()
    tree = BVHTree.FromPolygons(donor_points, [list(p.vertices) for p in donor.data.polygons])
    distances = [tree.find_nearest(v.co)[3] for v in low.data.vertices]
    extent = max(max(p[i] for p in donor_points) - min(p[i] for p in donor_points) for i in range(3))
    rms = math.sqrt(sum(d*d for d in distances) / len(distances))
    if max(distances) > extent * .03 or rms > extent * .01:
        raise SystemExit(f"bind geometry mismatch: max={max(distances):.6f} rms={rms:.6f} extent={extent:.6f}; "
                         "use this character's high rig and explicit --low-up")

    # Remove extra imported armatures from low file (should be none)
    for a in _armatures(low_objs):
        bpy.data.objects.remove(a, do_unlink=True)

    # Both are now in the SAME Blender world space; preserve this on parenting.
    bpy.ops.object.select_all(action="DESELECT")
    low.select_set(True)
    arm.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.parent_set(type="ARMATURE", keep_transform=True)

    # Ensure vertex groups exist matching bones
    for bone in arm.data.bones:
        if bone.name not in low.vertex_groups:
            low.vertex_groups.new(name=bone.name)

    # Data transfer weights from donor
    bpy.ops.object.select_all(action="DESELECT")
    donor.select_set(True)
    low.select_set(True)
    bpy.context.view_layer.objects.active = low
    mod = low.modifiers.new(name="WeightTransfer", type="DATA_TRANSFER")
    mod.object = donor
    mod.use_vert_data = True
    mod.data_types_verts = {"VGROUP_WEIGHTS"}
    mod.vert_mapping = "POLYINTERP_NEAREST"
    bpy.ops.object.datalayout_transfer(modifier=mod.name)
    bpy.ops.object.modifier_apply(modifier=mod.name)

    # Some inferred rigs omit a tail and attach it to a nearby arm. Local,
    # recipe-defined geodesic masks can bind that appendage to an existing
    # stable bone without changing geometry, UVs or the animation skeleton.
    override_reports=[]
    # Explicit source-space regions can bind a rigid carapace without altering
    # geometry/UVs. The recipe must be fitted to this character's inspected mesh.
    if args.weight_overrides:
        rules=json.loads(args.weight_overrides.read_text(encoding='utf-8-sig')).get('skinRegions',[])
        extent=max(p.y for p in original_points)-min(p.y for p in original_points)
        points=[p/extent for p in original_points]
        for rule in rules:
            if rule['bone'] not in arm.data.bones:raise ValueError('Region bone missing')
            lo,hi=rule['min'],rule['max'];feather=rule.get('feather',0)
            if feather<0 or any(lo[c]>=hi[c] for c in range(3)):raise ValueError('Invalid rigid region')
            target=low.vertex_groups[rule['bone']];affected=[];full=[]
            for i,p in enumerate(points):
                distance=max(max(lo[c]-p[c],0,p[c]-hi[c]) for c in range(3))
                if distance>feather:continue
                weight=1 if distance==0 else 1-distance/feather
                weight=weight*weight*(3-2*weight)
                previous={g.group:g.weight for g in low.data.vertices[i].groups}
                for group,value in previous.items():low.vertex_groups[group].add([i],value*(1-weight),'REPLACE')
                target.add([i],previous.get(target.index,0)*(1-weight)+weight,'REPLACE')
                affected.append(i)
                if distance==0:full.append(i)
            if not full:raise ValueError('Rigid region selected no vertices')
            override_reports.append(dict(kind='source_space_region',bone=rule['bone'],affected=len(affected),full_weight_vertices=len(full),
                normalized_min=[min(points[i][c] for i in affected) for c in range(3)],
                normalized_max=[max(points[i][c] for i in affected) for c in range(3)]))

    # More precise geodesic appendage overrides run after broad spatial regions.
    if args.weight_overrides:
        overrides=json.loads(args.weight_overrides.read_text(encoding='utf-8-sig')).get('skinOverrides',[])
        extent=max(p.y for p in original_points)-min(p.y for p in original_points)
        points=[p/extent for p in original_points]
        adjacency=[[] for _ in points]
        for edge in low.data.edges:
            a,b=edge.vertices;distance=(points[a]-points[b]).length
            adjacency[a].append((b,distance));adjacency[b].append((a,distance))
        unseen=set(range(len(points)));components=[]
        while unseen:
            stack=[unseen.pop()];part=[]
            while stack:
                v=stack.pop();part.append(v)
                for n,_ in adjacency[v]:
                    if n in unseen:unseen.remove(n);stack.append(n)
            components.append(part)
        main=max(components,key=len)
        for rule in overrides:
            if rule['bone'] not in arm.data.bones:raise ValueError('Override bone missing')
            full,fade=rule['fullDistance'],rule['fadeDistance']
            if not 0<=full<fade:raise ValueError('Invalid geodesic override falloff')
            seed=Vector(rule['seed']);start=min(main,key=lambda i:(points[i]-seed).length_squared)
            geodesic=[float('inf')]*len(points);geodesic[start]=0;queue=[(0,start)]
            while queue:
                distance,v=heapq.heappop(queue)
                if distance!=geodesic[v]:continue
                for n,cost in adjacency[v]:
                    trial=distance+cost
                    if trial<geodesic[n]:geodesic[n]=trial;heapq.heappush(queue,(trial,n))
            for part in components:
                if part is main:continue
                for i in part:
                    geodesic[i]=min(geodesic[j]+(points[i]-points[j]).length for j in main)
            target=low.vertex_groups[rule['bone']];affected=[]
            for i,distance in enumerate(geodesic):
                if distance>=fade:continue
                weight=1 if distance<=full else (fade-distance)/(fade-full)
                weight=weight*weight*(3-2*weight)
                previous={g.group:g.weight for g in low.data.vertices[i].groups}
                for group,value in previous.items():low.vertex_groups[group].add([i],value*(1-weight),'REPLACE')
                target.add([i],previous.get(target.index,0)*(1-weight)+weight,'REPLACE')
                affected.append(i)
            override_reports.append(dict(bone=rule['bone'],affected=len(affected),seed_vertex=start,
                normalized_min=[min(points[i][c] for i in affected) for c in range(3)],
                normalized_max=[max(points[i][c] for i in affected) for c in range(3)]))

    # Unity uses at most four weights here. Normalize after limiting, and reject
    # unweighted vertices instead of silently exporting broken skinning.
    bpy.ops.object.vertex_group_limit_total(limit=4)
    bpy.ops.object.vertex_group_normalize_all(lock_active=False)
    unweighted = [v.index for v in low.data.vertices if sum(g.weight for g in v.groups) < .999]
    if unweighted:
        raise SystemExit(f"unweighted/unnormalized vertices: {unweighted[:20]}")

    # Delete donor / helpers; keep arm + low
    keep = {arm, low}
    for obj in list(bpy.data.objects):
        if obj not in keep:
            bpy.data.objects.remove(obj, do_unlink=True)

    # Armature modifier
    if not any(m.type == "ARMATURE" for m in low.modifiers):
        am = low.modifiers.new(name="Armature", type="ARMATURE")
        am.object = arm
    arm.data.pose_position = "POSE"
    for bone in arm.pose.bones:
        bone.matrix_basis = Matrix.Identity(4)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    bpy.ops.object.select_all(action="DESELECT")
    arm.select_set(True)
    low.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.export_scene.fbx(
        filepath=str(out_path),
        use_selection=True,
        object_types={"ARMATURE", "MESH"},
        add_leaf_bones=False,
        bake_anim=False,
        armature_nodetype="NULL",
        axis_forward="-Z",
        axis_up="Y",
        apply_scale_options="FBX_SCALE_ALL",
        path_mode="COPY",
        embed_textures=True,
        mesh_smooth_type="FACE",
    )
    summary = {
        "ok": True,
        "out": str(out_path),
        "bones": len(arm.data.bones),
        "low_verts": len(low.data.vertices),
        "low_polys": len(low.data.polygons),
        "groups": len(low.vertex_groups),
        "low_up": args.low_up,
        "alignment_max": max(distances),
        "alignment_rms": rms,
        "source_rig": str(src_path),
        "source_low": str(low_path),
        "skin_overrides": override_reports,
        "low_to_rig_world": [list(row) for row in alignment @ axis @ original_world],
    }
    out_path.with_suffix(".skin.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print("MT_SUMMARY " + json.dumps(summary, ensure_ascii=False))


if __name__ == "__main__":
    main()