
# The Original importer was created by Nick
# Copyright (c) 2020 Nick

import bpy
import bmesh
import os
import array
import math
from mathutils import *
from bpy_extras.image_utils import load_image
from .cast import Cast, Model, Animation, Curve, NotificationTrack, Mesh, Skeleton, Bone, Material, File


def utilityBuildPath(root, asset):
    if os.path.isabs(asset):
        return asset

    root = os.path.dirname(root)
    return os.path.join(root, asset)


def utilityAssignBSDFMaterialSlots(material, slots, path):
    
    

    
    material.node_tree.nodes.remove(
        material.node_tree.nodes["Principled BSDF"])
    
    shader = material.node_tree.nodes.new("ShaderNodeGroup")
    output = material.node_tree.nodes['Material Output']
    # グループシェーダーを作成
    shader.node_tree = bpy.data.node_groups['S/G-Blender']

    #テクスチャを定義
    switcher = {
        "albedo": "Diffuse map",
        "diffuse": "Diffuse map",
        "specular": "Specular map",
        "ao": "AO map",
        "cavity": "Cavity map",
        "gloss": "Glossiness map",
        "normal": "Normal map",
        "emissive": "Emission input"
    }
    

    # Loop and connect the slots
    for slot in slots:
        connection = slots[slot]
        if not connection.__class__ is File:
            continue
        if not slot in switcher:
            continue

        texture = material.node_tree.nodes.new("ShaderNodeTexImage") #画像ノードを作成
        try:
            texture.image = bpy.data.images.load(
                utilityBuildPath(path, connection.Path())) #画像を読み込み
        except RuntimeError:
            pass
        if texture.image:
            material.node_tree.links.new(
                shader.inputs[switcher[slot]], texture.outputs["Color"])

            material.node_tree.links.new(shader.outputs[0], output.inputs[0]) # texture.outputのカラーとinputをつなげる(link)
        else:
            material.node_tree.nodes.remove(texture)

def importSkeletonNode(name, skeleton):
    if skeleton is None:
        return None

    armature = bpy.data.armatures.new("Joints")
    armature.display_type = "STICK"

    skeletonObj = bpy.data.objects.new(name, armature)
    skeletonObj.show_in_front = True

    bpy.context.view_layer.active_layer_collection.collection.objects.link(
        skeletonObj)
    bpy.context.view_layer.objects.active = skeletonObj

    bpy.ops.object.mode_set(mode='EDIT')

    bones = skeleton.Bones()
    handles = [None] * len(bones)
    matrices = {}

    for i, bone in enumerate(bones):
        newBone = armature.edit_bones.new(bone.Name())
        newBone.tail = 0, 0.05, 0  # I am sorry but blender sucks

        tempQuat = bone.LocalRotation()  # Also sucks, WXYZ? => XYZW master race
        matRotation = Quaternion(
            (tempQuat[3], tempQuat[0], tempQuat[1], tempQuat[2])).to_matrix().to_4x4()
        matTranslation = Matrix.Translation(Vector(bone.LocalPosition()))

        matrices[bone.Name()] = matTranslation @ matRotation
        handles[i] = newBone

    for i, bone in enumerate(bones):
        if bone.ParentIndex() > -1:
            handles[i].parent = handles[bone.ParentIndex()]

    bpy.context.view_layer.objects.active = skeletonObj
    bpy.ops.object.mode_set(mode='POSE')

    for bone in skeletonObj.pose.bones:
        bone.matrix_basis.identity()
        bone.matrix = matrices[bone.name]

    bpy.ops.pose.armature_apply()
    return skeletonObj


def importMaterialNode(path, material):
    # If you already created the material, ignore this
    materialNew = bpy.data.materials.get(material.Name())
    if materialNew is not None:
        return material.Name(), materialNew

    materialNew = bpy.data.materials.new(name=material.Name())
    materialNew.use_nodes = True

    # Blender really only wants a BSDF shader node
    # so we're gonna give it one
    utilityAssignBSDFMaterialSlots(materialNew, material.Slots(), path)

    return material.Name(), materialNew


