#!/usr/bin/env python3
"""
Headless Blender auto-rigger for Papaya-Shooter.

Takes an UNRIGGED humanoid mesh (GLB/glTF/FBX/OBJ), fits a Mixamo-compatible
skeleton to it, binds with automatic weights, and exports a rigged GLB.

Because the output skeleton uses standard Mixamo bone names, ANY Mixamo
animation (or CC0 pack retargeted to Mixamo names) can be merged onto the
result with tools/merge_animations.py — rig once, reuse every animation.

Usage:
  blender --background --python tools/autorig.py -- <input> <output.glb> [target_height] [library.glb]

Skeleton source, in order of preference:
  1. If a 4th arg (or assets/characters/animations/_library.glb) exists, the
     armature from that animation-library GLB is used as the template. This
     guarantees every library clip fits the character perfectly — same rig,
     no retargeting.
  2. Otherwise a Mixamo-named skeleton is built from heuristics (works with
     Mixamo-skeleton clip files instead).

If the input already contains an armature, it is kept as-is (bones are only
renamed to Mixamo convention when obvious matches exist) and the mesh is
re-exported normalized. Use tools/merge_animations.py next either way.

Heuristics assume a roughly upright humanoid in T-pose or A-pose. For models
that fail (extreme proportions, non-humanoids), use an external auto-rigger
(Mixamo web, AccuRig, Tripo/UniRig) and feed the rigged GLB straight to
merge_animations.py instead.
"""
import bpy
import sys
import os
from mathutils import Vector

argv = sys.argv
argv = argv[argv.index("--") + 1:] if "--" in argv else []
if len(argv) < 2:
    print("Usage: blender --background --python tools/autorig.py -- <input> <output.glb> [target_height]")
    sys.exit(1)

INPUT = argv[0]
OUTPUT = argv[1]
TARGET_HEIGHT = float(argv[2]) if len(argv) > 2 else 1.75
_default_library = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "assets", "characters", "animations", "_library.glb")
LIBRARY = argv[3] if len(argv) > 3 else (_default_library if os.path.exists(_default_library) else "")

# Mixamo bone names Godot/our animation library expects.
MIX = "mixamorig:"

# Common bone-name aliases -> Mixamo names, used when a rig already exists.
BONE_ALIASES = {
    "hips": "Hips", "pelvis": "Hips",
    "spine": "Spine", "spine1": "Spine1", "spine2": "Spine2",
    "chest": "Spine1", "upperchest": "Spine2",
    "neck": "Neck", "head": "Head",
    "leftshoulder": "LeftShoulder", "rightshoulder": "RightShoulder",
    "leftarm": "LeftArm", "leftupperarm": "LeftArm",
    "rightarm": "RightArm", "rightupperarm": "RightArm",
    "leftforearm": "LeftForeArm", "leftlowerarm": "LeftForeArm",
    "rightforearm": "RightForeArm", "rightlowerarm": "RightForeArm",
    "lefthand": "LeftHand", "righthand": "RightHand",
    "leftupleg": "LeftUpLeg", "leftupperleg": "LeftUpLeg", "leftthigh": "LeftUpLeg",
    "rightupleg": "RightUpLeg", "rightupperleg": "RightUpLeg", "rightthigh": "RightUpLeg",
    "leftleg": "LeftLeg", "leftlowerleg": "LeftLeg", "leftshin": "LeftLeg", "leftcalf": "LeftLeg",
    "rightleg": "RightLeg", "rightlowerleg": "RightLeg", "rightshin": "RightLeg", "rightcalf": "RightLeg",
    "leftfoot": "LeftFoot", "rightfoot": "RightFoot",
    "lefttoebase": "LeftToeBase", "lefttoe": "LeftToeBase",
    "righttoebase": "RightToeBase", "righttoe": "RightToeBase",
}


def clear_scene():
    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()
    for block_list in (bpy.data.meshes, bpy.data.armatures, bpy.data.materials, bpy.data.actions):
        for block in list(block_list):
            if block.users == 0:
                block_list.remove(block)


