import json
import os
import random
import re
import time
import uuid
from collections import OrderedDict

import numpy as np
from PIL import Image

import folder_paths


_MEDIA_EXT_RE = re.compile(
    r"\.(png|jpe?g|webp|gif|bmp|tiff?|avif|mp4|mov|webm|mkv|m4v)$", re.IGNORECASE
)
_DATE_TOKEN_RE = re.compile(r"%date:([^%]+)%")
_SERVE_TOKENS = OrderedDict()
_SERVE_CAP = 256


def _clean_input_token_value(value):
    cleaned = ""
    if value is not None:
        cleaned = value if isinstance(value, str) else str(value)
        cleaned = _MEDIA_EXT_RE.sub("", cleaned.strip())
        cleaned = cleaned.replace("\\", "_").replace("/", "_")
    return cleaned


def _expand_date_tokens(value):
    if not isinstance(value, str) or "%date:" not in value:
        return value

    now = time.localtime()

    def pad(number, width):
        return str(number).zfill(width)

    def repl(match):
        fmt = match.group(1)

        def swap(token_match):
            token = token_match.group(0)
            if token == "yyyy":
                return pad(now.tm_year, 4)
            if token == "yy":
                return str(now.tm_year)[-2:]
            if token == "MM":
                return pad(now.tm_mon, 2)
            if token == "M":
                return str(now.tm_mon)
            if token == "dd":
                return pad(now.tm_mday, 2)
            if token == "d":
                return str(now.tm_mday)
            if token in ("hh", "HH"):
                return pad(now.tm_hour, 2)
            if token in ("h", "H"):
                return str(now.tm_hour)
            if token == "mm":
                return pad(now.tm_min, 2)
            if token == "m":
                return str(now.tm_min)
            if token == "ss":
                return pad(now.tm_sec, 2)
            if token == "s":
                return str(now.tm_sec)
            return token

        return re.sub(r"yyyy|yy|MM|M|dd|d|hh|h|HH|H|mm|m|ss|s", swap, fmt)

    return _DATE_TOKEN_RE.sub(repl, value)


def _expand_native_tokens(value):
    if not isinstance(value, str) or "%" not in value:
        return value

    now = time.localtime()
    replacements = (
        ("%year%", f"{now.tm_year:04}"),
        ("%month%", f"{now.tm_mon:02}"),
        ("%day%", f"{now.tm_mday:02}"),
        ("%hour%", f"{now.tm_hour:02}"),
        ("%minute%", f"{now.tm_min:02}"),
        ("%second%", f"{now.tm_sec:02}"),
    )
    for token, replacement in replacements:
        value = value.replace(token, replacement)
    return value


def _safe_pattern(value):
    value = str(value or "").replace("\\", "/")
    value = re.sub(r'[<>:"|?*]', "_", value)
    value = re.sub(r"/{2,}", "/", value).strip(" /.") or "image_%counter%"
    return value


def _register_serve_token(path):
    token = uuid.uuid4().hex
    _SERVE_TOKENS[token] = path
    while len(_SERVE_TOKENS) > _SERVE_CAP:
        _SERVE_TOKENS.popitem(last=False)
    return token


def resolve_serve_token(token):
    return _SERVE_TOKENS.get(str(token or ""))


def _is_within_directory(parent_path, child_path):
    try:
        return os.path.commonpath([parent_path, child_path]) == parent_path
    except ValueError:
        return False


def _next_counter(directory, filename_template):
    os.makedirs(directory, exist_ok=True)
    if "%counter%" not in filename_template:
        return 1

    parts = filename_template.split("%counter%")
    highest = 0
    for entry in os.listdir(directory):
        if not entry.startswith(parts[0]) or not entry.endswith(parts[-1]):
            continue
        middle = entry[len(parts[0]):]
        if parts[-1]:
            middle = middle[: -len(parts[-1])]
        if middle.isdigit():
            highest = max(highest, int(middle))
    return highest + 1


def _build_pnginfo(prompt=None, extra_pnginfo=None):
    try:
        pnginfo = Image.PngImagePlugin.PngInfo()
    except AttributeError:
        from PIL.PngImagePlugin import PngInfo

        pnginfo = PngInfo()
    if prompt is not None:
        pnginfo.add_text("prompt", json.dumps(prompt))
    if isinstance(extra_pnginfo, dict):
        for key, value in extra_pnginfo.items():
            pnginfo.add_text(str(key), json.dumps(value))
    return pnginfo


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,)}


