#!/usr/bin/env python3
"""
Auto-rig a humanoid mesh in Blender and add animations.
Handles multi-part meshes (body, hair, accessories) by separating loose parts,
rigging only the body, and parenting everything to the same skeleton.
Exports as GLB with proper skeleton for Godot.
"""
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:
    print("Usage: blender --background --python rig_and_animate.py -- <input.glb> <output.glb>")
    sys.exit(1)

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

# Import the mesh
print(f"Importing {input_path}...")
if input_path.endswith('.glb') or input_path.endswith('.gltf'):
    bpy.ops.import_scene.gltf(filepath=input_path)
elif input_path.endswith('.obj'):
    bpy.ops.import_scene.obj(filepath=input_path, use_split_objects=False)
elif input_path.endswith('.fbx'):
    bpy.ops.import_scene.fbx(filepath=input_path)

imported = bpy.context.selected_objects
print(f"Imported {len(imported)} objects")

# Find all mesh objects
meshes = [o for o in imported if o.type == 'MESH' and o.name != 'Icosphere']
if not meshes:
    print("ERROR: No mesh found")
    sys.exit(1)

# Join all meshes into one first
if len(meshes) > 1:
    bpy.context.view_layer.objects.active = meshes[0]
    bpy.ops.object.join()
    meshes = [bpy.context.active_object]

mesh_obj = meshes[0]
print(f"Combined mesh: {mesh_obj.name} ({len(mesh_obj.data.vertices)} verts)")

# Clean up mesh
bpy.context.view_layer.objects.active = mesh_obj
bpy.ops.object.mode_set(mode='EDIT')
bpy.ops.mesh.select_all(action='SELECT')
bpy.ops.mesh.remove_doubles()
bpy.ops.mesh.normals_make_consistent(inside=False)
bpy.ops.object.mode_set(mode='OBJECT')

# Center and scale
bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS')
mesh_obj.location = (0, 0, 0)

dims = mesh_obj.dimensions
max_dim = max(dims.x, dims.y, dims.z)
print(f"Original dimensions: {dims}")
if max_dim > 0 and max_dim < 10:
    target_height = 1.8
    scale = target_height / max_dim
    mesh_obj.scale = (scale, scale, scale)
    bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)
    print(f"Scaled by {scale:.2f}x. New dims: {mesh_obj.dimensions}")

# Separate mesh into loose parts (body, hair, twintails, etc.)
# This is key: we need to identify which vertices belong to the body vs hair/accessories
print("\n=== Separating mesh into parts ===")
bpy.context.view_layer.objects.active = mesh_obj
bpy.ops.object.mode_set(mode='EDIT')
bpy.ops.mesh.select_all(action='SELECT')
bpy.ops.mesh.separate(type='LOOSE')
bpy.ops.object.mode_set(mode='OBJECT')

# List separated parts
parts = [o for o in bpy.context.selected_objects if o.type == 'MESH']
print(f"Separated into {len(parts)} parts:")
for p in parts:
    print(f"  {p.name}: {len(p.data.vertices)} verts, dims={p.dimensions}")

# Identify the body part (largest by vertex count, roughly humanoid proportions)
# The body should be the largest part with roughly humanoid proportions
body_part = None
hair_parts = []
other_parts = []

for p in parts:
    verts = len(p.data.vertices)
    dims = p.dimensions
    height = max(dims.x, dims.y, dims.z)
    width = min(dims.x, dims.y, dims.z)
    
    # Body: largest part, height > width, reasonable humanoid proportions
    if height > 0.5 and verts > 1000:
        if body_part is None or verts > len(body_part.data.vertices):
            if body_part:
                other_parts.append(body_part)
            body_part = p
        else:
            other_parts.append(p)
    else:
        hair_parts.append(p)

if not body_part:
    # Fallback: just use the largest part
    body_part = max(parts, key=lambda p: len(p.data.vertices))
    hair_parts = [p for p in parts if p != body_part]
    other_parts = []

print(f"\nBody: {body_part.name} ({len(body_part.data.vertices)} verts)")
print(f"Hair/accessories: {[p.name for p in hair_parts]}")
print(f"Other: {[p.name for p in other_parts]}")

# Delete non-body parts (hair, twintails, accessories) - they'll be parented to bones later
# Actually, keep them but we won't skin them to the body bones
all_non_body = hair_parts + other_parts

# Create the armature
print("\n=== Creating Rig ===")
bpy.ops.object.armature_add(enter_editmode=True, location=(0, 0, 0))
armature = bpy.context.active_object
armature.name = "Rig"
armature.data.display_type = 'STICK'
armature.data.show_axes = False

