import bpy
import math

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

# Import the existing Miku GLB
input_path = "/home/ai-agent/Papay-Shooter/assets/characters/skins/miku_rigged_animated.glb"
output_path = "/home/ai-agent/Papay-Shooter/assets/characters/skins/miku_proper_anim.glb"

bpy.ops.import_scene.gltf(filepath=input_path)
print("Imported GLB")

# Find the armature
armatures = [obj for obj in bpy.data.objects if obj.type == 'ARMATURE']
if not armatures:
    print("ERROR: No armature found!")
    exit(1)

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

# Switch to pose mode
bpy.context.view_layer.objects.active = armature
bpy.ops.object.mode_set(mode='POSE')

def create_anim_dense(name, bone_rotations_func, length_frames=60):
    action = bpy.data.actions.new(name=name)
    armature.animation_data.action = action
    action.use_fake_user = True
    
    for frame in range(1, length_frames + 1):
        bpy.context.scene.frame_set(frame)
        rotations = bone_rotations_func(frame, length_frames)
        for bone_name, euler in rotations.items():
            if bone_name in armature.pose.bones:
                bone = armature.pose.bones[bone_name]
                bone.rotation_mode = 'XYZ'
                bone.rotation_euler = euler
                bone.keyframe_insert(data_path="rotation_euler", frame=frame)
    
    return action

# WALK: Per-frame keyframes (60 frames, 1 second)
def walk_fn(frame, length):
    t = frame / length * math.pi * 2
    return {
        "LeftUpperLeg": (math.sin(t) * 1.0, 0, 0),
        "RightUpperLeg": (math.sin(t + math.pi) * 1.0, 0, 0),
        "LeftUpperArm": (math.sin(t + math.pi) * 0.7, 0, 0),
        "RightUpperArm": (math.sin(t) * 0.7, 0, 0),
        "Spine": (0, math.sin(t) * 0.2, 0),
    }

create_anim_dense("Walk", walk_fn, 60)
print("Walk created")

# RUN: Faster (40 frames)
def run_fn(frame, length):
    t = frame / length * math.pi * 2
    return {
        "LeftUpperLeg": (math.sin(t) * 1.3, 0, 0),
        "RightUpperLeg": (math.sin(t + math.pi) * 1.3, 0, 0),
        "LeftUpperArm": (math.sin(t + math.pi) * 1.0, 0, 0),
        "RightUpperArm": (math.sin(t) * 1.0, 0, 0),
        "Spine": (0, math.sin(t) * 0.3, 0),
    }

create_anim_dense("Run", run_fn, 40)
print("Run created")

# JUMP: Single jump cycle (40 frames)
def jump_fn(frame, length):
    t = frame / length
    peak = max(0.0, 1.0 - abs(t - 0.5) * 3.0)
    return {
        "LeftUpperArm": (-2.5 * peak, 0, -0.3 * peak),
        "RightUpperArm": (-2.5 * peak, 0, 0.3 * peak),
        "LeftUpperLeg": (-0.5 * peak, 0, 0),
        "RightUpperLeg": (-0.5 * peak, 0, 0),
        "Spine": (0.15 * peak, 0, 0),
        "Hips": (0, 0.3 * peak, 0),
    }

create_anim_dense("Jump", jump_fn, 40)
print("Jump created")

# IDLE: Gentle breathing (60 frames)
def idle_fn(frame, length):
    t = frame / length * math.pi * 2
    return {
        "Spine": (0, math.sin(t) * 0.5, 0),
        "Head": (math.sin(t * 0.7) * 0.15, math.sin(t * 0.5) * 0.1, 0),
        "LeftUpperArm": (math.sin(t) * 0.2, 0, math.sin(t * 0.5) * 0.5),
        "RightUpperArm": (math.sin(t) * 0.2, 0, -math.sin(t * 0.5) * 0.5),
        "Hips": (0, math.sin(t) * 0.15, 0),
    }

create_anim_dense("Idle", idle_fn, 60)
print("Idle created")

# Remove original stripped actions, keep only our custom ones
for action in list(bpy.data.actions):
    if action.use_fake_user:
        # Rename Blender's auto-suffixed names
        if action.name == "Walk.001":
            action.name = "Walk"
        elif action.name == "Run.001":
            action.name = "Run"
        elif action.name == "Jump.001":
            action.name = "Jump"
        elif action.name == "Idle.002":
            action.name = "Idle"
    else:
        bpy.data.actions.remove(action)

print(f"Final actions: {[a.name for a in bpy.data.actions if a.use_fake_user]}")

# Export
bpy.context.scene.render.fps = 60
bpy.ops.object.mode_set(mode='OBJECT')

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,
    export_optimize_animation_size=False,
)

print(f"\nExported to {output_path}")
print("Done!")