"""Combine the Meshy cook and free gait clips with the game's motion library.

Run in Blender 4.5. Requires the downloaded GLBs in source/meshy_v2.
No network access, credentials, generation, or gameplay changes happen here.
"""
import bpy
import json
import math
from pathlib import Path
from mathutils import Matrix, Quaternion, Vector

ROOT = Path(__file__).resolve().parents[2]
SOURCE = ROOT / 'tools/character-art/source/meshy_v2'
OUT = ROOT / 'assets/models/player/intro_cook_meshy'
OUT.mkdir(parents=True, exist_ok=True)
bpy.ops.wm.read_factory_settings(use_empty=True)
scene = bpy.context.scene
scene.render.fps = 24

def set_action(obj, action):
    obj.animation_data_create()
    obj.animation_data.action = action
    if action and action.slots:
        obj.animation_data.action_slot = action.slots[0]

def import_rig(path):
    before = set(bpy.data.objects)
    bpy.ops.import_scene.gltf(filepath=str(path))
    objects = set(bpy.data.objects) - before
    return next(o for o in objects if o.type == 'ARMATURE'), objects

rig, target_objects = import_rig(SOURCE / 'cook_rigged.glb')
rig.name = 'IntroCookMeshyRig'
body = next(o for o in target_objects if o.type == 'MESH' and o.name == 'char1')
body.name = 'IntroCookMeshyBody'
# Meshy's rig export includes an unskinned Icosphere helper. It must not inflate
# Godot's visual bounds or be rendered alongside the character.
for obj in target_objects - {rig, body}:
    bpy.data.objects.remove(obj, do_unlink=True)
rig.animation_data_clear()
for pb in rig.pose.bones:
    pb.location = (0,0,0); pb.rotation_mode = 'QUATERNION'
    pb.rotation_quaternion = Quaternion(); pb.scale = (1,1,1)

# Rigging exports only base color. Recover the original PBR material without
# modifying its bitmap textures; both assets retain the generated UV layout.
before = set(bpy.data.objects)
bpy.ops.import_scene.gltf(filepath=str(SOURCE / 'cook_meshy_raw.glb'))
raw_objects = set(bpy.data.objects) - before
raw_mesh = next(o for o in raw_objects if o.type == 'MESH')
assert len(raw_mesh.data.materials) == 1 and len(body.data.materials) == 1
body.data.materials[0] = raw_mesh.data.materials[0]
for obj in raw_objects:
    bpy.data.objects.remove(obj, do_unlink=True)

motion_sources = []
gaits = {}
for filename, clip in [('walking.glb','Walking'), ('running.glb','Running')]:
    source, objects = import_rig(SOURCE / filename)
    action = source.animation_data.action
    if action is None:
        action = source.animation_data.nla_tracks[0].strips[0].action
    for track in source.animation_data.nla_tracks:
        track.mute = True
    set_action(source, action)
    gaits[clip] = (source, action)
    motion_sources.append(source)
    for obj in objects - {source}:
        bpy.data.objects.remove(obj, do_unlink=True)

actions_before = set(bpy.data.actions)
with bpy.data.libraries.load(str(ROOT / 'tools/character-art/source/intro_cook.blend'), link=False) as (available, loaded):
    loaded.objects = ['IntroCookRig']
    loaded.actions = ['Idle','Hurt','Triple_Combo_Attack','Dead','Skill_01']
old = loaded.objects[0]
scene.collection.objects.link(old)
motion_sources.append(old)
old_actions = {a.name.split('.')[0]: a for a in loaded.actions}
for track in old.animation_data.nla_tracks:
    track.mute = True

# Anatomical rest axes explicitly compensate for A-pose versus T-pose. Applying
# only rest rotation deltas would leave the new character's arms outstretched.
NEXT = {'Hips':'Spine02','Spine02':'Spine01','Spine01':'Spine','Spine':'neck',
        'neck':'Head','Head':'head_end'}
for side in ['Left','Right']:
    for a,b in [('Shoulder','Arm'),('Arm','ForeArm'),('ForeArm','Hand'),
                ('UpLeg','Leg'),('Leg','Foot'),('Foot','ToeBase')]:
        NEXT[side+a] = side+b

