#!/usr/bin/env python3
"""
Add animations to an already-rigged Miku GLB.
Keyframes ALL bones for each animation to ensure proper track export.
"""
import bpy
import sys
import os
import math

argv = sys.argv
if '--' in argv:
    argv = argv[argv.index('--') + 1:]

input_path = argv[0] if len(argv) > 0 else ''
output_path = argv[1] if len(argv) > 1 else ''

if not input_path or not output_path:
    sys.exit(1)

bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()

print(f"Importing {input_path}...")
bpy.ops.import_scene.gltf(filepath=input_path)

armatures = [o for o in bpy.data.objects if o.type == 'ARMATURE']
meshes = [o for o in bpy.data.objects if o.type == 'MESH' and o.name != 'Icosphere']

if not armatures:
    print("ERROR: No armature")
    sys.exit(1)

arm = armatures[0]
print(f"Armature: {arm.name} ({len(arm.data.bones)} bones)")

bpy.context.view_layer.objects.active = arm
bpy.ops.object.mode_set(mode='POSE')

bone_names = [b.name for b in arm.data.bones if b.name != "neutral_bone"]
print(f"Bones: {bone_names}")

def create_animation(name, frame_end, keyframes_func):
    """Create an animation. keyframes_func(frame, bone_name) -> (loc, rot) or None."""
    print(f"  Creating {name}...")
    action = bpy.data.actions.new(name=name)
    arm.animation_data.action = action
    action.frame_range = (1, frame_end)
    
    # Keyframe ALL bones at every frame to ensure tracks are exported
    for frame in range(1, frame_end + 1):
        for bone_name in bone_names:
            bone = arm.pose.bones[bone_name]
            result = keyframes_func(frame, bone_name)
            if result is None:
                # Default: no movement
                bone.location = (0, 0, 0)
                bone.rotation_euler = (0, 0, 0)
            else:
                loc, rot = result
                bone.location = loc if loc else (0, 0, 0)
                bone.rotation_mode = 'XYZ'
                bone.rotation_euler = rot if rot else (0, 0, 0)
            
            bone.keyframe_insert(data_path="location", frame=frame)
            bone.keyframe_insert(data_path="rotation_euler", frame=frame)
    
    print(f"    {name}: {frame_end} frames, {len(bone_names)} bones")

# Idle: visible swaying (large enough to not be optimized away)
def idle_kf(frame, bone_name):
    t = (frame - 1) / 30.0
    phase = t * 2 * math.pi
    if bone_name == "Spine":
        return None, (math.sin(phase) * 0.5, 0, 0)
    elif bone_name == "LeftUpperArm":
        return None, (0, 0, 0.5 + math.sin(phase) * 0.3)
    elif bone_name == "RightUpperArm":
        return None, (0, 0, -0.5 - math.sin(phase) * 0.3)
    elif bone_name == "Neck":
        return None, (math.sin(phase) * 0.3, 0, 0)
    elif bone_name == "LeftLowerArm":
        return None, (-0.5, 0, 0)
    elif bone_name == "RightLowerArm":
        return None, (-0.5, 0, 0)
    return None, None

create_animation("Idle", 30, idle_kf)

# Walk
def walk_kf(frame, bone_name):
    t = (frame - 1) / 24.0
    phase = t * 2 * math.pi
    if bone_name == "Hips":
        return (0, 0, abs(math.sin(phase * 2)) * 0.02), None
    elif bone_name == "LeftUpperLeg":
        return None, (math.sin(phase) * 0.5, 0, 0)
    elif bone_name == "LeftLowerLeg":
        return None, (max(0, -math.sin(phase) * 0.4 + 0.3), 0, 0)
    elif bone_name == "RightUpperLeg":
        return None, (math.sin(phase + math.pi) * 0.5, 0, 0)
    elif bone_name == "RightLowerLeg":
        return None, (max(0, -math.sin(phase + math.pi) * 0.4 + 0.3), 0, 0)
    elif bone_name == "LeftUpperArm":
        return None, (-math.sin(phase) * 0.4, 0, 0.1)
    elif bone_name == "LeftLowerArm":
        return None, (-0.5 + max(0, math.sin(phase) * 0.3), 0, 0)
    elif bone_name == "RightUpperArm":
        return None, (-math.sin(phase + math.pi) * 0.4, 0, -0.1)
    elif bone_name == "RightLowerArm":
        return None, (-0.5 + max(0, math.sin(phase + math.pi) * 0.3), 0, 0)
    elif bone_name == "Spine":
        return None, (0, math.sin(phase) * 0.08, 0)
    return None, None

