import bpy, math, os, sys
from mathutils import Vector
from pathlib import Path

ROOT=Path(r"D:/UD/AI/AIGame")
OUT=ROOT/"Tools/output/companions/BudBud"
ASSET=ROOT/"AIClient/Assets/Bundles/Characters/Animal/BudBud"
# clean
bpy.ops.object.select_all(action='SELECT'); bpy.ops.object.delete(use_global=False)
for datablocks in (bpy.data.meshes, bpy.data.curves, bpy.data.armatures, bpy.data.materials, bpy.data.cameras, bpy.data.lights):
    pass

# one atlas material; Unity receives BudBud_base/mix via SetupCharacter later
mat=bpy.data.materials.new('BudBud_Atlas'); mat.use_nodes=True
nodes=mat.node_tree.nodes; links=mat.node_tree.links
for n in list(nodes): nodes.remove(n)
outn=nodes.new('ShaderNodeOutputMaterial'); bs=nodes.new('ShaderNodeBsdfPrincipled'); tex=nodes.new('ShaderNodeTexImage')
try:
    img=bpy.data.images.load(str(ASSET/'BudBud_base.png'), check_existing=True); tex.image=img
except Exception: pass
links.new(tex.outputs['Color'],bs.inputs['Base Color']); links.new(bs.outputs['BSDF'],outn.inputs['Surface'])

# category -> UV atlas swatch center
UV={'green':(0.18,0.82),'cream':(0.50,0.82),'dark':(0.82,0.82),'pink':(0.18,0.32),'yellow':(0.50,0.32),'blue':(0.82,0.32)}

def uv_color(obj, cat):
    if not obj.data.uv_layers: obj.data.uv_layers.new(name='UVMap')
    u,v=UV[cat]
    for loop in obj.data.uv_layers.active.data:
        loop.uv=(u,v)
    obj.data.materials.append(mat)

def sphere(name, loc, scale, cat, bone, seg=10, rings=6):
    bpy.ops.mesh.primitive_uv_sphere_add(segments=seg, ring_count=rings, location=loc)
    o=bpy.context.object; o.name=name; o.scale=scale
    bpy.ops.object.transform_apply(location=False,rotation=False,scale=True)
    uv_color(o,cat)
    vg=o.vertex_groups.new(name=bone); vg.add(list(range(len(o.data.vertices))),1.0,'REPLACE')
    return o

def cone(name, loc, radius1, radius2, depth, rot, cat, bone, verts=10):
    bpy.ops.mesh.primitive_cone_add(vertices=verts, radius1=radius1, radius2=radius2, depth=depth, location=loc, rotation=rot)
    o=bpy.context.object; o.name=name; uv_color(o,cat)
    vg=o.vertex_groups.new(name=bone); vg.add(list(range(len(o.data.vertices))),1.0,'REPLACE'); return o

def create_rig():
    data=bpy.data.armatures.new('BudBud_Armature'); arm=bpy.data.objects.new('Armature',data); bpy.context.collection.objects.link(arm)
    bpy.context.view_layer.objects.active=arm; arm.select_set(True); bpy.ops.object.mode_set(mode='EDIT')
    def eb(name,head,tail,parent=None):
        b=data.edit_bones.new(name); b.head=head; b.tail=tail
        if parent: b.parent=data.edit_bones.get(parent)
        return b
    eb('bone_0',(0,0.02,0),(0,0.18,0))
    eb('bone_1',(0,0.18,0),(0,0.42,0),'bone_0')
    eb('bone_2',(0,0.42,0),(0,0.62,0),'bone_1')
    eb('bone_3',(-0.18,0.36,0),(-0.30,0.36,0),'bone_1')
    eb('bone_4',(0.18,0.36,0),(0.30,0.36,0),'bone_1')
    eb('bone_5',(-0.10,0.18,0),(-0.10,0.03,0),'bone_0')
    eb('bone_6',(0.10,0.18,0),(0.10,0.03,0),'bone_0')
    eb('bone_7',(0,0.57,0),(0,0.78,0),'bone_2')
    bpy.ops.object.mode_set(mode='OBJECT'); arm.show_in_front=True
    return arm

