import os
import random

import numpy as np
from PIL import Image

import folder_paths


def _tensor_image_to_pil_image(tensor):
    image_tensor = tensor[0]
    if hasattr(image_tensor, "mul") and hasattr(image_tensor, "clamp"):
        image_array = image_tensor.mul(255).clamp(0, 255)
        if hasattr(image_array, "byte"):
            image_array = image_array.byte()
        image_array = image_array.cpu().numpy()
        return Image.fromarray(image_array)

    image_array = 255.0 * image_tensor.cpu().numpy()
    return Image.fromarray(np.clip(image_array, 0, 255).astype(np.uint8))


class DumasImageCompareNode:
    DESCRIPTION = (
        "Dumas Image Compare shows the difference between two images directly on "
        "the node. Connect one or two IMAGE inputs to compare before/after "
        "results, model variants, or processing stages without breaking a "
        "workflow when one branch is bypassed."
    )
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("new image",)
    FUNCTION = "compare_images"
    OUTPUT_NODE = True
    CATEGORY = "Dumas/Image"

    def __init__(self):
        self.output_dir = folder_paths.get_temp_directory()
        self.type = "temp"
        self.prefix_append = "_dumascmp_" + "".join(
            random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5)
        )
        self.compress_level = 4

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "image1": (
                    "IMAGE",
                    {
                        "tooltip": (
                            "First image to compare. Optional so muted or bypassed "
                            "branches do not trigger a missing-input error."
                        )
                    },
                ),
                "image2": (
                    "IMAGE",
                    {
                        "tooltip": (
                            "Second image to compare. Optional so the node can still "
                            "display a single available image."
                        )
                    },
                ),
            }
        }

    def compare_images(self, image1=None, image2=None):
        present = []
        if image1 is not None:
            present.append((1, image1))
        if image2 is not None:
            present.append((2, image2))
        results = []

        if present:
            first_tensor = present[0][1]
            prefix = "dumas_compare" + self.prefix_append
            first_image = first_tensor[0]
            full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
                prefix,
                self.output_dir,
                first_image.shape[1],
                first_image.shape[0],
            )
            join_path = os.path.join
            for slot, tensor in present:
                image = _tensor_image_to_pil_image(tensor)
                file_name = f"{filename}_{counter:05}_.png"
                image.save(
                    join_path(full_output_folder, file_name),
                    compress_level=self.compress_level,
                )
                results.append(
                    {
                        "filename": file_name,
                        "subfolder": subfolder,
                        "type": self.type,
                        "slot": slot,
                    }
                )
                counter += 1

        new_image = image2 if image2 is not None else image1
        return {"ui": {"images": results}, "result": (new_image,)}


NODE_CLASS_MAPPINGS = {
    "DumasImageCompare": DumasImageCompareNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DumasImageCompare": "Dumas Image Compare",
}