def import_model(path):
    ext = os.path.splitext(path)[1].lower()
    if ext in (".glb", ".gltf"):
        bpy.ops.import_scene.gltf(filepath=path)
    elif ext == ".fbx":
        bpy.ops.import_scene.fbx(filepath=path)
    elif ext == ".obj":
        if hasattr(bpy.ops.wm, "obj_import"):
            bpy.ops.wm.obj_import(filepath=path)
        else:
            bpy.ops.import_scene.obj(filepath=path)
    else:
        print(f"ERROR: unsupported format {ext}")
        sys.exit(1)


def normalize_alias(name):
    key = name.lower().replace("mixamorig:", "").replace("_", "").replace(" ", "").replace(".", "")
    key = key.replace("l_", "left").replace("r_", "right")
    return BONE_ALIASES.get(key)


def rename_existing_rig(arm):
    renamed = 0
    for bone in arm.data.bones:
        target = normalize_alias(bone.name)
        if target and not bone.name.startswith(MIX):
            bone.name = MIX + target
            renamed += 1
    print(f"Renamed {renamed} bones to Mixamo convention")


def mesh_slice_width(obj, y_frac, height, min_z):
    """Max |x| of vertices within a thin horizontal slice at y_frac of height."""
    z_lo = min_z + height * (y_frac - 0.03)
    z_hi = min_z + height * (y_frac + 0.03)
    max_x = 0.0
    mat = obj.matrix_world
    for v in obj.data.vertices:
        co = mat @ v.co
        if z_lo <= co.z <= z_hi:
            max_x = max(max_x, abs(co.x))
    return max_x


def fit_arms_to_mesh(arm, mesh):
    """Rotate each arm bone chain to lie along the mesh's actual arm direction.

    Sketchfab humanoids are usually modelled in an A-pose (arms angled down),
    but the library skeleton rests in a T-pose (arms horizontal). If we bind the
    A-pose mesh to T-pose arm bones, bone-heat weighting fails (bones sit outside
    the arms) and the animations shove the already-drooped arms across the body.

    Reorienting the arm bones down the real arm axis makes the rest pose fit the
    mesh, so bone-heat succeeds. The animations are then retargeted from the
    library's T-pose rest to this fitted rest in merge_animations.py.
    """
    mw = mesh.matrix_world
    verts = [mw @ v.co for v in mesh.data.vertices]
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.mode_set(mode="EDIT")
    eb = arm.data.edit_bones
    for side, sx in (("L", 1.0), ("R", -1.0)):
        chain = [eb.get(f"DEF-upper_arm.{side}"), eb.get(f"DEF-forearm.{side}"),
                 eb.get(f"DEF-hand.{side}")]
        chain = [b for b in chain if b]
        if not chain:
            continue
        shoulder = chain[0].head.copy()
        # Hand tip = farthest mesh vertex to this side, near arm height.
        arm_z = shoulder.z
        cand = [v for v in verts if (v.x * sx) > abs(shoulder.x) * 0.8
                and abs(v.z - arm_z) < 0.35]
        if not cand:
            continue
        hand_tip = max(cand, key=lambda v: v.x * sx)
        cur = chain[-1].tail - shoulder
        tgt = hand_tip - shoulder
        if cur.length < 1e-4 or tgt.length < 1e-4:
            continue
        rot = cur.normalized().rotation_difference(tgt.normalized())
        # Rotate the whole chain about the shoulder joint.
        for b in chain:
            b.head = shoulder + rot @ (b.head - shoulder)
            b.tail = shoulder + rot @ (b.tail - shoulder)
        sh = eb.get(f"DEF-shoulder.{side}")
        if sh:
            sh.tail = chain[0].head
    bpy.ops.object.mode_set(mode="OBJECT")
    print("Fitted arm bones to mesh A-pose")