# Delete default bone
for bone in armature.data.edit_bones:
    armature.data.edit_bones.remove(bone)

# Define humanoid skeleton (Mixamo-compatible names)
# Blender is Z-up, so head is at higher Z
bone_defs = [
    ("mixamorig:Hips",         (0.09, 0.0, 0.95),   (0.09, 0.0, 1.05)),
    ("mixamorig:Spine",        (0.09, 0.0, 1.05),   (0.09, 0.0, 1.20)),
    ("mixamorig:Spine1",       (0.09, 0.0, 1.20),   (0.09, 0.0, 1.35)),
    ("mixamorig:Spine2",       (0.09, 0.0, 1.35),   (0.09, 0.0, 1.45)),
    ("mixamorig:Neck",         (0.09, 0.0, 1.45),   (0.09, 0.0, 1.52)),
    ("mixamorig:Head",         (0.09, 0.0, 1.52),   (0.09, 0.0, 1.70)),
    ("mixamorig:LeftArm",      (-0.09, 0.0, 1.40),  (-0.18, 0.0, 1.22)),
    ("mixamorig:LeftForeArm",  (-0.18, 0.0, 1.22),  (-0.18, 0.0, 1.00)),
    ("mixamorig:LeftHand",     (-0.18, 0.0, 1.00),  (-0.18, 0.0, 0.78)),
    ("mixamorig:RightArm",     (0.27, 0.0, 1.40),   (0.18, 0.0, 1.22)),
    ("mixamorig:RightForeArm", (0.18, 0.0, 1.22),   (0.18, 0.0, 1.00)),
    ("mixamorig:RightHand",    (0.18, 0.0, 1.00),   (0.18, 0.0, 0.78)),
    ("mixamorig:LeftUpLeg",    (0.0, 0.0, 0.95),    (0.0, 0.0, 0.50)),
    ("mixamorig:LeftLeg",      (0.0, 0.0, 0.50),    (0.0, 0.0, 0.10)),
    ("mixamorig:LeftFoot",     (0.0, 0.0, 0.10),    (0.0, 0.0, 0.0)),
    ("mixamorig:RightUpLeg",   (0.18, 0.0, 0.95),   (0.18, 0.0, 0.50)),
    ("mixamorig:RightLeg",     (0.18, 0.0, 0.50),   (0.18, 0.0, 0.10)),
    ("mixamorig:RightFoot",    (0.18, 0.0, 0.10),   (0.18, 0.0, 0.0)),
]

bones = {}
for name, head, tail in bone_defs:
    bone = armature.data.edit_bones.new(name)
    bone.head = head
    bone.tail = tail
    bones[name] = bone

# Set parent hierarchy
bones["mixamorig:Spine"].parent = bones["mixamorig:Hips"]
bones["mixamorig:Spine1"].parent = bones["mixamorig:Spine"]
bones["mixamorig:Spine2"].parent = bones["mixamorig:Spine1"]
bones["mixamorig:Neck"].parent = bones["mixamorig:Spine2"]
bones["mixamorig:Head"].parent = bones["mixamorig:Neck"]
bones["mixamorig:LeftArm"].parent = bones["mixamorig:Spine2"]
bones["mixamorig:LeftForeArm"].parent = bones["mixamorig:LeftArm"]
bones["mixamorig:LeftHand"].parent = bones["mixamorig:LeftForeArm"]
bones["mixamorig:RightArm"].parent = bones["mixamorig:Spine2"]
bones["mixamorig:RightForeArm"].parent = bones["mixamorig:RightArm"]
bones["mixamorig:RightHand"].parent = bones["mixamorig:RightForeArm"]
bones["mixamorig:LeftUpLeg"].parent = bones["mixamorig:Hips"]
bones["mixamorig:LeftLeg"].parent = bones["mixamorig:LeftUpLeg"]
bones["mixamorig:LeftFoot"].parent = bones["mixamorig:LeftLeg"]
bones["mixamorig:RightUpLeg"].parent = bones["mixamorig:Hips"]
bones["mixamorig:RightLeg"].parent = bones["mixamorig:RightUpLeg"]
bones["mixamorig:RightFoot"].parent = bones["mixamorig:RightLeg"]

