import bpy
import bpy.ops
from bpy.types import Operator

class SMF_OT_flatten_materials(Operator):
    bl_idname = 'scene.flatten_materials'
    bl_label = 'Flatten materials'
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        return True

    def execute(self, context):
        # self.scene = context.scene
        # bpy.context.selected_objects
        # object.material_slots[0].material.node_tree.nodes (list)
        # bpy.ops.object.material_slot_select()
        # bpy.ops.object.mode_set(mode='EDIT') (mode='OBJECT)
        # bpy.types.ShaderNodeTexImage() ???

        selected_objects = [i for i in bpy.context.selected_objects]
        store_selected_objects = selected_objects.copy()
        last_image = None

        #select none
        bpy.ops.object.select_all(action='DESELECT')

        for o in selected_objects:
            imageMaterialNames = {}
            print("gettings materials from %s" % o.name)
            # select this object
            o.select_set(True)
            # switch to edit mode
            bpy.ops.object.mode_set(mode='EDIT')
            bpy.ops.mesh.select_all(action='DESELECT')
            
            for slot in o.material_slots:
                nodes = slot.material.node_tree.nodes

                # Check to not break more complex materials
                if not self.CheckIsSimpleMaterial(nodes): 
                    print("not simple enough")
                    continue

                image = self.getImage(nodes)
                if not last_image: last_image = image

                if last_image == image:
                    o.active_material_index = slot.slot_index
                    bpy.ops.object.material_slot_select()
                    continue

                # materials are not the same: Finish up last image and start again
                if last_image in imageMaterialNames:
                    print("re-assigning materials")
                    print("Hard Mode")
                    # select other material's mesh # active material index bpy.data.objects["name"].active_material_index : INT
                    o.active_material_index = o.material_slots[imageMaterialNames[last_image]].slot_index
                    # assign to this slot # bpy.ops.object.material_slot_assign()
                    bpy.ops.object.material_slot_assign()
                else:
                    print("new texture path: %s" % last_image.filepath)
                    imageMaterialNames[last_image] = o.material_slots[o.active_material_index].name
                    bpy.ops.object.material_slot_assign()
                bpy.ops.mesh.select_all(action='DESELECT')
                last_image = image
                
                # select this material
                o.active_material_index = slot.slot_index
                # select this materials faces
                bpy.ops.object.material_slot_select()
            
            bpy.ops.object.material_slot_assign()
            # switch to object mode
            bpy.ops.object.mode_set(mode='OBJECT')
            # remove unused slots
            bpy.ops.object.material_slot_remove_unused()
        
        # reselect all
        for o in selected_objects:
            o.select_set(True)

        return {'FINISHED'}

    def getImage(self, nodes):
        for node in nodes:
            if node.type == 'TEX_IMAGE':
                return node.image
        else:
            return False

    def CheckIsSimpleMaterial(self, nodes):
        if len(nodes) != 3: return False
        else: return True