def import_library_armature(height):
    """Import the animation library's armature as the rig template.

    Rigging to the exact skeleton the clips were authored on means clips need
    only a rest-pose retarget (handled in merge_animations.py).
    """
    before = set(bpy.data.objects)
    before_actions = set(bpy.data.actions)
    bpy.ops.import_scene.gltf(filepath=LIBRARY)
    new_objects = [o for o in bpy.data.objects if o not in before]

    arm = None
    for o in new_objects:
        if o.type == "ARMATURE":
            arm = o
    if not arm:
        print("WARNING: no armature in library, falling back to heuristic skeleton")
        for o in new_objects:
            bpy.data.objects.remove(o, do_unlink=True)
        return None

    # Drop the library's mannequin mesh and all bundled actions — we only
    # want the bare skeleton here; clips get merged in the next stage.
    for o in new_objects:
        if o is not arm:
            bpy.data.objects.remove(o, do_unlink=True)
    if arm.animation_data:
        arm.animation_data_clear()
    for a in [a for a in bpy.data.actions if a not in before_actions]:
        bpy.data.actions.remove(a)

    # Scale the rig uniformly so its height matches the character's.
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.select_all(action="DESELECT")
    arm.select_set(True)
    arm.location = (0, 0, 0)
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    rig_top = max((arm.matrix_world @ b.head_local).z for b in arm.data.bones)
    rig_top = max(rig_top, max((arm.matrix_world @ b.tail_local).z for b in arm.data.bones))
    if rig_top > 0.01:
        s = height / rig_top
        arm.scale = (s, s, s)
        bpy.ops.object.transform_apply(scale=True)
        print(f"Library rig scaled by {s:.3f} to height {height:.2f}")
    return arm


def build_mixamo_armature(height, min_z, shoulder_w, hip_w):
    """Create a Mixamo-named humanoid armature fitted to the mesh bounds.

    Proportions are standard humanoid ratios of total height; good enough for
    automatic weights on typical game characters.
    """
    bpy.ops.object.armature_add(enter_editmode=True, location=(0, 0, 0))
    arm = bpy.context.active_object
    arm.name = "Armature"
    eb = arm.data.edit_bones
    for b in list(eb):
        eb.remove(b)

    def z(frac):
        return min_z + height * frac

    def add(name, head, tail, parent=None, connect=False):
        b = eb.new(MIX + name)
        b.head = Vector(head)
        b.tail = Vector(tail)
        if parent:
            b.parent = eb[MIX + parent]
            b.use_connect = connect
        return b

    sw = shoulder_w * 0.75   # shoulder joint x
    hw = max(hip_w * 0.45, height * 0.055)  # hip joint x

    add("Hips",  (0, 0, z(0.53)), (0, 0, z(0.58)))
    add("Spine", (0, 0, z(0.58)), (0, 0, z(0.66)), "Hips", True)
    add("Spine1", (0, 0, z(0.66)), (0, 0, z(0.74)), "Spine", True)
    add("Spine2", (0, 0, z(0.74)), (0, 0, z(0.82)), "Spine1", True)
    add("Neck",  (0, 0, z(0.82)), (0, 0, z(0.87)), "Spine2", True)
    add("Head",  (0, 0, z(0.87)), (0, 0, z(1.00)), "Neck", True)

    for side, sx in (("Left", 1), ("Right", -1)):
        add(f"{side}Shoulder", (sx * sw * 0.25, 0, z(0.80)), (sx * sw, 0, z(0.80)), "Spine2")
        # Arms along +/-X (T-pose-ish); automatic weights tolerate A-pose meshes.
        arm_len = height * 0.16
        add(f"{side}Arm", (sx * sw, 0, z(0.80)), (sx * (sw + arm_len), 0, z(0.78)), f"{side}Shoulder", True)
        add(f"{side}ForeArm", (sx * (sw + arm_len), 0, z(0.78)), (sx * (sw + arm_len * 2), 0, z(0.76)), f"{side}Arm", True)
        add(f"{side}Hand", (sx * (sw + arm_len * 2), 0, z(0.76)), (sx * (sw + arm_len * 2.4), 0, z(0.75)), f"{side}ForeArm", True)

        add(f"{side}UpLeg", (sx * hw, 0, z(0.52)), (sx * hw, 0, z(0.29)), "Hips")
        add(f"{side}Leg", (sx * hw, 0, z(0.29)), (sx * hw, 0, z(0.06)), f"{side}UpLeg", True)
        add(f"{side}Foot", (sx * hw, 0, z(0.06)), (sx * hw, -height * 0.08, z(0.015)), f"{side}Leg", True)
        add(f"{side}ToeBase", (sx * hw, -height * 0.08, z(0.015)), (sx * hw, -height * 0.13, z(0.015)), f"{side}Foot", True)

    bpy.ops.object.mode_set(mode="OBJECT")
    return arm


