import os
from pathlib import Path
import torch
import numpy as np
from PIL import Image, ImageOps

class ImageBatchPath:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {"image_dir": ("STRING", {"default": "", "multiline": True, "placeholder": "Input directory containing images"})},
            "optional": {
                "batch_size": ("INT", {"default": 0, "min": 0, "step": 1, "tooltip": "Number of images to load (0 = all images)"}),
                "start_from": ("INT", {"default": 1, "min": 1, "step": 1, "tooltip": "Start from Nth image (1 = first image)"}),
                "sort_method": (["sequential", "reverse", "random"], {"default": "sequential", "tooltip": "Image loading order: sequential/reverse/random"})
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("IMAGE", "IMAGE_PATH")
    OUTPUT_IS_LIST = (True, True)
    FUNCTION = "load_image_batch"
    CATEGORY = "🧪AILab/📝JoyCaption"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        if 'sort_method' in kwargs and kwargs['sort_method'] == "random":
            return float("NaN")
        return hash(frozenset(kwargs))

    def load_image_batch(self, image_dir: str, batch_size: int = 0, start_from: int = 1, sort_method: str = "sequential"):
        if not os.path.isdir(image_dir):
            raise FileNotFoundError(f"Directory '{image_dir}' cannot be found.")
            
        valid_extensions = ['.jpg', '.jpeg', '.png', '.webp']
        image_files = [f for f in os.listdir(image_dir) if any(f.lower().endswith(ext) for ext in valid_extensions)]
        
        if not image_files:
            raise FileNotFoundError(f"No valid images found in '{image_dir}'.")

        if sort_method == "sequential":
            image_files.sort()
        elif sort_method == "reverse":
            image_files.sort(reverse=True)
        elif sort_method == "random":
            import random
            random.shuffle(image_files)

        start_index = min(start_from - 1, len(image_files) - 1)
        image_files = image_files[start_index:]
        if batch_size > 0:
            image_files = image_files[:batch_size]
        
        images = []
        image_paths = []
        for filename in image_files:
            img_path = os.path.join(image_dir, filename)
            try:
                image = ImageOps.exif_transpose(Image.open(img_path)).convert("RGB")
                image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
                images.append(image)
                image_paths.append(img_path)
            except Exception as e:
                print(f'\033[91mError loading {filename}: {str(e)}\033[0m')
                continue

        return (images, image_paths)

class CaptionSaver:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "string": ("STRING", {"forceInput": True}),
                "image_path": ("STRING", {"forceInput": True}),
            },
            "optional": {
                "image": ("IMAGE",),
                "custom_output_path": ("STRING", {"default": "", "multiline": True, "placeholder": "Custom output directory path. If empty, will use the directory of image_path"}),
                "custom_file_name": ("STRING", {"default": "", "multiline": False, "placeholder": "Custom filename (without extension)"}),
                "overwrite": ("BOOLEAN", {"default": True, "tooltip": "if true, will overwrite the existing txt file, if false, will add a number to the filename to make it unique"})
            }
        }
    
    RETURN_TYPES = ()
    FUNCTION = "save_caption"
    CATEGORY = "🧪AILab/📝JoyCaption"
    OUTPUT_NODE = True

    def get_unique_filename(self, base_path: Path) -> Path:
        """Get a unique filename by adding numbers if file exists."""
        if not base_path.exists():
            return base_path
        
        counter = 1
        while True:
            new_path = base_path.parent / f"{base_path.stem}_{counter:02d}{base_path.suffix}"
            if not new_path.exists():
                return new_path
            counter += 1

    def save_caption(self, string, image_path, image=None, custom_output_path="", custom_file_name="", overwrite=True):
        try:
            image_path = Path(image_path)
            save_dir = Path(custom_output_path.strip()) if custom_output_path.strip() else image_path.parent
            save_dir.mkdir(parents=True, exist_ok=True)
            
            try:
                base_name = custom_file_name.strip() if custom_file_name.strip() else image_path.stem
                txt_path = save_dir / f"{base_name}.txt"
                
                if not overwrite and txt_path.exists():
                    txt_path = self.get_unique_filename(txt_path)
                    base_name = txt_path.stem
                
                with open(txt_path, 'w', encoding='utf-8') as f:
                    f.write(string)
                print(f'\033[93m[{txt_path.name}]:\033[0m {string}')
                
                if image is not None and custom_output_path.strip():
                    try:
                        if isinstance(image, torch.Tensor):
                            while len(image.shape) > 3:
                                image = image.squeeze(0)
                            
                            image = (image.cpu().numpy() * 255).astype(np.uint8)
                            if image.shape[-1] != 3:
                                image = np.transpose(image, (1, 2, 0))
                            image = Image.fromarray(image)
                        
                        img_out_path = save_dir / f"{base_name}{image_path.suffix}"
                        image.save(img_out_path)
                        
                    except Exception as e:
                        print(f'\033[91mFailed to copy image: {str(e)}\033[0m')
                
            except Exception as e:
                print(f'\033[91mError saving caption: {str(e)}\033[0m')
            
        except Exception as e:
            print(f'\033[91mCritical error: {str(e)}\033[0m')
        
        return ()

NODE_CLASS_MAPPINGS = {
    "ImageBatchPath": ImageBatchPath,
    "CaptionSaver": CaptionSaver
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ImageBatchPath": "Image Batch Path 🖼️",
    "CaptionSaver": "Caption Saver 📝"
}