def head_world(obj, name):
    return obj.matrix_world @ obj.data.bones[name].head_local

def anatomical_rest(obj, bone):
    start = head_world(obj, bone.name)
    if bone.name in NEXT:
        direction = head_world(obj, NEXT[bone.name]) - start
    elif bone.parent:
        direction = start - head_world(obj, bone.parent.name)
    else:
        direction = Vector((0,0,1))
    y = direction.normalized()
    front = Vector((0,-1,0))
    if abs(y.dot(front)) > .96:
        front = Vector((0,0,1))
    z = (front - y*y.dot(front)).normalized()
    x = y.cross(z).normalized()
    z = x.cross(y).normalized()
    return Matrix((x,y,z)).transposed().to_quaternion()

ordered = sorted(rig.pose.bones, key=lambda b: len(b.parent_recursive))
target_world_q = rig.matrix_world.to_quaternion()
target_rest = {b.name: (rig.matrix_world @ b.matrix_local).to_quaternion() for b in rig.data.bones}
target_anatomy = {b.name: anatomical_rest(rig,b) for b in rig.data.bones}
result_actions = {}
jobs = [(name,*gaits[name],1.0) for name in ['Walking','Running']]
jobs.append(('RunFast',*gaits['Running'],1.4))
jobs += [(name,old,old_actions[name],1.0) for name in ['Idle','Hurt','Triple_Combo_Attack','Dead','Skill_01']]
for name, source, original, tempo in jobs:
    set_action(source, original)
    src_rest = {b.name: (source.matrix_world @ b.matrix_local).to_quaternion() for b in source.data.bones}
    calibration = {b.name: anatomical_rest(source,source.data.bones[b.name]) @ target_anatomy[b.name].inverted()
                   for b in rig.data.bones}
    action = bpy.data.actions.new(name + '_Meshy')
    set_action(rig, action)
    start, end = original.frame_range
    duration = max(2, round((end-start)/tempo))
    first_height = None
    height_scale = head_world(rig,'Hips').z / head_world(source,'Hips').z
    for frame in range(duration+1):
        sample = start + (end-start)*frame/duration
        scene.frame_set(int(sample), subframe=sample-int(sample))
        source_height = (source.matrix_world @ source.pose.bones['Hips'].matrix).translation.z
        if first_height is None:
            first_height = source_height
        for pb in ordered:
            source_pose = (source.matrix_world @ source.pose.bones[pb.name].matrix).to_quaternion()
            delta = source_pose @ src_rest[pb.name].inverted()
            wanted = target_world_q.inverted() @ delta @ calibration[pb.name] @ target_rest[pb.name]
            if name in {'Idle','Hurt'} and (pb.name == 'Hips' or any(part in pb.name for part in ['UpLeg','Leg','Foot','ToeBase'])):
                # The earlier prototype used a different ankle/sole angle.
                # Keep this mesh's planted stance during breathing and recoil.
                wanted = target_world_q.inverted() @ target_rest[pb.name]
            if pb.parent:
                relative = pb.parent.bone.matrix_local.inverted() @ pb.bone.matrix_local
                basis = (pb.parent.matrix @ relative).to_quaternion()
            else:
                basis = pb.bone.matrix_local.to_quaternion()
            pb.rotation_mode = 'QUATERNION'
            pb.rotation_quaternion = basis.inverted() @ wanted
            pb.location = (0,0,0); pb.scale = (1,1,1)
            if pb.name == 'Hips':
                bob = Vector((0,0,(source_height-first_height)*height_scale))
                pb.location = pb.bone.matrix_local.to_3x3().inverted() @ rig.matrix_world.to_3x3().inverted() @ bob
            bpy.context.view_layer.update()
            pb.keyframe_insert('rotation_quaternion', frame=frame)
            pb.keyframe_insert('location', frame=frame)
            pb.keyframe_insert('scale', frame=frame)
    result_actions[name] = action

rig.animation_data.action = None
for obj in motion_sources:
    bpy.data.objects.remove(obj, do_unlink=True)
