#### 3D Coat PBR to Blender Principled BSDF V 0.1.2

#### 3D Coat PBR to Blender Principled BSDF V 0.1.2
#### Plugin to import 3D Coat's generated PBR material files to Blender's Principal BSDF shader
####
#### This plugin will handle the generation of a new material for each selected object in the viewport.
####
#### 3D Coat Export Settings:
#### - Roughness/Metalness
#### - Uncheck "Use EXPORTER CONSTRUCTOR for per /channel packing"
#### - Set each extension to TGA
#### - Check the following ONLY
####    - Export Color
####    - Export Roughness
####    - Export Metalness
####    - Export Tangent Normal Map
####    - ExportAO
####    - ExportCurvature
#### - Set "Create Padding" to 8, should be enough.
####
#### Currently supported:
####    - Albedo/Color
####    - Metallic
####    - Roughness
####    - Normal
####
#### Usage:
#### 1. Select object(s) to apply new material to
#### 2. Select File => Import => 3D Coat PBR material
#### 3. Select the .OBJ or .FBX model. The plugin will look for each texture/map file in the same directory
#### 4. Profit!
####
#### Thanks to:
####    digman, for infos on how to set up the Principled BSDF shader correctly
####
#### Created by Haunted (Marcel Meisner)
#### Official 3D Coat Forum thread: http://3dcoat.com/forum/index.php?/topic/22025-3d-coat-pbr-to-blender-principled-bsdf-add-on/
####
#### Changes:
####
#### V 0.1.2
####    - Script now finds .tga, .png or .bmp files
####
#### V 0.1.1
####    - Added labels for the texture nodes
####    - Code clean up

bl_info = {
    'name': "3D Coat PBR material",
    'category': "Import",
    'author': 'Marcel Meisner (Haunted)',
    'version': (0, 1, 2),
    'blender': (2, 79, 0),
    'location': 'File > Import > 3D Coat PBR material',
    'description': "Import 3D Coat's generated PBR material files to Blender's Principal BSDF shader",
    'warning': '',
    'support': 'COMMUNITY',
}

import bpy
import os
from collections import namedtuple
from enum import Enum
import glob

# ImportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ImportHelper
from bpy.props import StringProperty, BoolProperty, EnumProperty
from bpy.types import Operator

class ImportFilesCollection(bpy.types.PropertyGroup):
    name = StringProperty(
            name="File Path",
            description="Filepath used for importing the file",
            maxlen=1024,
            subtype='FILE_PATH',
            )
bpy.utils.register_class(ImportFilesCollection)

class EMapType(Enum):
    diffuse     = 1
    metallic    = 2
    roughness   = 3
    normal      = 4

AvailableMaps = namedtuple('AvailableMaps', 'diffuse metallic roughness normal')


def CreatePBRMaterial(obj, availableMaps):
    matName = obj.name + '_3d_coat_pbr'

    print("Creating material '" + matName + "'.")

    mat = bpy.data.materials.new(matName)
    mat.use_nodes = True
    obj.data.materials.append(mat)

    material_output = mat.node_tree.nodes.get("Material Output")
    pbrNode = mat.node_tree.nodes.new('ShaderNodeBsdfPrincipled')
    mat.node_tree.links.new(pbrNode.outputs[0], material_output.inputs[0])
    mat.node_tree.nodes.remove(mat.node_tree.nodes.get('Diffuse BSDF'))

    material_output.location = 200,0

    if(availableMaps.diffuse != None):
        AddTexture(mat, pbrNode, EMapType.diffuse, availableMaps.diffuse)
    if(availableMaps.metallic != None):
        AddTexture(mat, pbrNode, EMapType.metallic, availableMaps.metallic)
    if(availableMaps.roughness != None):
        AddTexture(mat, pbrNode, EMapType.roughness, availableMaps.roughness)
    if(availableMaps.normal != None):
        AddTexture(mat, pbrNode, EMapType.normal, availableMaps.normal)


def AddTexture(material, pbrNode, type, mapPath):
    material_output = material.node_tree.nodes.get('Material Output')
    tree = material.node_tree
    if type == EMapType.diffuse:
        print('Adding diffuse map.')
        node = CreateNewTextureNode(material, -300, 200, mapPath)
        node.label = 'Albedo/Diffuse'
        tree.links.new(node.outputs[0], pbrNode.inputs[0])
    elif type == EMapType.metallic:
        print('Adding metallic map.')
        node = CreateNewTextureNode(material, -500, 0, mapPath)
        node.label = 'Metallic'
        tree.links.new(node.outputs[0], pbrNode.inputs[5])
    elif type == EMapType.roughness:
        print('Adding roughness map.')
        node = CreateNewTextureNode(material, -300, -200, mapPath)
        node.label = 'Roughness'
        tree.links.new(node.outputs[0], pbrNode.inputs[7])
    elif type == EMapType.normal:
        print('Adding normal map.')
        nm_node = CreateNewNormalMapNode(material, -300, -500, mapPath)
        node = CreateNewTextureNode(material, -500, -500, mapPath)
        node.label = 'Normal'
        tree.links.new(node.outputs[0], nm_node.inputs[1])
        tree.links.new(nm_node.outputs[0], pbrNode.inputs["Normal"])