arm=create_rig()
parts=[]
# body and head; front is -Z
parts += [sphere('Body',(0,0.29,0),(0.27,0.28,0.22),'green','bone_1',10,6)]
parts += [sphere('Head',(0,0.52,-0.005),(0.29,0.24,0.23),'green','bone_2',10,6)]
# cream muzzle and belly
parts += [sphere('Muzzle',(0,0.48,-0.205),(0.13,0.095,0.045),'cream','bone_2',10,6)]
parts += [sphere('Belly',(0,0.27,-0.205),(0.16,0.18,0.035),'cream','bone_1',10,6)]
# eyes and pupils
parts += [sphere('EyeL',(-0.105,0.565,-0.205),(0.052,0.068,0.035),'cream','bone_2',10,6),sphere('EyeR',(0.105,0.565,-0.205),(0.052,0.068,0.035),'cream','bone_2',10,6)]
parts += [sphere('PupilL',(-0.105,0.565,-0.237),(0.026,0.036,0.015),'dark','bone_2',8,6),sphere('PupilR',(0.105,0.565,-0.237),(0.026,0.036,0.015),'dark','bone_2',8,6)]
# nose
parts += [sphere('Nose',(0,0.505,-0.253),(0.028,0.021,0.016),'pink','bone_2',8,6)]
# leaf crown (broad petals)
parts += [sphere('LeafTop',(0,0.73,0),(0.14,0.12,0.07),'green','bone_7',10,6)]
for i,x in enumerate((-0.105,0.105)):
    p=sphere('LeafSide'+str(i),(x,0.70,0),(0.105,0.17,0.055),'green','bone_7',10,6); p.rotation_euler[2]=(-0.38 if x<0 else 0.38); parts.append(p)
# arms, legs and cream paws
parts += [sphere('ArmL',(-0.285,0.34,0),(0.075,0.16,0.075),'green','bone_3',10,6),sphere('ArmR',(0.285,0.34,0),(0.075,0.16,0.075),'green','bone_4',10,6)]
parts += [sphere('LegL',(-0.105,0.11,0),(0.09,0.14,0.10),'green','bone_5',10,6),sphere('LegR',(0.105,0.11,0),(0.09,0.14,0.10),'green','bone_6',10,6)]
parts += [sphere('FootL',(-0.105,0.035,-0.055),(0.12,0.045,0.13),'cream','bone_5',10,6),sphere('FootR',(0.105,0.035,-0.055),(0.12,0.045,0.13),'cream','bone_6',10,6)]
# armature modifiers
for o in parts:
    mod=o.modifiers.new('Armature','ARMATURE'); mod.object=arm
# front flower accent (small yellow center on leaf crown)
parts += [sphere('Flower',(0,0.77,-0.03),(0.055,0.055,0.035),'yellow','bone_7',8,5)]
mod=parts[-1].modifiers.new('Armature','ARMATURE'); mod.object=arm
# Make armature active only for export
# Keep the exported mesh grounded while preserving the same bind pose for every clip.
for _o in parts + [arm]: _o.location.y += 0.03

def select_only(objs):
    bpy.ops.object.select_all(action='DESELECT')
    for o in objs:o.select_set(True)
    bpy.context.view_layer.objects.active=objs[0]

def export(path, objs, action=None, anim=False):
    select_only(objs)
    bpy.ops.export_scene.fbx(filepath=str(path), use_selection=True, object_types={'ARMATURE','MESH'} if any(o.type=='MESH' for o in objs) else {'ARMATURE'}, add_leaf_bones=False, bake_space_transform=False, axis_forward='-Z', axis_up='Y', apply_unit_scale=True, bake_anim=anim, bake_anim_use_all_actions=False, bake_anim_use_nla_strips=False, bake_anim_force_startend_keying=True, primary_bone_axis='Y', secondary_bone_axis='X')

# bind rigged
export(OUT/'BudBud_Rigged.fbx',[arm]+parts,anim=False)
# create and export source clips with same armature skeleton