def importModelNode(model, path):
    # Extract the name of this model from the path
    modelName = os.path.splitext(os.path.basename(path))[0]

    # Import skeleton for binds, materials for meshes
    skeletonObj = importSkeletonNode(modelName, model.Skeleton())
    materialArray = {key: value for (key, value) in (
        importMaterialNode(path, x) for x in model.Materials())}

    meshes = model.Meshes()
    for mesh in meshes:
        newMesh = bpy.data.meshes.new("polySurfaceMesh")
        blendMesh = bmesh.new()

        vertexColorLayer = blendMesh.loops.layers.color.new("color1")
        vertexWeightLayer = blendMesh.verts.layers.deform.new()
        vertexUVLayers = [blendMesh.loops.layers.uv.new(
            "map%d" % x) for x in range(mesh.UVLayerCount())]

        vertexPositions = mesh.VertexPositionBuffer()
        for x in range(0, len(vertexPositions), 3):
            blendMesh.verts.new(
                Vector((vertexPositions[x], vertexPositions[x + 1], vertexPositions[x + 2])))
        blendMesh.verts.ensure_lookup_table()

        faceLookupMap = [1, 2, 0]
        vertexNormalLayer = []

        vertexNormals = mesh.VertexNormalBuffer()
        vertexColors = mesh.VertexColorBuffer()
        vertexUVs = [mesh.VertexUVLayerBuffer(
            x) for x in range(mesh.UVLayerCount())]

        def vertexToFaceVertex(face):
            for x, loop in enumerate(face.loops):
                vertexIndex = faces[faceStart + faceLookupMap[x]]

                if vertexNormals is not None:
                    vertexNormalLayer.append((vertexNormals[vertexIndex * 3], vertexNormals[(
                        vertexIndex * 3) + 1], vertexNormals[(vertexIndex * 3) + 2]))

                for uvLayer in range(mesh.UVLayerCount()):
                    uv = Vector(
                        (vertexUVs[uvLayer][vertexIndex * 2], vertexUVs[uvLayer][(vertexIndex * 2) + 1]))
                    uv.y = 1.0 - uv.y

                    loop[vertexUVLayers[uvLayer]].uv = uv

                if vertexColors is not None:
                    loop[vertexColorLayer] = [
                        (vertexColors[vertexIndex] >> i & 0xff) / 255.0 for i in (24, 16, 8, 0)]

        faces = mesh.FaceBuffer()
        for faceStart in range(0, len(faces), 3):
            indices = [blendMesh.verts[faces[faceStart + faceLookupMap[0]]],
                       blendMesh.verts[faces[faceStart + faceLookupMap[1]]], blendMesh.verts[faces[faceStart + faceLookupMap[2]]]]

            try:
                newLoop = blendMesh.faces.new(indices)
            except ValueError:
                continue
            else:
                vertexToFaceVertex(newLoop)

        maximumInfluence = mesh.MaximumWeightInfluence()
        if maximumInfluence > 0:
            weightBoneBuffer = mesh.VertexWeightBoneBuffer()
            weightValueBuffer = mesh.VertexWeightValueBuffer()
            for x, vert in enumerate(blendMesh.verts):
                if (weightValueBuffer[x * maximumInfluence] > 0.0):
                    vert[vertexWeightLayer][weightBoneBuffer[x * maximumInfluence]
                                            ] = weightValueBuffer[x * maximumInfluence]

        blendMesh.to_mesh(newMesh)
        newMesh.create_normals_split()

        if len(vertexNormalLayer) > 0:
            for x, _loop in enumerate(newMesh.loops):
                newMesh.loops[x].normal = vertexNormalLayer[x]

        newMesh.validate(clean_customdata=False)
        clnors = array.array('f', [0.0] * (len(newMesh.loops) * 3))

        newMesh.loops.foreach_get("normal", clnors)
        newMesh.polygons.foreach_set(
            "use_smooth", [True] * len(newMesh.polygons))
        newMesh.normals_split_custom_set(tuple(zip(*(iter(clnors),) * 3)))
        newMesh.use_auto_smooth = True

        meshObj = bpy.data.objects.new("CastMesh", newMesh)
        bpy.context.view_layer.active_layer_collection.collection.objects.link(
            meshObj)
        bpy.context.view_layer.objects.active = meshObj

        meshMaterial = mesh.Material()
        if meshMaterial is not None:
            meshObj.data.materials.append(materialArray[meshMaterial.Name()])

        for bone in skeletonObj.pose.bones:
            meshObj.vertex_groups.new(name=bone.name)

        meshObj.parent = skeletonObj
        modifier = meshObj.modifiers.new('Armature Rig', 'ARMATURE')
        modifier.object = skeletonObj
        modifier.use_bone_envelopes = False
        modifier.use_vertex_groups = True


def importRootNode(node, path):
    for child in node.ChildrenOfType(Model):
        importModelNode(child, path)
    # for child in node.ChildrenOfType(Animation):
    #     importAnimationNode(child, path)


def importCast(path):
    cast = Cast()
    cast.load(path)

    for root in cast.Roots():
        importRootNode(root, path)