def CreateNewTextureNode(material, x, y, imgPath):
        node = material.node_tree.nodes.new('ShaderNodeTexImage')
        node.image = bpy.data.images.load(imgPath)
        node.location= x, y
        return node


def CreateNewNormalMapNode(material, x, y, imgPath):
        node = material.node_tree.nodes.new('ShaderNodeNormalMap')
        node.location= x, y
        return node


class Import3DCoatPBR(Operator, ImportHelper):
    """Generate a Principled BSDF material for each selected object from 3D Coat's exported PBR files."""
    bl_idname = "node.pbr_3dcoat"  # important since its how bpy.ops.import_test.some_data is constructed
    bl_label = "Import 3DCoat PBR maps"

    # ImportHelper mixin class uses this
    filename_ext    = ".obj"

    filter_glob = StringProperty(
            default="*.obj;*.fbx",
            options={'HIDDEN'},
            )

    files = bpy.props.CollectionProperty(type=ImportFilesCollection)

    def execute(self, context):
        if bpy.context.scene.render.engine != 'CYCLES':
            self.report({'ERROR'},'Add-on works with Cycles Renderer only (for now).')
            return {'CANCELLED'}

        for i, f in enumerate(self.files, 1):
            suffix_color    = 'color'
            suffix_metal    = 'metalness'
            suffix_rough    = 'rough'
            suffix_nmap     = 'nmap'

            # Setup base filename and base path
            filename = os.path.splitext(f.name)[0]
            fdir = os.path.dirname(self.properties.filepath)
            filepath = os.path.join(fdir,filename)

            #Check for selected objects
            layerIndex = 0  #Chosen layer
            objects = bpy.context.scene.objects

            selectedObjects = [ob for ob in objects if ob.layers[layerIndex]
            and ob.select
            and ob.type == 'MESH']

            extension = '.tga'
            pathCheckName = filepath + '_*_' + suffix_color

            if len(glob.glob(pathCheckName + '.tga')) > 0:
                extension = '.tga'
            elif len(glob.glob(pathCheckName + '.png')) > 0:
                extension = '.png'
            elif len(glob.glob(pathCheckName + '.bmp')) > 0:
                extension = '.bmp'

            if len(selectedObjects) > 0:
                # Flags for textures
                hasDiffuse      = False
                hasMetallic     = False
                hasRoughness    = False
                hasNormalmap    = False

                # Search for these files on disc
                diffuse     = glob.glob(filepath + '_*_' + suffix_color + extension)
                metallic    = glob.glob(filepath + '_*_' + suffix_metal + extension)
                roughness   = glob.glob(filepath + '_*_' + suffix_rough + extension)
                normalmap   = glob.glob(filepath + '_*_' + suffix_nmap + extension)

                diffuse = diffuse[0] if len(diffuse) > 0 else ''
                metallic = metallic[0] if len(metallic) > 0 else ''
                roughness = roughness[0] if len(roughness) > 0 else ''
                normalmap = normalmap[0] if len(normalmap) > 0 else ''

                # Check if files are existing
                print("Searching PBR maps...")
                availableMaps = AvailableMaps(
                    diffuse if os.path.isfile(diffuse) else None,
                    metallic if os.path.isfile(metallic) else None,
                    roughness if os.path.isfile(roughness) else None,
                    normalmap if os.path.isfile(normalmap) else None,
                    )

                #Output
                print("Diffuse map: \t" + ("found=" + availableMaps.diffuse  if availableMaps.diffuse != None else "not found!"))
                print("Metallic map: \t" + ("found=" + availableMaps.metallic  if availableMaps.metallic != None  else "not found!"))
                print("Roughness map: \t" + ("found=" + availableMaps.roughness if availableMaps.roughness != None else "not found!"))
                print("Normal map: \t" + ("found=" + availableMaps.normal if availableMaps.normal != None  else "not found!"))

                #Create material
                for obj in selectedObjects:
                    CreatePBRMaterial(obj, availableMaps)
            else:
                self.report({'ERROR'},'No object selected to import the PBR material to.')
                return {'CANCELLED'}

        return {'FINISHED'}


# Only needed if you want to add into a dynamic menu
def menu_func_import(self, context):
    self.layout.operator(Import3DCoatPBR.bl_idname, text="3D Coat PBR material (.tga)")


def register():
    bpy.utils.register_class(Import3DCoatPBR)
    bpy.types.INFO_MT_file_import.append(menu_func_import)


def unregister():
    bpy.utils.unregister_class(Import3DCoatPBR)
    bpy.types.INFO_MT_file_import.remove(menu_func_import)


if __name__ == "__main__":
    register()

    #bpy.ops.node.pbr_3dcoat('INVOKE_DEFAULT')