create_animation("Walk", 24, walk_kf)

# Run
def run_kf(frame, bone_name):
    t = (frame - 1) / 20.0
    phase = t * 2 * math.pi
    if bone_name == "Hips":
        return (0, 0, abs(math.sin(phase * 2)) * 0.04), None
    elif bone_name == "LeftUpperLeg":
        return None, (math.sin(phase) * 0.9, 0, 0)
    elif bone_name == "LeftLowerLeg":
        return None, (max(0.1, -math.sin(phase) * 0.7 + 0.4), 0, 0)
    elif bone_name == "RightUpperLeg":
        return None, (math.sin(phase + math.pi) * 0.9, 0, 0)
    elif bone_name == "RightLowerLeg":
        return None, (max(0.1, -math.sin(phase + math.pi) * 0.7 + 0.4), 0, 0)
    elif bone_name == "LeftUpperArm":
        return None, (-math.sin(phase) * 0.8, 0, 0.15)
    elif bone_name == "LeftLowerArm":
        return None, (-1.2 + max(0, math.sin(phase) * 0.4), 0, 0)
    elif bone_name == "RightUpperArm":
        return None, (-math.sin(phase + math.pi) * 0.8, 0, -0.15)
    elif bone_name == "RightLowerArm":
        return None, (-1.2 + max(0, math.sin(phase + math.pi) * 0.4), 0, 0)
    elif bone_name == "Spine":
        return None, (0.15, math.sin(phase) * 0.1, 0)
    return None, None

create_animation("Run", 20, run_kf)

# Jump
def jump_kf(frame, bone_name):
    if frame <= 5:
        t = (frame - 1) / 4.0
        if bone_name == "Hips":
            return (0, 0, -t * 0.3), None
        elif bone_name in ("LeftUpperLeg", "RightUpperLeg"):
            return None, (-t * 0.7, 0, 0)
        elif bone_name in ("LeftLowerLeg", "RightLowerLeg"):
            return None, (t * 1.2, 0, 0)
        elif bone_name == "LeftUpperArm":
            return None, (0, 0, -t * 0.8)
        elif bone_name == "RightUpperArm":
            return None, (0, 0, t * 0.8)
    elif frame <= 8:
        t = (frame - 5) / 2.0
        if bone_name == "Hips":
            return (0, 0, -0.3 + t * 0.3), None
        elif bone_name in ("LeftUpperLeg", "RightUpperLeg"):
            return None, (-0.7 + t * 0.7, 0, 0)
        elif bone_name in ("LeftLowerLeg", "RightLowerLeg"):
            return None, (1.2 - t * 0.4, 0, 0)
        elif bone_name == "LeftUpperArm":
            return None, (-t * 2.5, 0, -0.8 - t)
        elif bone_name == "RightUpperArm":
            return None, (-t * 2.5, 0, 0.8 + t)
    elif frame <= 15:
        if bone_name == "LeftUpperArm":
            return None, (-2.5, 0, -1.8)
        elif bone_name == "RightUpperArm":
            return None, (-2.5, 0, 1.8)
        elif bone_name in ("LeftLowerArm", "RightLowerArm"):
            return None, (-0.3, 0, 0)
        elif bone_name == "LeftUpperLeg":
            return None, (0.3, 0, 0)
        elif bone_name == "LeftLowerLeg":
            return None, (0.6, 0, 0)
        elif bone_name == "RightUpperLeg":
            return None, (0.2, 0, 0)
        elif bone_name == "RightLowerLeg":
            return None, (0.4, 0, 0)
    else:
        t = (frame - 15) / 4.0
        if bone_name == "Hips":
            return (0, 0, -t * 0.3), None
        elif bone_name in ("LeftUpperLeg", "RightUpperLeg"):
            return None, (-t * 0.6, 0, 0)
        elif bone_name in ("LeftLowerLeg", "RightLowerLeg"):
            return None, (t * 1.0, 0, 0)
        elif bone_name == "LeftUpperArm":
            return None, (-2.5 + t * 2.5, 0, -1.8 + t * 1.8)
        elif bone_name == "RightUpperArm":
            return None, (-2.5 + t * 2.5, 0, 1.8 - t * 1.8)
    return None, None