for name in ["mixamorig:Spine", "mixamorig:Spine1", "mixamorig:Spine2", "mixamorig:Neck",
             "mixamorig:Head", "mixamorig:LeftArm", "mixamorig:LeftForeArm", "mixamorig:LeftHand",
             "mixamorig:RightArm", "mixamorig:RightForeArm", "mixamorig:RightHand",
             "mixamorig:LeftUpLeg", "mixamorig:LeftLeg", "mixamorig:LeftFoot",
             "mixamorig:RightUpLeg", "mixamorig:RightLeg", "mixamorig:RightFoot"]:
    bones[name].use_connect = True

print(f"Created {len(bones)} bones")
bpy.ops.object.mode_set(mode='OBJECT')

# Skin ONLY the body part to the armature with automatic weights
print("\n=== Skinning body mesh ===")
bpy.ops.object.select_all(action='DESELECT')
body_part.select_set(True)
armature.select_set(True)
bpy.context.view_layer.objects.active = armature
bpy.ops.object.parent_set(type='ARMATURE_AUTO')
print(f"Body '{body_part.name}' parented with automatic weights")

# Parent hair/accessories to the armature (no skinning, just follow)
# They'll move with the Head bone
print("\n=== Parenting hair/accessories ===")
for part in all_non_body:
    # Find the closest bone to this part's center
    part_center = part.location
    closest_bone = "mixamorig:Head"  # Default to head for hair
    closest_dist = float('inf')
    
    for bone_name, bone in bones.items():
        bone_head = armature.matrix_world @ bone.head
        dist = (part_center - bone_head).length
        if dist < closest_dist:
            closest_dist = dist
            closest_bone = bone_name
    
    # Parent to armature with empty weights (just follow the bone)
    bpy.ops.object.select_all(action='DESELECT')
    part.select_set(True)
    armature.select_set(True)
    bpy.context.view_layer.objects_active = armature
    bpy.ops.object.parent_set(type='ARMATURE_NAME', keep_transform=True)
    print(f"  '{part.name}' -> parented to '{closest_bone}'")

# Now create animations
print("\n=== Creating animations ===")
bpy.context.view_layer.objects.active = armature
bpy.ops.object.mode_set(mode='POSE')

def set_bone_rest(armature):
    for bone in armature.pose.bones:
        bone.location = (0, 0, 0)
        bone.rotation_euler = (0, 0, 0)
        bone.scale = (1, 1, 1)

def key_bone(bone_name, frame, loc=None, rot=None):
    bone = armature.pose.bones.get(bone_name)
    if not bone:
        return
    if loc:
        bone.location = loc
        bone.keyframe_insert(data_path="location", frame=frame)
    if rot:
        bone.rotation_mode = 'XYZ'
        bone.rotation_euler = rot
        bone.keyframe_insert(data_path="rotation_euler", frame=frame)

# Idle animation
print("  Idle...")
if not armature.animation_data:
    armature.animation_data_create()

action_idle = bpy.data.actions.new(name="Idle")
armature.animation_data.action = action_idle
action_idle.frame_range = (1, 30)
set_bone_rest(armature)

key_bone("mixamorig:Spine", 1, rot=(0.02, 0, 0))
key_bone("mixamorig:Spine", 15, rot=(-0.02, 0, 0))
key_bone("mixamorig:Spine", 30, rot=(0.02, 0, 0))
key_bone("mixamorig:LeftArm", 1, rot=(0, 0, 0.05))
key_bone("mixamorig:LeftArm", 15, rot=(0, 0, -0.05))
key_bone("mixamorig:LeftArm", 30, rot=(0, 0, 0.05))
key_bone("mixamorig:RightArm", 1, rot=(0, 0, -0.05))
key_bone("mixamorig:RightArm", 15, rot=(0, 0, 0.05))
key_bone("mixamorig:RightArm", 30, rot=(0, 0, -0.05))

# Walk animation
print("  Walk...")
action_walk = bpy.data.actions.new(name="Walk")
armature.animation_data.action = action_walk
action_walk.frame_range = (1, 24)
set_bone_rest(armature)

for f in range(1, 25):
    t = (f - 1) / 24
    phase = t * 2 * math.pi
    key_bone("mixamorig:Hips", f, loc=(0, 0, abs(math.sin(phase * 2)) * 0.01))
    key_bone("mixamorig:LeftUpLeg", f, rot=(math.sin(phase) * 0.4, 0, 0))
    key_bone("mixamorig:LeftLeg", f, rot=(max(0, -math.sin(phase) * 0.3 + 0.2), 0, 0))
    key_bone("mixamorig:RightUpLeg", f, rot=(math.sin(phase + math.pi) * 0.4, 0, 0))
    key_bone("mixamorig:RightLeg", f, rot=(max(0, -math.sin(phase + math.pi) * 0.3 + 0.2), 0, 0))
    key_bone("mixamorig:LeftArm", f, rot=(-math.sin(phase) * 0.3, 0, 0.05))
    key_bone("mixamorig:LeftForeArm", f, rot=(-0.3 + max(0, math.sin(phase) * 0.2), 0, 0))
    key_bone("mixamorig:RightArm", f, rot=(-math.sin(phase + math.pi) * 0.3, 0, -0.05))
    key_bone("mixamorig:RightForeArm", f, rot=(-0.3 + max(0, math.sin(phase + math.pi) * 0.2), 0, 0))
    key_bone("mixamorig:Spine", f, rot=(0, math.sin(phase) * 0.05, 0))

