# Custom bird rig + hand-keyed clips for powder_gull (Meshy mesh, no Meshy skeleton).
# blender -b -P gull_rig.py -- model.glb out.glb [out.blend]
import bpy, sys, math, mathutils
from mathutils import Vector, Quaternion

a = sys.argv[sys.argv.index("--") + 1:]
src, dst = a[0], a[1]
blend = a[2] if len(a) > 2 else None
bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=src)
me = [o for o in bpy.context.scene.objects if o.type == "MESH"][0]
bpy.context.view_layer.objects.active = me
me.select_set(True)
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
# weld UV-seam splits so weights stay continuous
import bmesh
bm = bmesh.new(); bm.from_mesh(me.data)
bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=1e-4)
bm.to_mesh(me.data); bm.free()

# ---- skeleton (Blender space: Z up, front -Y, character left = +X) ----
B = {
    "Root":   ((0, 0, -0.82), (0, 0, -0.62), None),
    "Hips":   ((0, -0.05, -0.30), (0, -0.05, 0.0), "Root"),
    "Chest":  ((0, -0.05, 0.0), (0, -0.05, 0.22), "Hips"),
    "Neck":   ((0, -0.05, 0.22), (0, -0.05, 0.32), "Chest"),
    "Head":   ((0, -0.05, 0.32), (0, -0.05, 0.78), "Neck"),
    "Tail":   ((-0.05, 0.15, -0.32), (-0.15, 0.45, -0.66), "Hips"),
}
for s, k in (("L", 1), ("R", -1)):
    lx = 0.11 if k > 0 else -0.17
    B["Thigh." + s] = ((lx, -0.07, -0.30), (lx, -0.06, -0.52), "Hips")
    B["Shin." + s] = ((lx, -0.06, -0.52), (lx + 0.02 * k, -0.12, -0.71), "Thigh." + s)
    B["Foot." + s] = ((lx + 0.02 * k, -0.12, -0.71), (lx + 0.06 * k, -0.32, -0.80), "Shin." + s)
    wy = 0.13 if k > 0 else 0.19
    B["Wing." + s] = ((0.27 * k, 0.0, 0.2), (0.6 * k, wy, 0.15), "Chest")
    B["Hand." + s] = ((0.6 * k, wy, 0.15), (0.95 * k, wy + 0.05, 0.2), "Wing." + s)

arm_data = bpy.data.armatures.new("GullRig")
arm = bpy.data.objects.new("GullRig", arm_data)
bpy.context.scene.collection.objects.link(arm)
bpy.context.view_layer.objects.active = arm
bpy.ops.object.mode_set(mode="EDIT")
for name, (h, t, p) in B.items():
    eb = arm_data.edit_bones.new(name); eb.head = h; eb.tail = t
for name, (h, t, p) in B.items():
    if p:
        eb = arm_data.edit_bones[name]; eb.parent = arm_data.edit_bones[p]
        eb.use_connect = (Vector(arm_data.edit_bones[p].tail) - Vector(h)).length < 1e-4
bpy.ops.object.mode_set(mode="OBJECT")


# ---- weights by zones ----
def smooth(t, a0, a1):
    """0 below a0, 1 above a1, smoothstep between."""
    u = max(0.0, min(1.0, (t - a0) / (a1 - a0)))
    return u * u * (3 - 2 * u)


groups = {n: me.vertex_groups.new(name=n) for n in B if n != "Root"}
for v in me.data.vertices:
    x, y, z = v.co
    w = {}
    ax = abs(x); s = "L" if x > 0 else "R"
    if z > 0.30:
        w = {"Head": 1.0}
    elif y > 0.2 and z < -0.3:
        w = {"Tail": 1.0}
    elif ax > 0.27 and z > -0.12:
        # wing: chest -> wing -> hand, blended at root and elbow
        r = smooth(ax, 0.27, 0.36); e = smooth(ax, 0.55, 0.66)
        w = {"Chest": 1 - r, "Wing." + s: r * (1 - e), "Hand." + s: r * e}
    elif z < -0.30:
        legside = "L" if x > -0.04 else "R"
        k = smooth(z, -0.36, -0.28)  # blend into hips at the top
        sh = 1 - smooth(z, -0.56, -0.48); ft = 1 - smooth(z, -0.74, -0.68)
        w = {"Hips": k, "Thigh." + legside: (1 - k) * (1 - sh),
             "Shin." + legside: (1 - k) * sh * (1 - ft), "Foot." + legside: (1 - k) * ft}
    else:
        c = smooth(z, -0.08, 0.08); hd = smooth(z, 0.22, 0.30)
        w = {"Hips": 1 - c, "Chest": c * (1 - hd), "Head": c * hd}
    for n, val in w.items():
        if val > 1e-3:
            groups[n].add([v.index], val, "REPLACE")
