from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
import os
import folder_paths
import node_helpers
from PIL import Image, ImageOps, ImageSequence
import numpy as np
import torch
import comfy.model_management

# Import the original LoadImage node for delegation
try:
    from nodes import LoadImage as OriginalLoadImage
except ImportError:
    raise ImportError(
        "Unable to import ComfyUI's LoadImage node. "
        "This custom node requires ComfyUI to be installed and loaded first."
    )

class LoadImageWithMetadata(ComfyNodeABC):
    """Load an image from disk while preserving ComfyUI workflow metadata from PNG.

    This is a drop-in replacement for ComfyUI's built-in Load Image node,
    but it preserves the embedded workflow metadata instead of stripping it.
    """

    DESCRIPTION = """Load image from disk with ComfyUI workflow metadata preserved.

    Works with PNG files that contain ComfyUI metadata.
    """

    CATEGORY = "image/load"

    @classmethod
    def INPUT_TYPES(s) -> InputTypeDict:
        return OriginalLoadImage.INPUT_TYPES()

    RETURN_TYPES = (*OriginalLoadImage.RETURN_TYPES, IO.STRING)
    try:
        RETURN_NAMES = (*OriginalLoadImage.RETURN_NAMES, "workflow_metadata")
    except AttributeError:
        # RETURN_NAMES appears to be optional, fall back to using the return
        # types to be consistent with the original, but still name our new
        # return type since "STRING" isn't particularly meaningful:
        RETURN_NAMES = (*OriginalLoadImage.RETURN_TYPES, "workflow_metadata")

    FUNCTION = "load_image"
    OUTPUT_NODE = True

    def load_image(self, image):
        """Load image and preserve metadata."""

        # Call original implementation
        ret = OriginalLoadImage.load_image(self, image)

        # Load the original image to get metadata
        image_path = folder_paths.get_annotated_filepath(image)
        img = Image.open(image_path)
        metadata = img.info.get("workflow", None)

        return (*ret, metadata)