from pyedm_platform_selector import pyedm
from logger import log
from tree_node import TreeNode
import bpy
import animation as anim
from math_tools import IDENTITY_MATRIX

def build_bone_id(armature_name, bone_name):
    return f'{armature_name} : {bone_name}'

class BoneNode(TreeNode):
    def __init__(self, pbone: bpy.types.PoseBone, armature: bpy.types.Armature) -> None:
        super().__init__()
        self.name = build_bone_id(armature.name, pbone.name)
        self.armature = armature
        self.pbone = pbone
        self.bone = pbone.bone # rest bone
        self.edm_node = None
        self.update_matrices()

    def build_bones(self, parent: pyedm.Node):
        self.edm_node = parent = anim.extract_bone_animation(parent, self)
        for i in self.children:
            i.build_bones(parent)

    def update_matrices(self):
        self.mat_inv = (self.bone.matrix_local).inverted()
        if self.pbone.parent:
            self.mat = self.pbone.parent.matrix.inverted() @ self.pbone.matrix
        else:
            self.mat = self.pbone.matrix

# we need to build tree because blender holds bones in a flat list
def build_bones_tree(parent, bones_list, armature):
    bones = [BoneNode(b, armature) for b in bones_list]

    for b in bones:
        for c in b.pbone.children:
            for bb in bones:
                if c == bb.pbone:
                    b.add_child(bb)
                    break

    roots = [b for b in bones if b.parent == None]

    for r in roots:
        r.build_bones(parent)

    return bones

def export_armature(armature, parent, bones_dict):
    # Pose bones contain pose data for the current frame, while regular bone objects store the default state of the bone.
    pose_bones = armature.pose.bones
    if not pose_bones:
        log.warning(f'Armature {armature.name} is empty')
        return
    

    #root = pyedm.Node(f'Armature {armature.name} Root')
    #root = pyedm.Transform(f'Armature {armature.name} Root', armature.matrix_local)
    #parent = parent.addChild(root)
    
    bone_nodes = build_bones_tree(parent, pose_bones, armature)
    
    bones_dict.update({(x.name, x) for x in bone_nodes})
    
    return parent