me.parent = arm
mod = me.modifiers.new("Armature", "ARMATURE"); mod.object = arm


# ---- animation helpers ----
def rot(axis, deg):
    return Quaternion(Vector(axis), math.radians(deg))


def bone_axis_rot(bone, world_axis, deg):
    """Rotation about a world axis expressed in the bone's rest frame."""
    m = arm.data.bones[bone].matrix_local.to_3x3().inverted()
    return Quaternion((m @ Vector(world_axis)).normalized(), math.radians(deg))


FOLD = 55  # wings lowered along the body outside of flight


def folded(pose):
    out = dict(pose)
    for sd in ("L", "R"):
        out["Wing." + sd] = flap(-FOLD, sd) + [(Z, 18 if sd == "L" else -18)] + list(pose.get("Wing." + sd, []))
        out["Hand." + sd] = flap(-15, sd) + list(pose.get("Hand." + sd, []))
    return out


def make_action(name, frames, keys, loop=True, fold=True):
    if fold:
        keys = {f: folded(p) for f, p in keys.items()}
    """keys: {frame: {bone: (world_axis, deg) list or 'loc': Vector}}"""
    act = bpy.data.actions.new(name)
    arm.animation_data_create(); arm.animation_data.action = act
    bones = [b.name for b in arm.pose.bones]
    for f in sorted(keys):
        pose = keys[f]
        for pb in arm.pose.bones:
            q = Quaternion()
            for axis, deg in pose.get(pb.name, []):
                q = bone_axis_rot(pb.name, axis, deg) @ q
            pb.rotation_mode = "QUATERNION"; pb.rotation_quaternion = q
            pb.location = arm.data.bones[pb.name].matrix_local.to_3x3().inverted() @ pose.get("loc:" + pb.name, Vector())
            pb.keyframe_insert("rotation_quaternion", frame=f)
            pb.keyframe_insert("location", frame=f)
    act.use_fake_user = True
    track = arm.animation_data.nla_tracks.new(); track.name = name
    track.strips.new(name, int(min(keys)), act)
    arm.animation_data.action = None
    return act


X, Y, Z = (1, 0, 0), (0, 1, 0), (0, 0, 1)
UP = Vector((0, 0, 1))


def flap(deg_up, s):  # wing raise about the forward axis; L raises with -Y rotation
    return [(Y, -deg_up if s == "L" else deg_up)]


clips = {}
# Idle: breathing, head glance, wings settle slightly
clips["Idle"] = make_action("Idle", 48, {
    1: {"Chest": [(X, 0)], "Head": [(Z, 0)], "Wing.L": flap(0, "L"), "Wing.R": flap(0, "R")},
    13: {"Chest": [(X, 3)], "Head": [(Z, 12)], "Wing.L": flap(-6, "L"), "Wing.R": flap(-6, "R")},
    25: {"Chest": [(X, 0)], "Head": [(Z, 0)], "Wing.L": flap(0, "L"), "Wing.R": flap(0, "R")},
    37: {"Chest": [(X, 3)], "Head": [(Z, -12)], "Wing.L": flap(-6, "L"), "Wing.R": flap(-6, "R")},
    49: {"Chest": [(X, 0)], "Head": [(Z, 0)], "Wing.L": flap(0, "L"), "Wing.R": flap(0, "R")},
})


def walk_pose(ph):
    sw = math.sin(ph); lift = max(0, math.sin(ph)); liftR = max(0, -math.sin(ph))
    return {"Thigh.L": [(X, -28 * sw)], "Thigh.R": [(X, 28 * sw)],
            "Shin.L": [(X, 25 * lift)], "Shin.R": [(X, 25 * liftR)],
            "Hips": [(Y, 7 * sw)], "Chest": [(Y, -4 * sw)],
            "Head": [(X, -6 * abs(math.cos(ph)))],
            "Wing.L": flap(-8 + 5 * sw, "L"), "Wing.R": flap(-8 - 5 * sw, "R"),
            "loc:Root": Vector((0, 0, 0.045 + 0.025 * abs(math.cos(ph))))}


clips["Walk"] = make_action("Walk", 24, {1 + i * 3: walk_pose(i * math.pi / 4) for i in range(9)})


