# coding: utf-8
# Copyright (C) 2023 Ilia Kurochkin <brothermechanic@yandex.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
'''
DESCRIPTION.
Set up materials.
'''
__version__ = '0.1'

import logging
import os
import random
import bpy

logger = logging.getLogger(__name__)


def assign_pbr_material(obj_names, textures_path):
    ''' Assign material with PBR textures pack. '''

    textures = sorted(os.listdir(textures_path))

    for obj_name in obj_names:

        part_name = obj_name.rpartition('_')[0]

        obj = bpy.data.objects[obj_name]
        if len(obj.material_slots) > 0:
            # should be only one material so clear all exist slots
            obj.data.materials.clear()

        if part_name in bpy.data.materials:
            # if material is already in scene
            obj.data.materials.append(bpy.data.materials[part_name])
            continue

        # create Principled BSDF based material
        bmat = bpy.data.materials.new(name=f'M_{part_name}')
        bmat.use_nodes = True
        principled_node = bmat.node_tree.nodes["Principled BSDF"]
        random_color = (
            round(random.uniform(0.1, 0.9), 3),
            round(random.uniform(0.1, 0.9), 3),
            round(random.uniform(0.1, 0.9), 3),
            1.0)
        bmat.diffuse_color = random_color
        principled_node.inputs['Base Color'].default_value = random_color

        for texture in textures:
            if not texture.startswith('T_' + part_name):
                continue
            if '_ao.' in texture:
                image_obj = bpy.data.images.load(os.path.join(textures_path, texture))
                image_obj.colorspace_settings.name = 'Linear'
                texture_ao = bmat.node_tree.nodes.new(type="ShaderNodeTexImage")
                texture_ao.name = texture_ao.label = 'ambient_occlusion'
                texture_ao.location = -1000, 500
                texture_ao.image = bpy.data.images[image_obj.name]

                mix_node = bmat.node_tree.nodes.new(type="ShaderNodeMix")
                mix_node.location = -500, 500
                mix_node.data_type = 'RGBA'
                mix_node.blend_type = 'MULTIPLY'
                mix_node.inputs['Factor'].default_value = 0.5

                bmat.node_tree.links.new(texture_ao.outputs['Color'], mix_node.inputs[7])
                bmat.node_tree.links.new(mix_node.outputs[2], principled_node.inputs['Base Color'])


            if '_d.' in texture:
                image_obj = bpy.data.images.load(os.path.join(textures_path, texture))
                image_obj.colorspace_settings.name = 'sRGB'
                texture_d = bmat.node_tree.nodes.new(type="ShaderNodeTexImage")
                texture_d.name = texture_d.label = 'diffuse'
                texture_d.location = -1000, 0
                texture_d.image = bpy.data.images[image_obj.name]

                if texture_ao:
                    bmat.node_tree.links.new(texture_d.outputs['Color'], mix_node.inputs[6])
                else:
                    bmat.node_tree.links.new(texture_d.outputs['Color'], principled_node.inputs['Base Color'])

            if '_m.' in texture:
                image_obj = bpy.data.images.load(os.path.join(textures_path, texture))
                image_obj.colorspace_settings.name = 'Linear'
                texture_m = bmat.node_tree.nodes.new(type="ShaderNodeTexImage")
                texture_m.name = texture_m.label = 'metallic'
                texture_m.location = -1000, -500
                texture_m.image = bpy.data.images[image_obj.name]

                bmat.node_tree.links.new(texture_m.outputs['Color'], principled_node.inputs['Metallic'])

            if '_r.' in texture:
                image_obj = bpy.data.images.load(os.path.join(textures_path, texture))
                image_obj.colorspace_settings.name = 'Linear'
                texture_r = bmat.node_tree.nodes.new(type="ShaderNodeTexImage")
                texture_r.name = texture_r.label = 'roughness'
                texture_r.location = -1000, -1000
                texture_r.image = bpy.data.images[image_obj.name]

                bmat.node_tree.links.new(texture_r.outputs['Color'], principled_node.inputs['Roughness'])

            if '_n.' in texture:
                image_obj = bpy.data.images.load(os.path.join(textures_path, texture))
                image_obj.colorspace_settings.name = 'Non-Color'
                texture_n = bmat.node_tree.nodes.new(type="ShaderNodeTexImage")
                texture_n.name = texture_n.label = 'normal'
                texture_n.location = -1000, -1500
                texture_n.image = bpy.data.images[image_obj.name]

                normal_node = bmat.node_tree.nodes.new(type="ShaderNodeNormalMap")
                normal_node.location = -500, -1500

                bmat.node_tree.links.new(texture_n.outputs['Color'], normal_node.inputs['Color'])
                bmat.node_tree.links.new(normal_node.outputs['Normal'], principled_node.inputs['Normal'])

        obj.data.materials.append(bmat)

    logger.info('Shading of %s objects is finished!', len(obj_names))
    return True