def _count_weighted_verts(mesh):
    return sum(1 for v in mesh.data.vertices if len(v.groups) > 0)


# Bones that must never deform the mesh: the root/master and any rig control
# or mechanism bones. A deform-flagged root runs up the body centre-line, so
# without this it captures inner-leg/heel verts and stretches them to centre.
_NON_DEFORM_HINTS = ("root", "master", "mch-", "mch_", "ctrl", "org-", "-ik",
                     "_ik", "pole", "target", "properties")


def _deform_segments(arm):
    segs = []
    for b in arm.data.bones:
        if not b.use_deform:
            continue
        if any(h in b.name.lower() for h in _NON_DEFORM_HINTS):
            continue
        head = arm.matrix_world @ b.head_local
        tail = arm.matrix_world @ b.tail_local
        segs.append((b.name, head, tail))
    return segs


def rigid_nearest_bone_weights(mesh, arm):
    """Weight each vertex to its nearest deform bones (distance-falloff blend).

    A guaranteed-to-work substitute when bone-heat weighting fails — which it
    does on most imported characters, whose joined hair/clothing/body meshes
    have the interior and overlapping geometry the heat solver chokes on.

    Each vertex is blended across its nearest few bone segments with an
    inverse-distance falloff, so joints deform smoothly instead of tearing.
    Control/root bones are excluded (see _NON_DEFORM_HINTS) so the mesh follows
    real limbs, not the centre-line master bone.
    """
    segs = _deform_segments(arm)
    if not segs:
        return 0

    for vg in list(mesh.vertex_groups):
        mesh.vertex_groups.remove(vg)
    groups = {name: mesh.vertex_groups.new(name=name) for name, _, _ in segs}

    def dist_to_seg(p, a, b):
        ab = b - a
        denom = ab.dot(ab)
        t = 0.0 if denom == 0 else max(0.0, min(1.0, (p - a).dot(ab) / denom))
        return (p - (a + ab * t)).length

    K = 4          # blend across up to this many nearest bones
    FALLOFF = 3.0  # higher = tighter to the single nearest bone
    mw = mesh.matrix_world
    for v in mesh.data.vertices:
        p = mw @ v.co
        dists = sorted(((dist_to_seg(p, s[1], s[2]), s[0]) for s in segs),
                       key=lambda d: d[0])[:K]
        nearest = max(dists[0][0], 1e-5)
        # Inverse-distance weights, relative to the nearest bone.
        raw = [((nearest / max(d, 1e-5)) ** FALLOFF, name) for d, name in dists]
        total = sum(w for w, _ in raw)
        for w, name in raw:
            groups[name].add([v.index], w / total, "REPLACE")
    return _count_weighted_verts(mesh)


def bind_mesh_to_armature(mesh, arm):
    """Bind mesh to armature so the glTF exporter writes a COMPLETE skin
    (skin object + node.skin reference + per-vertex JOINTS/WEIGHTS).

    1. Try Blender automatic (bone-heat) weights for smooth deformation.
    2. If that assigns (almost) nothing — common on layered hair/clothing
       meshes where bone-heat fails — fall back to rigid nearest-bone.

    Keep the standard ARMATURE_AUTO result: an Armature modifier plus vertex
    groups (parent_type stays OBJECT). That is exactly what the exporter needs
    to write the vertex weights. (An earlier version stripped the modifier and
    used parent_type='ARMATURE'; that produced an ORPHAN skin with no weights,
    so the mesh rendered its bind pose — a permanent T-pose — in game.)
    """
    bpy.ops.object.select_all(action="DESELECT")
    mesh.select_set(True)
    arm.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.parent_set(type="ARMATURE_AUTO")

    weighted = _count_weighted_verts(mesh)
    total = len(mesh.data.vertices)
    if weighted < total * 0.5:
        print(f"Automatic weights covered {weighted}/{total} verts — "
              "falling back to rigid nearest-bone weights")
        weighted = rigid_nearest_bone_weights(mesh, arm)
    print(f"Bound mesh: {weighted}/{total} verts weighted")