create_animation("Jump", 20, jump_kf)

# Crouch
def crouch_kf(frame, bone_name):
    if bone_name == "Hips":
        return (0, 0, -0.5), None
    elif bone_name == "Spine":
        return None, (0.4, 0, 0)
    elif bone_name in ("LeftUpperLeg", "RightUpperLeg"):
        return None, (-0.9, 0, 0)
    elif bone_name in ("LeftLowerLeg", "RightLowerLeg"):
        return None, (1.5, 0, 0)
    elif bone_name == "LeftUpperArm":
        return None, (-0.6, 0, 0.3)
    elif bone_name == "RightUpperArm":
        return None, (-0.6, 0, -0.3)
    return None, None

create_animation("Crouch", 1, crouch_kf)

# Death
def death_kf(frame, bone_name):
    if frame <= 14:
        t = (frame - 1) / 13.0
        if bone_name == "Hips":
            return (0, 0, -t * 0.6), (-t * 1.3, 0, t * 0.3)
        elif bone_name == "Spine":
            return None, (-t * 0.6, t * 0.4, 0)
        elif bone_name == "Chest":
            return None, (-t * 0.5, 0, 0)
        elif bone_name in ("LeftUpperArm", "RightUpperArm"):
            sign = -1 if "Left" in bone_name else 1
            return None, (t * 2.2, 0, sign * -t * 0.6)
        elif bone_name in ("LeftUpperLeg", "RightUpperLeg"):
            return None, (-t * 0.4, 0, 0)
    else:
        t = (frame - 14) / 15.0
        if bone_name == "Hips":
            return (0, 0, -0.6 - t * 0.3), (-1.3 - t * 0.6, 0.8 * (1 - t), 0.3 + t * 0.4)
        elif bone_name == "Spine":
            return None, (-0.8 - t * 0.4, 0.5 * (1 - t), 0)
        elif bone_name == "Chest":
            return None, (-0.7 - t * 0.3, 0, 0)
        elif bone_name == "LeftHand":
            return None, (t * 0.4, t * 0.6, -0.8 - t * 0.6)
        elif bone_name == "RightHand":
            return None, (t * 0.5, -t * 0.7, 0.9 + t * 0.5)
    return None, None

create_animation("Death", 30, death_kf)

print(f"\nTotal animations: {len(bpy.data.actions)}")
for action in bpy.data.actions:
    print(f"  {action.name} ({action.frame_range[0]:.0f}-{action.frame_range[1]:.0f})")

bpy.ops.object.mode_set(mode='OBJECT')

bpy.ops.object.select_all(action='DESELECT')
arm.select_set(True)
for m in meshes:
    m.select_set(True)
bpy.context.view_layer.objects.active = arm

print(f"\nExporting to {output_path}...")
bpy.ops.export_scene.gltf(
    filepath=output_path,
    export_format='GLB',
    export_apply=True,
    export_animations=True,
    export_animation_mode='ACTIONS',
    export_skins=True,
    export_yup=True,
)
print(f"Exported! {os.path.getsize(output_path)} bytes")

import struct, json

with open(output_path, 'rb') as f:
    data = f.read()

json_len = struct.unpack('<I', data[12:16])[0]
json_data = data[20:20+json_len]
gltf = json.loads(json_data)

nodes = gltf.get('nodes', [])
skeleton_idx = None
for i, n in enumerate(nodes):
    if n.get('name') == 'MikuRig':
        skeleton_idx = i
        break

if skeleton_idx is not None:
    skins = gltf.get('skins', [])
    if skins:
        skins[0]['skeleton'] = skeleton_idx

new_json = json.dumps(gltf, separators=(',', ':')).encode('utf-8')
while len(new_json) % 4 != 0:
    new_json += b' '

bin_start = 20 + json_len
bin_len = struct.unpack('<I', data[bin_start:bin_start+4])[0]
new_length = 12 + 8 + len(new_json) + 8 + bin_len

new_data = b'glTF'
new_data += struct.pack('<I', 2)
new_data += struct.pack('<I', new_length)
new_data += struct.pack('<I', len(new_json))
new_data += b'JSON'
new_data += new_json
new_data += data[bin_start:]

with open(output_path, 'wb') as f:
    f.write(new_data)

print(f"Final: {os.path.getsize(output_path)} bytes")
print("=== Complete ===")