import os
import json
import hashlib
import numpy as np
import torch
from PIL import Image, ImageOps
import folder_paths


class TZ_LoadPNGWithMetadata:
    @classmethod
    def INPUT_TYPES(cls):
        input_dir = folder_paths.get_input_directory()
        files = []

        if os.path.isdir(input_dir):
            for f in os.listdir(input_dir):
                if f.lower().endswith(".png"):
                    files.append(f)

        files.sort()

        if not files:
            files = ["example.png"]

        return {
            "required": {
                "image": (files, {"image_upload": True}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING")
    RETURN_NAMES = ("image", "mask", "image_path", "raw_metadata")
    FUNCTION = "load_image"
    CATEGORY = "Herve/metadata"

    @classmethod
    def IS_CHANGED(cls, image):
        image_path = folder_paths.get_annotated_filepath(image)
        m = hashlib.sha256()
        m.update(image.encode("utf-8"))
        if os.path.exists(image_path):
            stat = os.stat(image_path)
            m.update(str(stat.st_mtime_ns).encode("utf-8"))
            m.update(str(stat.st_size).encode("utf-8"))
        return m.hexdigest()

    @classmethod
    def VALIDATE_INPUTS(cls, image):
        if not image.lower().endswith(".png"):
            return "Este nodo solo acepta archivos PNG"
        image_path = folder_paths.get_annotated_filepath(image)
        if not os.path.isfile(image_path):
            return f"No se encontró el archivo: {image}"
        return True

    def load_image(self, image):
        image_path = folder_paths.get_annotated_filepath(image)

        img = Image.open(image_path)
        img = ImageOps.exif_transpose(img)

        info = getattr(img, "info", {}) or {}
        raw_metadata = json.dumps(info, ensure_ascii=False, indent=2)

        image_rgb = img.convert("RGB")
        image_np = np.array(image_rgb).astype(np.float32) / 255.0
        image_tensor = torch.from_numpy(image_np)[None,]

        if "A" in img.getbands():
            mask_np = np.array(img.getchannel("A")).astype(np.float32) / 255.0
            mask = 1.0 - torch.from_numpy(mask_np)
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32)

        return (image_tensor, mask, image_path, raw_metadata)


NODE_CLASS_MAPPINGS = {
    "TZ_LoadPNGWithMetadata": TZ_LoadPNGWithMetadata
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TZ_LoadPNGWithMetadata": "TZ Load PNG With Metadata"
}