import bpy

class MaterialUtil:

    ''' --------- Basic --------- '''
    @staticmethod
    def change_object_material(obj, mat):
        if obj.data.materials:
            obj.data.materials[0] = mat
        else:
            obj.data.materials.append(mat)
    
    ''' ------- Materials ------- '''
    @staticmethod
    def create_normal_material():
        normal_mat = bpy.data.materials.new(name="NormalMaterial")
        normal_mat.use_nodes = True

        nodes = normal_mat.node_tree.nodes
        links = normal_mat.node_tree.links

        nodes.clear()

        geometry_node = nodes.new(type="ShaderNodeNewGeometry")
        vector_transform_node = nodes.new(type="ShaderNodeVectorTransform")
        separate_xyz_node = nodes.new(type="ShaderNodeSeparateXYZ")
        multiply_node_x = nodes.new(type="ShaderNodeMath")
        multiply_node_y = nodes.new(type="ShaderNodeMath")
        multiply_node_z = nodes.new(type="ShaderNodeMath")
        combine_xyz_node = nodes.new(type="ShaderNodeCombineXYZ")
        light_path_node = nodes.new(type="ShaderNodeLightPath")
        emission_node_1 = nodes.new(type="ShaderNodeEmission")
        emission_node_2 = nodes.new(type="ShaderNodeEmission")
        mix_shader_node_1 = nodes.new(type="ShaderNodeMixShader")
        mix_shader_node_2 = nodes.new(type="ShaderNodeMixShader")
        material_output_node = nodes.new(type="ShaderNodeOutputMaterial")

        vector_transform_node.vector_type = 'VECTOR'
        vector_transform_node.convert_from = 'WORLD'
        vector_transform_node.convert_to = 'CAMERA'

        multiply_node_x.operation = 'MULTIPLY'
        multiply_node_x.inputs[1].default_value = 1.0

        multiply_node_y.operation = 'MULTIPLY'
        multiply_node_y.inputs[1].default_value = 1.0

        multiply_node_z.operation = 'MULTIPLY'
        multiply_node_z.inputs[1].default_value = -1.0

        emission_node_1.inputs['Strength'].default_value = 1.0
        emission_node_2.inputs['Strength'].default_value = 1.0

        mix_shader_node_2.inputs['Fac'].default_value = 0.5

        links.new(geometry_node.outputs['Normal'], vector_transform_node.inputs['Vector'])
        links.new(vector_transform_node.outputs['Vector'], separate_xyz_node.inputs['Vector'])
        links.new(separate_xyz_node.outputs['X'], multiply_node_x.inputs[0])
        links.new(separate_xyz_node.outputs['Y'], multiply_node_y.inputs[0])
        links.new(separate_xyz_node.outputs['Z'], multiply_node_z.inputs[0])
        links.new(multiply_node_x.outputs['Value'], combine_xyz_node.inputs['X'])
        links.new(multiply_node_y.outputs['Value'], combine_xyz_node.inputs['Y'])
        links.new(multiply_node_z.outputs['Value'], combine_xyz_node.inputs['Z'])
        links.new(combine_xyz_node.outputs['Vector'], emission_node_1.inputs['Color'])
        links.new(light_path_node.outputs['Is Camera Ray'], mix_shader_node_1.inputs['Fac'])
        links.new(emission_node_1.outputs['Emission'], mix_shader_node_1.inputs[2])
        links.new(mix_shader_node_1.outputs['Shader'], mix_shader_node_2.inputs[1])
        links.new(emission_node_2.outputs['Emission'], mix_shader_node_2.inputs[2])
        links.new(mix_shader_node_2.outputs['Shader'], material_output_node.inputs['Surface'])
        return normal_mat
    
    @staticmethod
    def create_mask_material(color=(1.0, 1.0, 1.0)): 
        mask_mat = bpy.data.materials.new(name="MaskMaterial")
        mask_mat.use_nodes = True

        nodes = mask_mat.node_tree.nodes
        links = mask_mat.node_tree.links

        nodes.clear()
        emission_node = nodes.new(type="ShaderNodeEmission")
        emission_node.inputs['Color'].default_value = (*color, 1.0) 
        emission_node.inputs['Strength'].default_value = 1.0
        material_output_node = nodes.new(type="ShaderNodeOutputMaterial")
        links.new(emission_node.outputs['Emission'], material_output_node.inputs['Surface'])

        return mask_mat

    

# -------- debug --------
if __name__ == "__main__":
    cube = bpy.data.objects.get("Cube")
    normal_mat = MaterialUtil.create_normal_material()
    MaterialUtil.change_object_material(cube, normal_mat)