bl_info = {
    "name": "Material Image Remapper",
    "author": "Chipp Walters, Anthony Aragues",
    "version": (1, 0, 1),
    "blender": (2, 82, 0),
    "location": "View3D > Tool Shelf > ♥U",
    "description": "Finds duplicate images in all scene materials and remaps them to one",
    "warning": "",
    "wiki_url": "https://www.chippwalters.com",
    "category": "Edit",
}

import bpy
import os
import hashlib
import tempfile

# Ensure the temp directory exists
temp_dir = tempfile.gettempdir()
os.makedirs(temp_dir, exist_ok=True)

def save_and_hash_image(image, temp_dir, scene):
    """Save the image to a temporary file and calculate its hash."""
    temp_filename = os.path.join(temp_dir, image.name)
    temp_filepath = bpy.path.abspath(temp_filename)
    
    # Attempt to save the image to the temp directory
    try:
        # Ensure the scene parameter is correctly passed as a keyword argument
        image.save_render(filepath=temp_filepath, scene=scene)
        print(f"*** Image saved to: {temp_filepath}")
    except Exception as e:
        print(f"*** Error saving image {image.name}: {e}")
        return None

    # Calculate the hash of the saved image
    try:
        with open(temp_filepath, 'rb') as f:
            img_data = f.read()
        img_hash = hashlib.md5(img_data).hexdigest()
        print(f"*** Hash for {image.name}: {img_hash}")
    except Exception as e:
        print(f"*** Error hashing image {image.name}: {e}")
        return None

    # Optionally, clean up the temporary file
    os.remove(temp_filepath)
    print(f"*** Temporary file removed: {temp_filepath}")
    
    return img_hash


def image_hash(image, temp_dir):
    """Generate a hash for an image, handling both packed and unpacked images."""
    if image.packed_file:
        # Pass the current scene from the context to save_and_hash_image
        return save_and_hash_image(image, temp_dir, bpy.context.scene)
    else:
        # For unpacked images, directly hash the file data
        try:
            abs_path = bpy.path.abspath(image.filepath)
            with open(abs_path, 'rb') as f:
                img_data = f.read()
            img_hash = hashlib.md5(img_data).hexdigest()
            return img_hash
        except Exception as e:
            print(f"*** Error processing image {image.name}: {e}")
            return None

def check_image_nodes(node, images_info, seen_hashes, temp_dir):
    """Recursive function to check nodes for image textures, including in node groups, avoiding duplicates based on hash."""
    if node.type == 'TEX_IMAGE':
        image = node.image
        if image is not None:
            img_hash = image_hash(image, temp_dir)
            if img_hash and img_hash not in seen_hashes:
                seen_hashes.add(img_hash)  # Mark this hash as seen
                image_info = {
                    'name': image.name,
                    'filepath': image.filepath if not image.packed_file else "Packed Image",
                    'exists': os.path.exists(image.filepath) if not image.packed_file else "N/A (packed image)",
                    'packed': image.packed_file is not None,
                    'hash': img_hash
                }
                images_info.append(image_info)
    elif node.type == 'GROUP':
        # Dive into node groups
        for group_node in node.node_tree.nodes:
            check_image_nodes(group_node, images_info, seen_hashes, temp_dir)

def process_materials(temp_dir):
    """Process all materials in the scene."""
    all_images = []
    hash_to_image_info = {}  # Maps hash values to image info for detecting duplicates
    duplicates = {}  # Dictionary to hold duplicate image mappings (old image -> unique image)

    for mat in bpy.data.materials:
        if mat.use_nodes:
            images_info = []
            seen_hashes = set()  # Keep track of image hashes
            for node in mat.node_tree.nodes:
                check_image_nodes(node, images_info, seen_hashes, temp_dir)
            for image_info in images_info:
                img_hash = image_info['hash']
                if img_hash in hash_to_image_info:
                    # If the hash is already present, this is a duplicate
                    old_image_info = image_info
                    unique_image_info = hash_to_image_info[img_hash]
                    duplicates[old_image_info['name']] = unique_image_info['name']
                else:
                    # New unique image, track it
                    hash_to_image_info[img_hash] = image_info
                all_images.append(image_info)  # Track all processed images

    # Print out all unique images found in the scene 
    print("\n*** All images processed in the scene: " + str(len(duplicates)))
    for image in all_images:
        print(f"- {image['name']} (Filepath: {image['filepath']}, Packed: {image['packed']}, Hash: {image['hash']})")

    # Print duplicates
    num_replaced = 0
    if duplicates:
        
        for old_image, unique_image in duplicates.items():
            if old_image != unique_image:
                print(old_image,unique_image)
                print("\n*** Duplicate images found:")
                print(f"- {old_image} -> {unique_image}")
                # Replace all occurrences of old_image with unique_image in materials
                replace_duplicate_image(old_image, unique_image)
                num_replaced += 1
    else:
        print("*** No duplicate images found.")
    #num_replaced = 1
    
    #num_replaced -= 1
    return num_replaced

