import os
import logging

import numpy as np
import torch
from PIL import Image

import folder_paths


class LoadDatasetItemAtIndex:
    """Load image, caption, and save_prefix from a dataset folder at a given index.

    Images and captions are paired by filename stem (e.g. abc.png + abc.txt).
    Items are sorted alphabetically by caption filename.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "folder": (sorted(folder_paths.get_input_subfolders()),),
                "index": ("INT", {"default": 0, "min": 0, "max": 999999, "control_after_generate": True}),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING", "STRING")
    RETURN_NAMES = ("image", "caption", "save_prefix")
    FUNCTION = "load"
    CATEGORY = "dataset"

    def load(self, folder, index):
        input_dir = folder_paths.get_input_directory()
        folder_path = os.path.join(input_dir, folder)

        if not os.path.isdir(folder_path):
            raise ValueError(f"Folder not found: {folder_path}")

        # Use caption files as the canonical list (they define the dataset items)
        txt_files = sorted(
            f for f in os.listdir(folder_path) if f.lower().endswith(".txt")
        )

        if not txt_files:
            raise ValueError(f"No .txt caption files found in {folder_path}")

        if index >= len(txt_files):
            raise ValueError(
                f"Index {index} out of range (folder has {len(txt_files)} items)"
            )

        txt_file = txt_files[index]
        stem = os.path.splitext(txt_file)[0]
        txt_path = os.path.join(folder_path, txt_file)

        # Load caption
        with open(txt_path, "r", encoding="utf-8") as f:
            caption = f.read().strip()

        save_prefix = f"{folder}/{stem}"

        # Load matching image
        valid_image_exts = [".png", ".jpg", ".jpeg", ".webp"]
        img_tensor = None
        for ext in valid_image_exts:
            img_path = os.path.join(folder_path, stem + ext)
            if os.path.exists(img_path):
                img = Image.open(img_path).convert("RGB")
                img_array = np.array(img).astype(np.float32) / 255.0
                img_tensor = torch.from_numpy(img_array)[None,]
                break

        if img_tensor is None:
            raise ValueError(
                f"No matching image found for caption {txt_file} in {folder_path}"
            )

        logging.info(
            f"LoadDatasetItemAtIndex: loaded {stem} (index {index}/{len(txt_files)})"
        )
        return (img_tensor, caption, save_prefix)


class SaveDatasetItem:
    """Save an image and/or text caption to the output directory using the exact
    save_prefix as the filename (no counter suffix). save_prefix may contain a
    subdirectory component, e.g. 'my_dataset/image_001'.

    Both inputs are optional — connect whichever outputs you need to save.
    The image is saved as .png with workflow metadata, the caption as .txt.
    """

    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_prefix": ("STRING", {"default": "output/item"}),
            },
            "optional": {
                "image": ("IMAGE",),
                "caption": ("STRING", {"forceInput": True}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "save"
    CATEGORY = "dataset"

    def save(self, save_prefix, image=None, caption=None, prompt=None, extra_pnginfo=None):
        if image is None and caption is None:
            raise ValueError("At least one of image or caption must be connected")

        # Sanitize: prevent directory traversal
        save_prefix = save_prefix.replace("..", "").strip("/")
        dirpath = os.path.join(self.output_dir, os.path.dirname(save_prefix))
        os.makedirs(dirpath, exist_ok=True)

        ui_images = []

        # Save image
        if image is not None:
            filepath = os.path.join(self.output_dir, save_prefix + ".png")

            i = 255.0 * image[0].cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))

            metadata = None
            try:
                from comfy.cli_args import args
                if not args.disable_metadata:
                    from PIL.PngImagePlugin import PngInfo
                    metadata = PngInfo()
                    if prompt is not None:
                        import json
                        metadata.add_text("prompt", json.dumps(prompt))
                    if extra_pnginfo is not None:
                        import json
                        for x in extra_pnginfo:
                            metadata.add_text(x, json.dumps(extra_pnginfo[x]))
            except Exception:
                pass

            img.save(filepath, pnginfo=metadata, compress_level=4)
            logging.info(f"SaveDatasetItem: saved image {filepath}")

            subfolder = os.path.dirname(save_prefix)
            filename = os.path.basename(save_prefix) + ".png"
            ui_images.append({"filename": filename, "subfolder": subfolder, "type": "output"})

        # Save caption
        if caption is not None:
            filepath = os.path.join(self.output_dir, save_prefix + ".txt")
            with open(filepath, "w", encoding="utf-8") as f:
                f.write(caption)
            logging.info(f"SaveDatasetItem: saved caption {filepath}")

        return {"ui": {"images": ui_images}}


NODE_CLASS_MAPPINGS = {
    "LoadDatasetItemAtIndex": LoadDatasetItemAtIndex,
    "SaveDatasetItem": SaveDatasetItem,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadDatasetItemAtIndex": "Load Dataset Item At Index",
    "SaveDatasetItem": "Save Dataset Item",
}