"""Load every image in a local folder and emit them one by one."""

from __future__ import annotations

import os
import re
from pathlib import Path

import numpy as np
import torch
from PIL import Image, ImageOps, UnidentifiedImageError
from comfy_api.latest import io


def _resolve_folder(folder_path: str) -> Path:
    raw_path = str(folder_path or "").strip().strip('"').strip("'")
    if not raw_path:
        raise ValueError("加载图像（文件夹）：请输入文件夹路径")

    expanded = os.path.expandvars(os.path.expanduser(raw_path))
    folder = Path(expanded)
    if not folder.is_absolute():
        folder = Path.cwd() / folder
    folder = folder.resolve()

    if not folder.exists():
        raise ValueError(f"加载图像（文件夹）：文件夹不存在：{folder}")
    if not folder.is_dir():
        raise ValueError(f"加载图像（文件夹）：路径不是文件夹：{folder}")
    return folder


def _natural_sort_key(path: Path):
    """Sort image2 before image10 while remaining case-insensitive."""
    return tuple(
        int(part) if part.isdigit() else part.casefold()
        for part in re.split(r"(\d+)", path.name)
    )


def _list_image_files(folder: Path) -> list[Path]:
    # Pillow's registry reflects the formats supported by the current runtime,
    # including optional formats supplied by installed Pillow plugins.
    Image.init()
    supported_extensions = {suffix.casefold() for suffix in Image.registered_extensions()}
    image_files = sorted(
        (
            path
            for path in folder.iterdir()
            if path.is_file() and path.suffix.casefold() in supported_extensions
        ),
        key=_natural_sort_key,
    )
    if not image_files:
        raise ValueError(f"加载图像（文件夹）：文件夹中没有可读取的图片：{folder}")
    return image_files


def _load_image_tensor(path: Path) -> torch.Tensor:
    try:
        with Image.open(path) as opened:
            image = ImageOps.exif_transpose(opened)
            image.seek(0)
            if image.mode == "I":
                image = image.point(lambda value: value * (1 / 255))
            rgb_image = image.convert("RGB")
            array = np.asarray(rgb_image, dtype=np.float32) / 255.0
    except (OSError, ValueError, UnidentifiedImageError) as exc:
        raise ValueError(f"加载图像（文件夹）：无法读取图片 {path.name}：{exc}") from exc

    # ComfyUI IMAGE tensors use [batch, height, width, channels]. Each list
    # item is kept as a separate batch of one so original dimensions survive.
    return torch.from_numpy(array).unsqueeze(0)


class LoadImagesFromFolder(io.ComfyNode):
    """Load local images in natural filename order as a ComfyUI output list."""

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="O1keyLoadImagesFromFolder",
            display_name="加载图像（文件夹）",
            category="image",
            description=(
                "读取本地文件夹第一层中的所有图片，按文件名自然顺序逐张输出。"
                "每张图片保留原始分辨率，可直接连接普通图像处理节点。"
            ),
            search_aliases=[
                "文件夹图片",
                "批量加载图片",
                "folder images",
                "load images from folder",
            ],
            inputs=[
                io.String.Input(
                    "文件夹路径",
                    default="",
                    placeholder=r"例如：D:\images",
                ),
            ],
            outputs=[
                io.Image.Output(display_name="图像", is_output_list=True),
            ],
        )

    @classmethod
    def fingerprint_inputs(cls, 文件夹路径: str):
        """Invalidate ComfyUI's cache when the folder's image set changes."""
        try:
            folder = _resolve_folder(文件夹路径)
            return tuple(
                (path.name, path.stat().st_size, path.stat().st_mtime_ns)
                for path in _list_image_files(folder)
            )
        except (OSError, ValueError):
            # Execution will provide the user-facing validation error.
            return str(文件夹路径 or "")

    @classmethod
    def execute(cls, 文件夹路径: str) -> io.NodeOutput:
        folder = _resolve_folder(文件夹路径)
        image_files = _list_image_files(folder)
        images = []

        for index, path in enumerate(image_files, start=1):
            tensor = _load_image_tensor(path)
            images.append(tensor)
            height, width = tensor.shape[1:3]
            print(
                f"加载图像（文件夹）：{index}/{len(image_files)} "
                f"{path.name} ({width}×{height})"
            )

        print(f"加载图像（文件夹）：已从 {folder} 加载 {len(images)} 张图片")
        return io.NodeOutput(images)