def main():
    clear_scene()
    print(f"Importing {INPUT}...")
    import_model(INPUT)

    meshes = [o for o in bpy.data.objects if o.type == "MESH"]
    armatures = [o for o in bpy.data.objects if o.type == "ARMATURE"]
    if not meshes:
        print("ERROR: no mesh in input")
        sys.exit(1)

    # Join meshes so weights/normalization apply uniformly.
    bpy.ops.object.select_all(action="DESELECT")
    for m in meshes:
        m.select_set(True)
    bpy.context.view_layer.objects.active = meshes[0]
    if len(meshes) > 1:
        bpy.ops.object.join()
    mesh = bpy.context.active_object

    # Detach from the importer's hierarchy (Sketchfab wraps meshes in scaled/
    # rotated empties) so the mesh sits in clean world space, then bake its
    # own transform. Otherwise bind + normalization fight the parent transform.
    bpy.ops.object.select_all(action="DESELECT")
    mesh.select_set(True)
    bpy.context.view_layer.objects.active = mesh
    if mesh.parent:
        bpy.ops.object.parent_clear(type="CLEAR_KEEP_TRANSFORM")
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    for o in list(bpy.data.objects):
        if o.type == "EMPTY":
            bpy.data.objects.remove(o, do_unlink=True)

    if armatures:
        print("Input already rigged — keeping existing skeleton, renaming bones")
        rename_existing_rig(armatures[0])
    else:
        # Normalize: feet on ground, centered, target height.
        bb = [mesh.matrix_world @ Vector(c) for c in mesh.bound_box]
        min_z = min(c.z for c in bb)
        max_z = max(c.z for c in bb)
        cur_h = max_z - min_z
        if cur_h <= 0:
            print("ERROR: degenerate mesh bounds")
            sys.exit(1)
        s = TARGET_HEIGHT / cur_h
        mesh.scale = (s, s, s)
        bpy.ops.object.transform_apply(scale=True)
        bb = [mesh.matrix_world @ Vector(c) for c in mesh.bound_box]
        min_z = min(c.z for c in bb)
        cx = (min(c.x for c in bb) + max(c.x for c in bb)) / 2
        cy = (min(c.y for c in bb) + max(c.y for c in bb)) / 2
        mesh.location.x -= cx
        mesh.location.y -= cy
        mesh.location.z -= min_z
        bpy.ops.object.transform_apply(location=True)
        min_z = 0.0
        height = TARGET_HEIGHT

        arm = None
        from_library = False
        if LIBRARY:
            print(f"Using animation library skeleton: {LIBRARY}")
            arm = import_library_armature(height)
            from_library = arm is not None
        if not arm:
            shoulder_w = mesh_slice_width(mesh, 0.80, height, min_z)
            hip_w = mesh_slice_width(mesh, 0.53, height, min_z)
            print(f"Fitted: height={height:.2f} shoulder_w={shoulder_w:.2f} hip_w={hip_w:.2f}")
            arm = build_mixamo_armature(height, min_z, shoulder_w, hip_w)

        # Fit the arm bones to the mesh's real (usually A-pose) arm direction so
        # bone-heat weighting fits and the rest pose matches the geometry. The
        # library-clip retarget in merge_animations.py compensates for the
        # rest-pose change.
        if from_library:
            fit_arms_to_mesh(arm, mesh)

        bind_mesh_to_armature(mesh, arm)

    print(f"Exporting {OUTPUT}...")
    os.makedirs(os.path.dirname(os.path.abspath(OUTPUT)), exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=OUTPUT,
        export_format="GLB",
        export_yup=True,
        # NOTE: export_apply must stay False — applying modifiers on a skinned
        # mesh bakes away the Armature modifier and drops the skin binding
        # (glTF then exports bones as plain nodes with skins:0). Transforms are
        # already applied in code before binding, so nothing is lost here.
        export_apply=False,
        export_animations=True,
        export_skins=True,
    )
    print("Done.")


main()