def load(self, context, filepath=""):
    # シェーダーをアペンド
    shader_path = bpy.context.preferences.addons[__package__].preferences.apex_sgshader_path
    try:

        file_path = shader_path
        inner_path = 'NodeTree'
        object_name = 'S/G-Blender'

        bpy.ops.wm.append(
            filepath=os.path.join(file_path, inner_path, object_name),
            directory=os.path.join(file_path, inner_path),
            filename=object_name
            )
    except:
        
        self.report({'ERROR'}, 'Set the Shader path in AddonPreferences first.')
        return False
    # Parse and load cast nodes
    importCast(filepath)

    # Update the scene, reset view mode before returning.
    bpy.context.view_layer.update()
    bpy.ops.object.mode_set(mode="OBJECT")
    return True

[((16, 7, 16, 27), 'os.path.isabs', 'os.path.isabs', ({(16, 21, 16, 26): 'asset'}, {}), '(asset)', False, 'import os\n'), ((19, 11, 19, 32), 'os.path.dirname', 'os.path.dirname', ({(19, 27, 19, 31): 'root'}, {}), '(root)', False, 'import os\n'), ((20, 11, 20, 36), 'os.path.join', 'os.path.join', ({(20, 24, 20, 28): 'root', (20, 30, 20, 35): 'asset'}, {}), '(root, asset)', False, 'import os\n'), ((75, 15, 75, 47), 'bpy.data.armatures.new', 'bpy.data.armatures.new', ({(75, 38, 75, 46): '"""Joints"""'}, {}), "('Joints')", False, 'import bpy\n'), ((78, 18, 78, 54), 'bpy.data.objects.new', 'bpy.data.objects.new', ({(78, 39, 78, 43): 'name', (78, 45, 78, 53): 'armature'}, {}), '(name, armature)', False, 'import bpy\n'), ((81, 4, 82, 20), 'bpy.context.view_layer.active_layer_collection.collection.objects.link', 'bpy.context.view_layer.active_layer_collection.collection.objects.link', ({(82, 8, 82, 19): 'skeletonObj'}, {}), '(\n    skeletonObj)', False, 'import bpy\n'), ((85, 4, 85, 40), 'bpy.ops.object.mode_set', 'bpy.ops.object.mode_set', (), '', False, 'import bpy\n'), ((108, 4, 108, 40), 'bpy.ops.object.mode_set', 'bpy.ops.object.mode_set', (), '', False, 'import bpy\n'), ((114, 4, 114, 33), 'bpy.ops.pose.armature_apply', 'bpy.ops.pose.armature_apply', ({}, {}), '()', False, 'import bpy\n'), ((279, 4, 279, 35), 'bpy.context.view_layer.update', 'bpy.context.view_layer.update', ({}, {}), '()', False, 'import bpy\n'), ((280, 4, 280, 42), 'bpy.ops.object.mode_set', 'bpy.ops.object.mode_set', (), '', False, 'import bpy\n'), ((145, 18, 145, 56), 'bpy.data.meshes.new', 'bpy.data.meshes.new', ({(145, 38, 145, 55): '"""polySurfaceMesh"""'}, {}), "('polySurfaceMesh')", False, 'import bpy\n'), ((146, 20, 146, 31), 'bmesh.new', 'bmesh.new', ({}, {}), '()', False, 'import bmesh\n'), ((223, 18, 223, 59), 'bpy.data.objects.new', 'bpy.data.objects.new', ({(223, 39, 223, 49): '"""CastMesh"""', (223, 51, 223, 58): 'newMesh'}, {}), "('CastMesh', newMesh)", False, 'import bpy\n'), ((224, 8, 225, 20), 'bpy.context.view_layer.active_layer_collection.collection.objects.link', 'bpy.context.view_layer.active_layer_collection.collection.objects.link', ({(225, 12, 225, 19): 'meshObj'}, {}), '(meshObj)', False, 'import bpy\n'), ((136, 33, 136, 55), 'os.path.basename', 'os.path.basename', ({(136, 50, 136, 54): 'path'}, {}), '(path)', False, 'import os\n'), ((267, 21, 267, 69), 'os.path.join', 'os.path.join', ({(267, 34, 267, 43): 'file_path', (267, 45, 267, 55): 'inner_path', (267, 57, 267, 68): 'object_name'}, {}), '(file_path, inner_path, object_name)', False, 'import os\n'), ((268, 22, 268, 57), 'os.path.join', 'os.path.join', ({(268, 35, 268, 44): 'file_path', (268, 46, 268, 56): 'inner_path'}, {}), '(file_path, inner_path)', False, 'import os\n')]