class DumasSaveImageNode:
    DESCRIPTION = (
        "Dumas Save Image writes images to any folder, with filename tokens "
        "such as %input%, %input2%, %date:yyyy-MM-dd%, %counter%, %width%, "
        "%height%, and %batch_num%."
    )
    RETURN_TYPES = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "Dumas/Image"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE", {"tooltip": "Image batch to save."}),
                "folder": (
                    "STRING",
                    {
                        "default": "",
                        "multiline": False,
                        "tooltip": (
                            "Destination folder. Leave empty to save in ComfyUI's "
                            "output directory."
                        ),
                    },
                ),
                "pattern": (
                    "STRING",
                    {
                        "default": "image_%date:yyyy-MM-dd%_%counter%",
                        "multiline": False,
                        "tooltip": (
                            "Filename pattern with optional subfolders. Tokens: "
                            "%input%, %input2%, %date:yyyy-MM-dd%, %counter%, "
                            "%width%, %height%, %batch_num%."
                        ),
                    },
                ),
                "format": (["png", "jpg"], {"default": "png"}),
                "quality": (
                    "INT",
                    {"default": 100, "min": 1, "max": 100, "step": 1},
                ),
                "embed_workflow": ("BOOLEAN", {"default": True}),
                "save_on_run": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "name": (
                    "STRING",
                    {
                        "forceInput": True,
                        "tooltip": "Optional text inserted by the %input% token.",
                    },
                ),
                "name_2": (
                    "STRING",
                    {
                        "forceInput": True,
                        "tooltip": "Optional text inserted by the %input2% token.",
                    },
                ),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    @classmethod
    def IS_CHANGED(cls, **_kwargs):
        return float("nan")

    def save_images(
        self,
        images,
        folder,
        pattern,
        format,
        quality,
        embed_workflow,
        save_on_run,
        name=None,
        name_2=None,
        prompt=None,
        extra_pnginfo=None,
    ):
        if not save_on_run:
            return {"ui": {"images": []}}

        width = int(images.shape[2])
        height = int(images.shape[1])
        output_dir = folder_paths.get_output_directory()
        target_dir = os.path.abspath(folder.strip()) if str(folder or "").strip() else output_dir
        os.makedirs(target_dir, exist_ok=True)

        resolved_pattern = str(pattern or "image_%date:yyyy-MM-dd%_%counter%")
        resolved_pattern = resolved_pattern.replace("%input%", _clean_input_token_value(name))
        resolved_pattern = resolved_pattern.replace("%input2%", _clean_input_token_value(name_2))
        resolved_pattern = _expand_date_tokens(resolved_pattern)
        resolved_pattern = _expand_native_tokens(resolved_pattern)
        resolved_pattern = resolved_pattern.replace("%width%", str(width))
        resolved_pattern = resolved_pattern.replace("%height%", str(height))
        resolved_pattern = _safe_pattern(resolved_pattern)

        extension = ".jpg" if format == "jpg" else ".png"
        quality = max(1, min(100, int(quality)))
        ui_images = []

        for batch_index in range(images.shape[0]):
            frame_pattern = resolved_pattern.replace("%batch_num%", str(batch_index))
            frame_parts = [part for part in frame_pattern.split("/") if part]
            sub_dirs = frame_parts[:-1]
            filename_template = (frame_parts[-1] if frame_parts else "image_%counter%") + extension
            frame_dir = os.path.join(target_dir, *sub_dirs)
            counter = _next_counter(frame_dir, filename_template)
            filename = filename_template.replace("%counter%", str(counter).zfill(3))
            image = _tensor_image_to_pil_image(images[batch_index : batch_index + 1])
            full_path = os.path.join(frame_dir, filename)

            if format == "jpg":
                image = image.convert("RGB")
                image.save(full_path, "JPEG", quality=quality)
            else:
                pnginfo = None
                if embed_workflow:
                    pnginfo = _build_pnginfo(prompt=prompt, extra_pnginfo=extra_pnginfo)
                image.save(full_path, "PNG", pnginfo=pnginfo)

            if _is_within_directory(output_dir, full_path):
                subfolder = os.path.relpath(frame_dir, output_dir)
                ui_images.append(
                    {
                        "filename": filename,
                        "subfolder": "" if subfolder == "." else subfolder.replace("\\", "/"),
                        "type": "output",
                    }
                )
            else:
                ui_images.append(
                    {
                        "filename": filename,
                        "subfolder": frame_dir.replace("\\", "/"),
                        "type": "external",
                        "token": _register_serve_token(full_path),
                    }
                )

        return {"ui": {"images": ui_images}}


NODE_CLASS_MAPPINGS = {
    "DumasImageCompare": DumasImageCompareNode,
    "DumasSaveImage": DumasSaveImageNode,
}

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