def clear_actions(): arm.animation_data_clear()
def pose_action(name, seconds):
    act=bpy.data.actions.new(name); arm.animation_data_create(); arm.animation_data.action=act
    # bind pose at t=0 and end; 30fps
    for b in arm.pose.bones: b.rotation_mode='XYZ'
    frames=[0, int(seconds*15), int(seconds*30)]
    for f in frames:
        t=f/30.0
        for b in arm.pose.bones:
            b.location=(0,0,0); b.rotation_euler=(0,0,0); b.scale=(1,1,1)
        if name=='Idle':
            s=math.sin(math.pi*t/seconds)
            arm.pose.bones['bone_1'].scale=(1+0.012*s,1+0.018*s,1+0.012*s)
            arm.pose.bones['bone_7'].rotation_euler[1]=0.12*s
        elif name=='Walk':
            ph=math.sin(2*math.pi*t/seconds)
            arm.pose.bones['bone_5'].rotation_euler[0]=0.28*ph; arm.pose.bones['bone_6'].rotation_euler[0]=-0.28*ph
            arm.pose.bones['bone_3'].rotation_euler[0]=-0.20*ph; arm.pose.bones['bone_4'].rotation_euler[0]=0.20*ph
            arm.pose.bones['bone_0'].location[1]=0.012*(1-math.cos(2*math.pi*t/seconds))
        else:
            u=max(0,min(1,t/seconds)); strike=math.sin(math.pi*u)**2
            arm.pose.bones['bone_3'].rotation_euler[0]=math.radians(-55)*strike; arm.pose.bones['bone_4'].rotation_euler[0]=math.radians(-55)*strike
            arm.pose.bones['bone_3'].location[2]=-0.025*strike; arm.pose.bones['bone_4'].location[2]=-0.025*strike
            arm.pose.bones['bone_7'].rotation_euler[2]=0.18*strike
        for b in arm.pose.bones:
            b.keyframe_insert('location',frame=f); b.keyframe_insert('rotation_euler',frame=f); b.keyframe_insert('scale',frame=f)
    act.frame_start=0; act.frame_end=frames[-1]
    return act
for name in ('Idle','Walk','Attack'):
    clear_actions(); act=pose_action(name,2.0); export(OUT/'motions'/name/'animated.fbx',[arm],anim=True); act=None
# static export by remove armature modifiers and armature
for o in parts:
    if o.type=='MESH':
        for m in list(o.modifiers): o.modifiers.remove(m)
arm.hide_render=True; arm.hide_viewport=True
export(ASSET/'BudBud.fbx',parts,anim=False)
# also copy rigged into output source location handled above
# preview render
bpy.ops.object.select_all(action='DESELECT')
for o in parts:o.select_set(True)
# camera
bpy.ops.object.camera_add(location=(0,0.48,-2.0), rotation=(math.radians(82),0,math.radians(28)))
cam=bpy.context.object; bpy.context.scene.camera=cam
# point camera
def look_at(obj, target): obj.rotation_euler=(Vector(target)-obj.location).to_track_quat('-Z','Y').to_euler()
look_at(cam,(0,0.40,0)); cam.data.lens=58
bpy.ops.object.light_add(type='AREA', location=(-1.2,1.4,-1.3)); key=bpy.context.object; key.data.energy=650; key.data.shape='DISK'; key.data.size=2.2; look_at(key,(0,.4,0))
bpy.ops.object.light_add(type='AREA', location=(1.3,.7,.8)); fill=bpy.context.object; fill.data.energy=350; fill.data.size=1.8; look_at(fill,(0,.4,0))
scene=bpy.context.scene; scene.render.engine='BLENDER_EEVEE'; scene.render.resolution_x=640; scene.render.resolution_y=640; scene.render.resolution_percentage=100
scene.render.image_settings.file_format='PNG'; scene.render.filepath=str(OUT/'BudBud_preview.png'); scene.world.color=(0.92,0.95,0.96)
bpy.ops.render.render(write_still=True)
print('BUDBUD_EXPORT_OK')