import os
import re
import numpy as np
from datetime import datetime
from PIL import Image, ImageCms
import folder_paths

# Standard Apple ImageIO DQT Quantization Tables for JPEG (Luminance & Chrominance)
APPLE_DQT_LUMA = [
    2, 1, 1, 2, 2, 4, 5, 6,
    1, 1, 1, 2, 3, 6, 6, 6,
    1, 1, 2, 2, 4, 6, 7, 6,
    1, 2, 2, 3, 5, 9, 8, 7,
    2, 2, 4, 6, 7, 11, 10, 8,
    2, 3, 5, 6, 8, 10, 11, 9,
    5, 6, 8, 9, 10, 12, 12, 10,
    7, 9, 9, 10, 11, 10, 10, 10
]

APPLE_DQT_CHROMA = [
    2, 2, 2, 4, 8, 8, 8, 8,
    2, 2, 3, 5, 8, 8, 8, 8,
    2, 3, 5, 8, 8, 8, 8, 8,
    4, 5, 8, 8, 8, 8, 8, 8,
    8, 8, 8, 8, 8, 8, 8, 8,
    8, 8, 8, 8, 8, 8, 8, 8,
    8, 8, 8, 8, 8, 8, 8, 8,
    8, 8, 8, 8, 8, 8, 8, 8
]

def process_filename_prefix(prefix):
    def replace_date(match):
        fmt_str = match.group(1)
        replacements = [
            ("yyyy", "%Y"), ("YY", "%y"), ("yy", "%y"),
            ("MM", "%m"), ("dd", "%d"),
            ("HH", "%H"), ("hh", "%H"),
            ("mm", "%M"), ("ss", "%S")
        ]
        for old, new in replacements:
            fmt_str = fmt_str.replace(old, new)
        return datetime.now().strftime(fmt_str)

    prefix = re.sub(r"%date:([^%]+)%", replace_date, prefix)
    return prefix

class SaveImageJPEG:
    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"
        self.prefix_append = ""
        
        # Generate binary sRGB IEC61966-2.1 ICC profile
        srgb_profile = ImageCms.createProfile("sRGB")
        self.srgb_icc_bytes = ImageCms.ImageCmsProfile(srgb_profile).tobytes()

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "Images": ("IMAGE", {"tooltip": "Input images tensor to save."}),
                "Filename Prefix": ("STRING", {"default": "iPhone_JPEG"}),
                "Quality": ("INT", {"default": 88, "min": 1, "max": 100, "step": 1, "tooltip": "JPEG Quality factor (Apple default is 85-90%)."}),
                "Use Apple JPEG Tables": ("BOOLEAN", {"default": True, "tooltip": "Use custom Apple DQT quantization tables."}),
                "Subsampling": (["4:2:0", "4:2:2", "4:4:4"], {"default": "4:2:0", "tooltip": "Chroma subsampling mode (Apple default is 4:2:0)."}),
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "📸 shots by Faded/Utils"
    DESCRIPTION = "Saves JPEG images matching Apple's compression profile with dynamic date formatting."

    def save_images(self, Images, **kwargs):
        filename_prefix = kwargs.get("Filename Prefix", "iPhone_JPEG")
        quality = kwargs.get("Quality", 88)
        use_apple_jpeg_tables = kwargs.get("Use Apple JPEG Tables", True)
        subsampling = kwargs.get("Subsampling", "4:2:0")
        
        filename_prefix += self.prefix_append
        filename_prefix = process_filename_prefix(filename_prefix)
        
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
            filename_prefix, self.output_dir, Images[0].shape[2], Images[0].shape[1]
        )
        
        results = list()
        qtables_param = [APPLE_DQT_LUMA, APPLE_DQT_CHROMA] if use_apple_jpeg_tables else None

        for (batch_number, image) in enumerate(Images):
            # Convert PyTorch tensor (float 0..1) to PIL Image (uint8 0..255 RGB)
            i = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            
            file = f"{filename}_{counter:05}_.jpg"
            file_path = os.path.join(full_output_folder, file)
            
            img.save(
                file_path,
                format="JPEG",
                quality=quality,
                subsampling=subsampling,
                progressive=False,
                qtables=qtables_param,
                icc_profile=self.srgb_icc_bytes,
            )
            
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
            })
            counter += 1

        return {"ui": {"images": results}}

NODE_CLASS_MAPPINGS = {
    "SaveImageJPEG": SaveImageJPEG
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SaveImageJPEG": "Save Image (JPEG)"
}