def fly_pose(ph):
    up = math.sin(ph); lag = math.sin(ph - 0.9)
    return {"Wing.L": flap(42 * up, "L"), "Wing.R": flap(42 * up, "R"),
            "Hand.L": flap(28 * lag, "L"), "Hand.R": flap(28 * lag, "R"),
            "Hips": [(X, -18)], "Thigh.L": [(X, 70)], "Thigh.R": [(X, 70)],
            "Shin.L": [(X, -40)], "Shin.R": [(X, -40)], "Tail": [(X, -10)],
            "Head": [(X, 14)],
            "loc:Root": Vector((0, 0, 0.35 - 0.05 * up))}


clips["Fly"] = make_action("Fly", 16, {1 + i * 2: fly_pose(i * math.pi / 4) for i in range(9)}, fold=False)

# Attack (dive bomb): wings up, crouch, lunge + peck, recover
clips["Attack"] = make_action("Attack", 22, {
    1: {},
    6: {"Wing.L": flap(55, "L"), "Wing.R": flap(55, "R"), "Hips": [(X, 10)], "Head": [(X, -15)],
        "Thigh.L": [(X, -15)], "Thigh.R": [(X, -15)], "loc:Root": Vector((0, 0.05, 0.0))},
    11: {"Wing.L": flap(-35, "L"), "Wing.R": flap(-35, "R"), "Hips": [(X, -35)], "Head": [(X, 35)],
         "Thigh.L": [(X, 35)], "Thigh.R": [(X, 35)], "loc:Root": Vector((0, -0.28, 0.14))},
    15: {"Wing.L": flap(-20, "L"), "Wing.R": flap(-20, "R"), "Hips": [(X, -25)], "Head": [(X, 20)],
         "loc:Root": Vector((0, -0.22, 0.06))},
    23: {},
}, loop=False)
# Hurt: recoil backwards, wings jolt
clips["Hurt"] = make_action("Hurt", 12, {
    1: {},
    4: {"Hips": [(X, 18)], "Head": [(X, -25), (Z, 15)], "Wing.L": flap(35, "L"), "Wing.R": flap(25, "R"),
        "loc:Root": Vector((0, 0.08, 0.07))},
    13: {},
}, loop=False)
# Death: topple onto its back, wings drop open
clips["Death"] = make_action("Death", 30, {
    1: {},
    8: {"Hips": [(X, 25)], "Head": [(X, -30)], "Wing.L": flap(40, "L"), "Wing.R": flap(40, "R"),
        "loc:Root": Vector((0, 0.05, 0.08))},
    20: {"Root": [(X, 80)], "Head": [(X, -20), (Z, 25)], "Wing.L": flap(-30, "L"), "Wing.R": flap(-30, "R"),
         "Thigh.L": [(X, 50)], "Thigh.R": [(X, 30)], "loc:Root": Vector((0, 0.1, 0.3))},
    30: {"Root": [(X, 88)], "Head": [(X, -15), (Z, 30)], "Wing.L": flap(-35, "L"), "Wing.R": flap(-35, "R"),
         "Thigh.L": [(X, 55)], "Thigh.R": [(X, 35)], "Shin.L": [(X, -30)], "loc:Root": Vector((0, 0.12, 0.45))},
}, loop=False)

for act in bpy.data.actions:
    for fc in act.fcurves if hasattr(act, "fcurves") else []:
        for kp in fc.keyframe_points:
            kp.interpolation = "BEZIER"

# ---- ground check: lowest evaluated vertex per clip at a few frames ----
dg = bpy.context.evaluated_depsgraph_get()
for name, act in clips.items():
    arm.animation_data.action = act
    lows = []
    for f in (1, 6, 12, 20):
        bpy.context.scene.frame_set(f)
        dg = bpy.context.evaluated_depsgraph_get()
        e = me.evaluated_get(dg); m = e.to_mesh()
        lows.append(round(min((e.matrix_world @ vv.co).z for vv in m.vertices), 3)); e.to_mesh_clear()
    print("GROUND", name, lows)
arm.animation_data.action = None

if blend:
    bpy.ops.wm.save_as_mainfile(filepath=blend)
bpy.ops.object.select_all(action="DESELECT")
me.select_set(True); arm.select_set(True)
bpy.ops.export_scene.gltf(filepath=dst, export_format="GLB", use_selection=True,
                          export_animations=True, export_animation_mode="ACTIONS")
print("EXPORTED", dst, list(clips))