def replace_duplicate_image(old_image_name, unique_image_name):
    """Replace all occurrences of old_image_name with unique_image_name in materials."""
    print(f"*** Replacing {old_image_name} with {unique_image_name} in materials...")
    for mat in bpy.data.materials:
        if mat.use_nodes:
            for node in mat.node_tree.nodes:
                replace_image_in_node(node, old_image_name, unique_image_name)

def replace_image_in_node(node, old_image_name, unique_image_name):
    """Replace old_image_name with unique_image_name in a single node."""
    if node.type == 'TEX_IMAGE' and node.image.name == old_image_name:
        print(f"*** Found {old_image_name} in node {node.name}")
        node.image = bpy.data.images[unique_image_name]
        print(f"*** Replaced {old_image_name} with {unique_image_name} in node {node.name}")
    elif node.type == 'GROUP':
        # Recursively replace images within node groups
        for group_node in node.node_tree.nodes:
            replace_image_in_node(group_node, old_image_name, unique_image_name)

class SimpleConfirmOperator(bpy.types.Operator):
    """Simple Operator to confirm remapping."""
    bl_idname = "object.simple_confirm_operator"
    bl_label = "Remap Material Images"
    bl_options = {'REGISTER', 'UNDO'}

    def execute(self, context):
        num_replaced = process_materials(temp_dir)
        self.report({'INFO'}, f"Number of images replaced: " + str(num_replaced))
        print("Images replaced " + str(num_replaced))
        return {'FINISHED'}

    def invoke(self, context, event):
        wm = context.window_manager
        return wm.invoke_confirm(self, event)


class CleanUpUnusedMatImagesOperator(bpy.types.Operator):
    bl_idname = "wm.clean_up_unused_mat_images"
    bl_label = "Clean Up Unused Mat Images"
    bl_options = {'REGISTER', 'UNDO'}

    def execute(self, context):
        # Remove unused images
        removed_images = self.remove_unused_images()

        if removed_images:
            print("Removed unused material images:")
            for image_name in removed_images:
                print(f"- {image_name}")
        else:
            print("No unused material images found.")
        
        self.report({'INFO'}, f"Number of images replaced: " + str(len(removed_images)))
        return {'FINISHED'}

    def invoke(self, context, event):
        return context.window_manager.invoke_confirm(self, event)

    def draw(self, context):
        self.layout.label(text="Are you sure you want to clean up unused material images?")
        # Additional UI components can be added here if necessary.

    def remove_unused_images(self):
        removed_image_names = []
        images_to_remove = [image for image in bpy.data.images if not image.users]
        for image in images_to_remove:
            removed_image_names.append(image.name)
            bpy.data.images.remove(image)
        return removed_image_names


class SimplePanel(bpy.types.Panel):
    """Panel to house the Remap Material Images button."""
    bl_label = "Remap Material Images Panel"
    bl_idname = "OBJECT_PT_simple_panel"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = 'Edit'

    def draw(self, context):
        layout = self.layout
        layout.operator("object.simple_confirm_operator", text="Remap Material Images")
        layout.operator("wm.clean_up_unused_mat_images", text="Delete Unused Images")

def register():
    bpy.utils.register_class(SimpleConfirmOperator)
    bpy.utils.register_class(SimplePanel)
    bpy.utils.register_class(CleanUpUnusedMatImagesOperator)

def unregister():
    bpy.utils.unregister_class(SimpleConfirmOperator)
    bpy.utils.unregister_class(SimplePanel)
    bpy.utils.unregister_class(CleanUpUnusedMatImagesOperator)

if __name__ == "__main__":
    register()