import os
import torch
import numpy as np
import tifffile
import folder_paths
from typing import Dict, Tuple
import OpenImageIO as oiio
from datetime import datetime
import sys

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2 as cv

class saver:
    """Image saver node with dynamic widget behavior"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # Common widgets that are always visible at the top
                "images": ("IMAGE",),
                "file_path": ("STRING", {"default": "ComfyUI"}),
                "version": ("INT", {"default": 1, "min": -1, "max": 999}),
                "start_frame": ("INT", {"default": 1001, "min": 0, "max": 99999999}),
                "frame_pad": ("INT", {"default": 4, "min": 1, "max": 8}),
                
                # This widget controls the visibility of others
                "file_type": (["exr", "png", "jpg", "webp", "tiff"], {"default": "png"}),
                
                # These widgets will be shown/hidden based on file_type
                "bit_depth": (["8", "16", "32"], {"default": "16"}),
                "exr_compression": (["none", "zip", "zips", "rle", "pxr24", "b44", "b44a", "dwaa", "dwab"], 
                                  {"default": "zips"}),
                "quality": ("INT", {"default": 95, "min": 1, "max": 100}),
                "save_as_grayscale": ("BOOLEAN", {"default": False})
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "COCO Tools/Savers"
    
    # This tells ComfyUI that file_type widget needs custom handling
    CUSTOM_WIDGETS = ["file_type"]

    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"

    @staticmethod
    def is_grayscale(image: np.ndarray) -> bool:
        """Check if an image is grayscale by either being single channel or having identical RGB channels."""
        if len(image.shape) == 2 or image.shape[-1] == 1:
            return True
        if image.shape[-1] == 3:
            return np.allclose(image[..., 0], image[..., 1]) and np.allclose(image[..., 1], image[..., 2])
        return False

    def _validate_format_bitdepth(self, file_type: str, bit_depth: int) -> Tuple[str, int]:
        valid_combinations = {
            "exr": [16, 32],  # OpenEXR supports half and full float
            "png": [8, 16, 32],  # Now supporting 32-bit PNGs through OpenImageIO
            "jpg": [8],
            "jpeg": [8],
            "webp": [8],
            "tiff": [8, 16, 32]
        }
        
        if bit_depth not in valid_combinations[file_type]:
            sys.stderr.write(f"Warning: {file_type} format only supports {valid_combinations[file_type]} bit depth. Adjusting to {valid_combinations[file_type][0]} bit.\n")
            bit_depth = valid_combinations[file_type][0]
        
        return file_type, bit_depth

    def increment_filename(self, filepath: str) -> str:
        base, ext = os.path.splitext(filepath)
        counter = 1
        new_filepath = f"{base}_{counter:05d}{ext}"
        while os.path.exists(new_filepath):
            counter += 1
            new_filepath = f"{base}_{counter:05d}{ext}"
        return new_filepath

    def _convert_bit_depth(self, img: np.ndarray, bit_depth: int) -> np.ndarray:
        """Convert image to specified bit depth."""
        if bit_depth == 8:
            return (np.clip(img, 0, 1) * 255).astype(np.uint8)
        elif bit_depth == 16:
            if isinstance(img.dtype, np.floating):
                # For floating point data (like EXR), preserve as float16
                return img.astype(np.float16)
            else:
                # For integer data, scale to 16-bit range
                return (np.clip(img, 0, 1) * 65535).astype(np.uint16)
        else:  # 32-bit
            return img.astype(np.float32)

    def _prepare_image_for_saving(self, img: np.ndarray, file_type: str, save_as_grayscale: bool = False) -> np.ndarray:
        """Prepare image for saving by handling color space and channel conversion."""
        # Handle single channel images
        if len(img.shape) == 2:
            img = img[..., np.newaxis]
        
        # Convert to grayscale if requested or if image is already grayscale
        if save_as_grayscale or self.is_grayscale(img):
            if img.shape[-1] == 3:
                # For EXR workflow, just take R channel to preserve original values
                img = img[..., 0:1]
        # Convert to BGR for OpenCV formats only (jpg, jpeg, webp) if not saving as grayscale
        elif file_type in ['jpg', 'jpeg', 'webp'] and img.shape[-1] >= 3:
            img = cv.cvtColor(img, cv.COLOR_RGB2BGR)
        
        return img

    def _prepare_image_data(self, img_tensor: torch.Tensor, file_type: str, bit_depth: int, 
                           save_as_grayscale: bool) -> np.ndarray:
        """Prepare image data for saving by handling conversion and preprocessing."""
        # Convert from torch tensor if needed
        if isinstance(img_tensor, torch.Tensor):
            # Remove batch dimension if present
            if img_tensor.ndim == 4 and img_tensor.shape[0] == 1:
                img_np = img_tensor.squeeze(0).cpu().numpy()
            else:
                img_np = img_tensor.cpu().numpy()
        else:
            img_np = img_tensor
        
        sys.stderr.write(f"Initial data - Shape: {img_np.shape}, Min: {img_np.min()}, Max: {img_np.max()}\n")
        
        # Prepare image (handles channel conversion if needed)
        img_np = self._prepare_image_for_saving(img_np, file_type, save_as_grayscale)
        
        # Convert bit depth for all formats
        img_np = self._convert_bit_depth(img_np, bit_depth)
        sys.stderr.write(f"After bit_depth conversion min: {img_np.min()}, max: {img_np.max()}\n")
        
        return img_np

    def _save_exr(self, img_np: np.ndarray, out_path: str, save_as_grayscale: bool, exr_compression: str, bit_depth: int) -> None:
        """Save image data as EXR file using OpenImageIO."""
        try:
            # Determine channels and prepare data
            if img_np.ndim == 2 or img_np.shape[-1] == 1 or save_as_grayscale:
                channels = 1
                if img_np.ndim == 3 and img_np.shape[-1] > 1:
                    exr_data = img_np[..., 0:1]  # Extract first channel for grayscale
                else:
                    exr_data = img_np[..., np.newaxis] if img_np.ndim == 2 else img_np
            else:
                channels = img_np.shape[-1]
                exr_data = img_np

            sys.stderr.write(f"\nSaving EXR - Shape: {exr_data.shape}, Channels: {channels}\n")
            sys.stderr.write(f"Value range - Min: {exr_data.min()}, Max: {exr_data.max()}\n")
            
            # Ensure data is contiguous and in the right format
            if bit_depth == 16:
                exr_data = np.ascontiguousarray(exr_data.astype(np.float16))
                pixel_type = oiio.HALF
            else:  # 32-bit
                exr_data = np.ascontiguousarray(exr_data.astype(np.float32))
                pixel_type = oiio.FLOAT
            
            # Create spec with detected channels
            spec = oiio.ImageSpec(
                exr_data.shape[1],  # width
                exr_data.shape[0],  # height
                channels,  # channels
                pixel_type  # Use appropriate pixel type based on bit depth
            )
            
            # Set EXR-specific attributes
            spec.attribute("compression", exr_compression)
            spec.attribute("Orientation", 1)
            spec.attribute("DateTime", datetime.now().isoformat())
            spec.attribute("Software", "COCO Tools")
            
            # Create and write image buffer
            buf = oiio.ImageBuf(spec)
            buf.set_pixels(oiio.ROI(), exr_data)
            
            if not buf.write(out_path):
                raise RuntimeError(f"Failed to write EXR: {oiio.geterror()}")
            
        except Exception as e:
            raise RuntimeError(f"OpenImageIO EXR save failed: {str(e)}")

    def _save_png(self, img_np: np.ndarray, out_path: str, bit_depth: int) -> None:
        """Save image data as PNG file using OpenImageIO."""
        try:
            # Determine data type based on bit depth
            if bit_depth == 8:
                dtype = "uint8"
                pixel_type = np.uint8
            elif bit_depth == 16:
                dtype = "uint16"
                pixel_type = np.uint16
            else:  # 32-bit
                dtype = "float"
                pixel_type = np.float32
            
            # Convert and validate array
            png_data = np.ascontiguousarray(img_np.astype(pixel_type))
            
            # Create image spec
            channels = 1 if png_data.ndim == 2 else png_data.shape[-1]
            spec = oiio.ImageSpec(
                png_data.shape[1],  # width
                png_data.shape[0],  # height
                channels,  # channels
                dtype
            )
            
            # Set PNG-specific attributes
            spec.attribute("compression", "zip")
            spec.attribute("png:compressionLevel", 9)
            
            if bit_depth == 32:
                spec.attribute("oiio:ColorSpace", "Linear")
            
            # Create and write image buffer
            buf = oiio.ImageBuf(spec)
            buf.set_pixels(oiio.ROI(), png_data)
            
            if not buf.write(out_path):
                raise RuntimeError(f"Failed to write PNG: {oiio.geterror()}")
        
        except Exception as e:
            raise RuntimeError(f"OpenImageIO PNG save error: {str(e)}")

    def _save_jpeg_or_webp(self, img_np: np.ndarray, out_path: str, quality: int) -> None:
        """Save image data as JPEG or WebP file using OpenCV."""
        cv.imwrite(out_path, img_np, [cv.IMWRITE_JPEG_QUALITY, quality])

    def _save_tiff(self, img_np: np.ndarray, out_path: str, bit_depth: int) -> None:
        """Save image data as TIFF file using tifffile."""
        # Convert BGR back to RGB for TIFF saving
        if img_np.shape[-1] >= 3:
            img_np = cv.cvtColor(img_np, cv.COLOR_BGR2RGB)
        
        # Handle TIFF saving with appropriate bit depth
        if bit_depth == 8:
            tifffile.imwrite(out_path, img_np.astype(np.uint8), photometric='rgb')
        elif bit_depth == 16:
            tifffile.imwrite(out_path, img_np.astype(np.uint16), photometric='rgb')
        else:  # 32-bit
            tifffile.imwrite(
                out_path,
                img_np.astype(np.float32),
                photometric='rgb',
                dtype=np.float32
            )

    def save_images(self, images, file_path, file_type, bit_depth,
                   quality=95, save_as_grayscale=False,
                   version=1, start_frame=1001, frame_pad=4,
                   prompt=None, extra_pnginfo=None, exr_compression="zips"):
        """Save a batch of images in various formats with support for versioning and frame numbering."""
        try:
            bit_depth = int(bit_depth)
            file_type, bit_depth = self._validate_format_bitdepth(file_type.lower(), bit_depth)
            file_ext = f".{file_type}"

            # Handle absolute vs relative paths
            if not os.path.isabs(file_path):
                full_output_folder, filename, _, subfolder, _ = folder_paths.get_save_image_path(
                    file_path, self.output_dir, images.shape[2], images.shape[3]
                )
                base_path = os.path.join(full_output_folder, filename)
            else:
                base_path = file_path
                os.makedirs(os.path.dirname(base_path), exist_ok=True)

            # Handle versioning
            version_str = f"_v{version:03}" if version >= 0 else ""
            
            # Process each image in the batch
            for i, img_tensor in enumerate(images):
                # Prepare image data
                img_np = self._prepare_image_data(img_tensor, file_type, bit_depth, save_as_grayscale)
                
                # Generate output filename
                frame_num = f".{str(start_frame + i).zfill(frame_pad)}" if file_type == "exr" else f"_{i:05d}"
                out_path = f"{base_path}{version_str}{frame_num}{file_ext}"
                
                # Use improved filename increment if file exists
                if os.path.exists(out_path):
                    out_path = self.increment_filename(out_path)

                # Save the image based on format
                if file_type == "exr":
                    self._save_exr(img_np, out_path, save_as_grayscale, exr_compression, bit_depth)
                elif file_type == "png":
                    self._save_png(img_np, out_path, bit_depth)
                elif file_type in ["jpg", "jpeg", "webp"]:
                    self._save_jpeg_or_webp(img_np, out_path, quality)
                elif file_type == "tiff":
                    self._save_tiff(img_np, out_path, bit_depth)

            # Return success
            return {"ui": {"images": []}}

        except Exception as e:
            sys.stderr.write(f"Error saving images: {e}\n")
            return {"ui": {"error": str(e)}}

# Register the node
NODE_CLASS_MAPPINGS = {
    "saver": saver,
}

# Optional: Register display names
NODE_DISPLAY_NAME_MAPPINGS = {
    "saver": "Image Saver"
}