for action in list(bpy.data.actions):
    if action not in result_actions.values():
        bpy.data.actions.remove(action)
for name, action in result_actions.items():
    action.name = name
    action.use_fake_user = True
    track = rig.animation_data.nla_tracks.new(); track.name = name; track.mute = True
    strip = track.strips.new(name, 0, action); strip.action_slot = action.slots[0]
for pb in rig.pose.bones:
    pb.location = (0,0,0); pb.rotation_quaternion = Quaternion(); pb.scale = (1,1,1)
bpy.context.view_layer.update()
bpy.ops.object.select_all(action='DESELECT')
rig.select_set(True); body.select_set(True); bpy.context.view_layer.objects.active = rig
bpy.ops.export_scene.gltf(filepath=str(OUT/'intro_cook_meshy.glb'), export_format='GLB', use_selection=True,
    export_animations=True, export_animation_mode='NLA_TRACKS', export_nla_strips=True,
    export_force_sampling=True, export_frame_range=False, export_anim_slide_to_zero=True,
    export_materials='EXPORT', export_yup=True)

# Reusable, packed Blender scene with the final rig and animation library.
floor_mat = bpy.data.materials.new('Preview ground'); floor_mat.diffuse_color = (.16,.18,.21,1)
bpy.ops.mesh.primitive_plane_add(size=200)
floor = bpy.context.object; floor.name = 'PreviewGround'; floor.data.materials.append(floor_mat)
world = bpy.data.worlds.new('PreviewWorld'); world.use_nodes = True
world.node_tree.nodes['Background'].inputs[0].default_value = (.25,.3,.36,1)
world.node_tree.nodes['Background'].inputs[1].default_value = .45
scene.world = world
center = Vector((0,0,.92))
for name,loc,power in [('Key',(-3,-4,4),600),('Fill',(4,-2,3),300),('Rim',(1,3,4),700)]:
    light = bpy.data.lights.new(name,'AREA'); light.energy = power; light.size = 4
    obj = bpy.data.objects.new(name,light); scene.collection.objects.link(obj); obj.location = loc
    obj.rotation_euler = (center-obj.location).to_track_quat('-Z','Y').to_euler()
camera_data = bpy.data.cameras.new('PreviewCamera'); camera = bpy.data.objects.new('PreviewCamera',camera_data)
scene.collection.objects.link(camera); camera.location = (2.6,-6,2.45)
camera.rotation_euler = (center-camera.location).to_track_quat('-Z','Y').to_euler()
camera_data.type = 'ORTHO'; camera_data.ortho_scale = 2.15; scene.camera = camera
scene.render.engine = 'CYCLES'; scene.cycles.samples = 24
scene.render.resolution_x = 900; scene.render.resolution_y = 1000; scene.render.resolution_percentage = 100
scene.render.image_settings.file_format = 'PNG'; scene.view_settings.view_transform = 'AgX'
set_action(rig,result_actions['Idle']); scene.frame_set(0); scene.frame_start = 0; scene.frame_end = 48
bpy.ops.file.pack_all()
bpy.ops.wm.save_as_mainfile(filepath=str(SOURCE/'intro_cook_meshy.blend'))
for name,phase in [('Idle',0),('Walking',.25),('RunFast',.25),('Triple_Combo_Attack',.25)]:
    set_action(rig,result_actions[name]); scene.frame_set(round(result_actions[name].frame_range[1]*phase))
    scene.render.filepath = str(SOURCE/(name+'_preview.png'))
    bpy.ops.render.render(write_still=True)
report = {'vertices':len(body.data.vertices),'polygons':len(body.data.polygons),'bones':len(rig.data.bones),
    'animations':list(result_actions),'generation_task':'01a0aa62-ac21-7582-b495-35755ded3c73',
    'rig_task':'01a0aa67-ce2b-70d7-9cec-dd2629be00f8','meshy_credits_used':35,
    'reference':'res://assets/story/intro/08_omelette_single.png'}
(OUT/'build_manifest.json').write_text(json.dumps(report,indent=2)+'\n',encoding='utf-8')
print('MESHY_COOK_BUILD',json.dumps(report))