"""
Blender 5.1 script: create walk + idle animations for matador_v02.
Run:  blender --background matador_v02.blend --python create_matador_anims.py
"""
import bpy, math, os


# ── setup ──────────────────────────────────────────────────────────────────────
arm_obj = next(o for o in bpy.data.objects if o.type == 'ARMATURE')
bpy.context.view_layer.objects.active = arm_obj
bpy.ops.object.mode_set(mode='POSE')

if arm_obj.animation_data is None:
    arm_obj.animation_data_create()

pb = arm_obj.pose.bones

# Ensure all animated bones use XYZ Euler
for bname in ['leg_L', 'leg_R', 'shin_L', 'shin_R',
              'arm_L', 'arm_R', 'forearm_L', 'forearm_R',
              'collarbone_L', 'collarbone_R']:
    pb[bname].rotation_mode = 'XYZ'

CYCLE = 30        # frames (1 s at 30 fps)
LEG_SWING  = 25   # degrees
KNEE_BEND  = 30
ARM_SWING  = 15
ARM_DOWN   = 55   # lower arms from T-pose around bone-local Y
ELBOW_BEND = 25


# ── WALK action ────────────────────────────────────────────────────────────────
# Remove old actions if they exist
for old_name in ['walk', 'idle']:
    if old_name in bpy.data.actions:
        bpy.data.actions.remove(bpy.data.actions[old_name])

walk = bpy.data.actions.new(name='walk')
walk.use_fake_user = True
walk.use_cyclic = True
arm_obj.animation_data.action = walk

# Clear all pose bone transforms first
bpy.ops.pose.select_all(action='SELECT')
bpy.ops.pose.rot_clear()

arm_down_l = math.radians(ARM_DOWN)
arm_down_r = math.radians(-ARM_DOWN)

for f in range(0, CYCLE + 1, 2):
    t = (f / CYCLE) * math.tau

    # Legs – swing forward/back on bone-local X
    pb['leg_L'].rotation_euler = (math.sin(t) * math.radians(LEG_SWING), 0, 0)
    pb['leg_L'].keyframe_insert(data_path='rotation_euler', frame=f)

    pb['leg_R'].rotation_euler = (math.sin(t + math.pi) * math.radians(LEG_SWING), 0, 0)
    pb['leg_R'].keyframe_insert(data_path='rotation_euler', frame=f)

    # Shins – only bend backward when leg is back
    val_l = max(0.0, -math.sin(t)) * math.radians(KNEE_BEND)
    val_r = max(0.0, -math.sin(t + math.pi)) * math.radians(KNEE_BEND)
    pb['shin_L'].rotation_euler = (val_l, 0, 0)
    pb['shin_L'].keyframe_insert(data_path='rotation_euler', frame=f)
    pb['shin_R'].rotation_euler = (val_r, 0, 0)
    pb['shin_R'].keyframe_insert(data_path='rotation_euler', frame=f)

    # Arms – lower from T-pose (Y) + swing (X, opposite phase to legs)
    swing_l = math.sin(t + math.pi) * math.radians(ARM_SWING)
    swing_r = math.sin(t) * math.radians(ARM_SWING)
    pb['arm_L'].rotation_euler = (swing_l, arm_down_l, 0)
    pb['arm_L'].keyframe_insert(data_path='rotation_euler', frame=f)
    pb['arm_R'].rotation_euler = (swing_r, arm_down_r, 0)
    pb['arm_R'].keyframe_insert(data_path='rotation_euler', frame=f)

    # Forearms – constant slight bend
    pb['forearm_L'].rotation_euler = (math.radians(ELBOW_BEND), 0, 0)
    pb['forearm_L'].keyframe_insert(data_path='rotation_euler', frame=f)
    pb['forearm_R'].rotation_euler = (math.radians(ELBOW_BEND), 0, 0)
    pb['forearm_R'].keyframe_insert(data_path='rotation_euler', frame=f)

    # Collarbones – identity
    pb['collarbone_L'].rotation_euler = (0, 0, 0)
    pb['collarbone_L'].keyframe_insert(data_path='rotation_euler', frame=f)
    pb['collarbone_R'].rotation_euler = (0, 0, 0)
    pb['collarbone_R'].keyframe_insert(data_path='rotation_euler', frame=f)

print("Created 'walk' action")


# ── IDLE action ────────────────────────────────────────────────────────────────
idle = bpy.data.actions.new(name='idle')
idle.use_fake_user = True
arm_obj.animation_data.action = idle

bpy.ops.pose.select_all(action='SELECT')
bpy.ops.pose.rot_clear()

# Arms lowered, everything else at rest – 2 frames for a valid action
for f in [0, 1]:
    pb['arm_L'].rotation_euler = (0, arm_down_l, 0)
    pb['arm_L'].keyframe_insert(data_path='rotation_euler', frame=f)
    pb['arm_R'].rotation_euler = (0, arm_down_r, 0)
    pb['arm_R'].keyframe_insert(data_path='rotation_euler', frame=f)
    pb['forearm_L'].rotation_euler = (math.radians(ELBOW_BEND), 0, 0)
    pb['forearm_L'].keyframe_insert(data_path='rotation_euler', frame=f)
    pb['forearm_R'].rotation_euler = (math.radians(ELBOW_BEND), 0, 0)
    pb['forearm_R'].keyframe_insert(data_path='rotation_euler', frame=f)
    for bname in ['collarbone_L', 'collarbone_R', 'leg_L', 'leg_R', 'shin_L', 'shin_R']:
        pb[bname].rotation_euler = (0, 0, 0)
        pb[bname].keyframe_insert(data_path='rotation_euler', frame=f)

print("Created 'idle' action")


# ── NLA tracks for export ─────────────────────────────────────────────────────
bpy.ops.object.mode_set(mode='OBJECT')

# Clear existing NLA tracks
if arm_obj.animation_data.nla_tracks:
    for t in list(arm_obj.animation_data.nla_tracks):
        arm_obj.animation_data.nla_tracks.remove(t)

tw = arm_obj.animation_data.nla_tracks.new()
tw.name = 'walk'
sw = tw.strips.new('walk', 0, walk)

ti = arm_obj.animation_data.nla_tracks.new()
ti.name = 'idle'
si = ti.strips.new('idle', 0, idle)

# ── export FBX ─────────────────────────────────────────────────────────────────
fbx_path = os.path.join(os.path.dirname(bpy.data.filepath), 'matador_v02.fbx')
bpy.ops.export_scene.fbx(
    filepath=fbx_path,
    use_selection=False,
    bake_anim=True,
    bake_anim_use_all_actions=True,
    bake_anim_use_nla_strips=False,
    bake_anim_force_startend_keying=True,
    bake_anim_step=1.0,
    bake_anim_simplify_factor=0.0,
    add_leaf_bones=False,
    primary_bone_axis='Y',
    secondary_bone_axis='X',
    apply_scale_options='FBX_SCALE_ALL',
)
print('Exported FBX to:', fbx_path)

bpy.ops.wm.save_as_mainfile(filepath=bpy.data.filepath)
print('Saved .blend')