# Run animation
print("  Run...")
action_run = bpy.data.actions.new(name="Run")
armature.animation_data.action = action_run
action_run.frame_range = (1, 20)
set_bone_rest(armature)

for f in range(1, 21):
    t = (f - 1) / 20
    phase = t * 2 * math.pi
    key_bone("mixamorig:Hips", f, loc=(0, 0, abs(math.sin(phase * 2)) * 0.03))
    key_bone("mixamorig:LeftUpLeg", f, rot=(math.sin(phase) * 0.8, 0, 0))
    key_bone("mixamorig:LeftLeg", f, rot=(max(0.1, -math.sin(phase) * 0.6 + 0.3), 0, 0))
    key_bone("mixamorig:RightUpLeg", f, rot=(math.sin(phase + math.pi) * 0.8, 0, 0))
    key_bone("mixamorig:RightLeg", f, rot=(max(0.1, -math.sin(phase + math.pi) * 0.6 + 0.3), 0, 0))
    key_bone("mixamorig:LeftArm", f, rot=(-math.sin(phase) * 0.7, 0, 0.1))
    key_bone("mixamorig:LeftForeArm", f, rot=(-1.0 + max(0, math.sin(phase) * 0.3), 0, 0))
    key_bone("mixamorig:RightArm", f, rot=(-math.sin(phase + math.pi) * 0.7, 0, -0.1))
    key_bone("mixamorig:RightForeArm", f, rot=(-1.0 + max(0, math.sin(phase + math.pi) * 0.3), 0, 0))
    key_bone("mixamorig:Spine", f, rot=(0.1, math.sin(phase) * 0.08, 0))

# Jump animation
print("  Jump...")
action_jump = bpy.data.actions.new(name="Jump")
armature.animation_data.action = action_jump
action_jump.frame_range = (1, 20)
set_bone_rest(armature)

for f in range(1, 6):
    t = (f - 1) / 4
    key_bone("mixamorig:Hips", f, loc=(0, 0, -t * 0.2))
    key_bone("mixamorig:LeftUpLeg", f, rot=(-t * 0.6, 0, 0))
    key_bone("mixamorig:LeftLeg", f, rot=(t * 1.0, 0, 0))
    key_bone("mixamorig:RightUpLeg", f, rot=(-t * 0.6, 0, 0))
    key_bone("mixamorig:RightLeg", f, rot=(t * 1.0, 0, 0))
    key_bone("mixamorig:LeftArm", f, rot=(0, 0, -t * 0.5))
    key_bone("mixamorig:RightArm", f, rot=(0, 0, t * 0.5))

for f in range(6, 9):
    t = (f - 6) / 2
    key_bone("mixamorig:Hips", f, loc=(0, 0, -0.2 + t * 0.2))
    key_bone("mixamorig:LeftUpLeg", f, rot=(-0.6 + t * 0.6, 0, 0))
    key_bone("mixamorig:LeftLeg", f, rot=(1.0 - t * 0.3, 0, 0))
    key_bone("mixamorig:RightUpLeg", f, rot=(-0.6 + t * 0.6, 0, 0))
    key_bone("mixamorig:RightLeg", f, rot=(1.0 - t * 0.3, 0, 0))
    key_bone("mixamorig:LeftArm", f, rot=(-t * 2.0, 0, -0.5 - t))
    key_bone("mixamorig:RightArm", f, rot=(-t * 2.0, 0, 0.5 + t))

for f in range(9, 16):
    key_bone("mixamorig:LeftArm", f, rot=(-2.0, 0, -1.5))
    key_bone("mixamorig:RightArm", f, rot=(-2.0, 0, 1.5))
    key_bone("mixamorig:LeftForeArm", f, rot=(-0.2, 0, 0))
    key_bone("mixamorig:RightForeArm", f, rot=(-0.2, 0, 0))
    key_bone("mixamorig:LeftUpLeg", f, rot=(0.2, 0, 0))
    key_bone("mixamorig:LeftLeg", f, rot=(0.5, 0, 0))
    key_bone("mixamorig:RightUpLeg", f, rot=(0.1, 0, 0))
    key_bone("mixamorig:RightLeg", f, rot=(0.3, 0, 0))

