import os
import glob
from typing import Union, BinaryIO, Dict, List, Tuple, Optional
import torch
import numpy as np
from PIL import Image, ImageOps
import folder_paths
from comfy.utils import ProgressBar


def load_image(image_path):
    """Load image and convert to tensor format expected by ComfyUI"""
    img = Image.open(image_path)
    img = ImageOps.exif_transpose(img)
    img = img.convert("RGB")
    img_array = np.array(img).astype(np.float32) / 255.0
    img_tensor = torch.from_numpy(img_array)[None,]
    return img_tensor


def tensor_to_pil(tensor):
    """Convert tensor back to PIL Image"""
    img_array = tensor.squeeze().cpu().numpy()
    img_array = (img_array * 255).astype(np.uint8)
    return Image.fromarray(img_array)


class BatchImageLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "folder_path": ("STRING", {"default": ""}),
                "file_pattern": ("STRING", {"default": "*.png,*.jpg,*.jpeg"}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING", "INT")
    RETURN_NAMES = ("images", "filenames", "count")
    FUNCTION = "load_batch"
    CATEGORY = "BatchNameLoop"

    def load_batch(self, folder_path: str, file_pattern: str):
        if not os.path.exists(folder_path):
            raise ValueError(f"Folder path does not exist: {folder_path}")
        
        patterns = [p.strip() for p in file_pattern.split(",")]
        image_files = []
        
        for pattern in patterns:
            files = glob.glob(os.path.join(folder_path, pattern))
            image_files.extend(files)
        
        image_files = sorted(list(set(image_files)))
        
        if not image_files:
            raise ValueError(f"No images found in {folder_path} with patterns {file_pattern}")
        
        images = []
        filenames = []
        
        pbar = ProgressBar(len(image_files))
        
        for i, img_path in enumerate(image_files):
            try:
                img_tensor = load_image(img_path)
                images.append(img_tensor)
                filenames.append(os.path.basename(img_path))
                pbar.update(i + 1)
            except Exception as e:
                print(f"Error loading image {img_path}: {e}")
                continue
        
        if not images:
            raise ValueError("No valid images could be loaded")
        
        batch_images = torch.cat(images, dim=0)
        
        return (batch_images, filenames, len(filenames))


class BatchImageIterator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "filenames": ("STRING",),
                "index": ("INT", {"default": 0, "min": 0}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "filename")
    FUNCTION = "iterate"
    CATEGORY = "BatchNameLoop"

    def iterate(self, images, filenames, index):
        if isinstance(filenames, str):
            filenames = [filenames]
        
        if index >= len(images) or index >= len(filenames):
            index = index % len(images)
        
        current_image = images[index:index+1]
        current_filename = filenames[index]
        
        return (current_image, current_filename)


class BatchImageSaver:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "filename_prefix": ("STRING", {"default": "batch_"}),
                "output_folder": ("STRING", {"default": ""}),
            },
            "optional": {
                "original_filename": ("STRING", {"default": ""}),
                "preserve_original_name": ("BOOLEAN", {"default": True}),
            }
        }

    OUTPUT_NODE = True
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("saved_paths",)
    FUNCTION = "save_batch"
    CATEGORY = "BatchNameLoop"

    def save_batch(self, images, filename_prefix, output_folder, original_filename="", preserve_original_name=True):
        if not output_folder:
            output_folder = folder_paths.get_output_directory()
        
        os.makedirs(output_folder, exist_ok=True)
        
        saved_paths = []
        
        for i in range(images.shape[0]):
            img_tensor = images[i:i+1]
            
            if preserve_original_name and original_filename:
                if isinstance(original_filename, list):
                    filename = original_filename[i] if i < len(original_filename) else f"{filename_prefix}{i:04d}.png"
                else:
                    name, ext = os.path.splitext(original_filename)
                    filename = f"{name}_{i:04d}{ext}" if images.shape[0] > 1 else original_filename
            else:
                filename = f"{filename_prefix}{i:04d}.png"
            
            filepath = os.path.join(output_folder, filename)
            
            pil_image = tensor_to_pil(img_tensor)
            pil_image.save(filepath)
            saved_paths.append(filepath)
        
        return (saved_paths,)


class BatchImageSingle:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "original_filename": ("STRING",),
                "output_folder": ("STRING", {"default": ""}),
            },
            "optional": {
                "filename_prefix": ("STRING", {"default": "processed_"}),
                "preserve_original_name": ("BOOLEAN", {"default": True}),
            }
        }

    OUTPUT_NODE = True
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("saved_path",)
    FUNCTION = "save_single"
    CATEGORY = "BatchNameLoop"

    def save_single(self, image, original_filename, output_folder, filename_prefix="processed_", preserve_original_name=True):
        if not output_folder:
            output_folder = folder_paths.get_output_directory()
        
        os.makedirs(output_folder, exist_ok=True)
        
        if preserve_original_name:
            name, ext = os.path.splitext(original_filename)
            filename = f"{filename_prefix}{name}{ext}" if filename_prefix else original_filename
        else:
            filename = f"{filename_prefix}output.png"
        
        filepath = os.path.join(output_folder, filename)
        
        pil_image = tensor_to_pil(image)
        pil_image.save(filepath)
        
        return (filepath,)