for f in range(16, 21):
    t = (f - 16) / 4
    key_bone("mixamorig:Hips", f, loc=(0, 0, -t * 0.2))
    key_bone("mixamorig:LeftUpLeg", f, rot=(-t * 0.5, 0, 0))
    key_bone("mixamorig:LeftLeg", f, rot=(t * 0.8, 0, 0))
    key_bone("mixamorig:RightUpLeg", f, rot=(-t * 0.5, 0, 0))
    key_bone("mixamorig:RightLeg", f, rot=(t * 0.8, 0, 0))
    key_bone("mixamorig:LeftArm", f, rot=(-2.0 + t * 2.0, 0, -1.5 + t * 1.5))
    key_bone("mixamorig:RightArm", f, rot=(-2.0 + t * 2.0, 0, 1.5 - t * 1.5))

# Crouch animation
print("  Crouch...")
action_crouch = bpy.data.actions.new(name="Crouch")
armature.animation_data.action = action_crouch
action_crouch.frame_range = (1, 1)
set_bone_rest(armature)

key_bone("mixamorig:Hips", 1, loc=(0, 0, -0.4))
key_bone("mixamorig:Spine", 1, rot=(0.3, 0, 0))
key_bone("mixamorig:LeftUpLeg", 1, rot=(-0.8, 0, 0))
key_bone("mixamorig:LeftLeg", 1, rot=(1.4, 0, 0))
key_bone("mixamorig:RightUpLeg", 1, rot=(-0.8, 0, 0))
key_bone("mixamorig:RightLeg", 1, rot=(1.4, 0, 0))
key_bone("mixamorig:LeftArm", 1, rot=(-0.5, 0, 0.2))
key_bone("mixamorig:RightArm", 1, rot=(-0.5, 0, -0.2))

# Death animation
print("  Death...")
action_death = bpy.data.actions.new(name="Death")
armature.animation_data.action = action_death
action_death.frame_range = (1, 30)
set_bone_rest(armature)

for f in range(1, 15):
    t = (f - 1) / 13
    key_bone("mixamorig:Hips", f, loc=(0, 0, -t * 0.5), rot=(-t * 1.2, 0, t * 0.2))
    key_bone("mixamorig:Spine", f, rot=(-t * 0.5, t * 0.3, 0))
    key_bone("mixamorig:Spine2", f, rot=(-t * 0.4, 0, 0))
    key_bone("mixamorig:LeftArm", f, rot=(t * 2.0, 0, -t * 0.5))
    key_bone("mixamorig:RightArm", f, rot=(t * 2.0, 0, t * 0.5))
    key_bone("mixamorig:LeftUpLeg", f, rot=(-t * 0.3, 0, 0))
    key_bone("mixamorig:RightUpLeg", f, rot=(-t * 0.3, 0, 0))

for f in range(15, 31):
    t = (f - 15) / 15
    key_bone("mixamorig:Hips", f, loc=(0, 0, -0.5 - t * 0.2), rot=(-1.2 - t * 0.5, 0.8 * (1 - t), 0.2 + t * 0.3))
    key_bone("mixamorig:Spine", f, rot=(-0.7 - t * 0.3, 0.5 * (1 - t), 0))
    key_bone("mixamorig:Spine2", f, rot=(-0.6 - t * 0.2, 0, 0))
    key_bone("mixamorig:LeftHand", f, rot=(t * 0.3, t * 0.5, -0.7 - t * 0.5))
    key_bone("mixamorig:RightHand", f, rot=(t * 0.4, -t * 0.6, 0.8 + t * 0.4))

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

# Switch to object mode
bpy.ops.object.mode_set(mode='OBJECT')

# Select all for export
bpy.ops.object.select_all(action='DESELECT')
armature.select_set(True)
body_part.select_set(True)
for part in all_non_body:
    part.select_set(True)
bpy.context.view_layer.objects.active = armature

# Export as GLB
print(f"\n=== Exporting 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")

# Patch skeleton reference
import struct
import 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') == 'Rig':
        skeleton_idx = i
        break

if skeleton_idx is not None:
    skins = gltf.get('skins', [])
    if skins:
        skins[0]['skeleton'] = skeleton_idx
        print(f"Set skeleton to {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, {len(gltf.get('animations', []))} animations")
print("=== Complete ===")