node.py

import os
import json
import torch
import numpy as np
from PIL import Image, ImageOps
from server import PromptServer
from aiohttp import web

# --- Hilfsfunktion: EXIF Sterne-Bewertung ---
def get_exif_rating(file_path):
    try:
        with Image.open(file_path) as img:
            exif = img.getexif()
            if exif:
                rating = exif.get(18246)
                if rating is not None:
                    if isinstance(rating, tuple): return int(rating[0])
                    return int(rating)
                
                rating_pct = exif.get(18249)
                if rating_pct is not None:
                    p = int(rating_pct)
                    if p >= 99: return 5
                    elif p >= 75: return 4
                    elif p >= 50: return 3
                    elif p >= 25: return 2
                    elif p > 0: return 1
    except Exception:
        pass
    return 0

# --- API Endpunkte ---
@PromptServer.instance.routes.get("/custom_gallery/list")
async def list_images(request):
    folder_path = request.rel_url.query.get("folder", "")
    min_rating_str = request.rel_url.query.get("min_rating", "0")
    page = int(request.rel_url.query.get("page", "1"))
    limit = int(request.rel_url.query.get("limit", "100"))
    
    try: min_rating = int(min_rating_str)
    except: min_rating = 0
        
    if not folder_path or not os.path.exists(folder_path):
        return web.json_response({"error": "Ordner nicht gefunden", "images": [], "folders": [], "total_pages": 1})
        
    valid_extensions = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
    images = []
    folders = []
    
    try:
        with os.scandir(folder_path) as it:
            for entry in it:
                if entry.is_dir():
                    folders.append(entry.name)
                elif entry.is_file():
                    ext = os.path.splitext(entry.name)[1].lower()
                    if ext in valid_extensions:
                        if min_rating > 0:
                            rating = get_exif_rating(entry.path)
                            if rating < min_rating:
                                continue 
                        images.append({"name": entry.name, "time": entry.stat().st_mtime})
                
        images.sort(key=lambda x: x["time"], reverse=True)
        folders.sort(key=str.lower)
        
        total_images = len(images)
        total_pages = max(1, (total_images + limit - 1) // limit)
        start_idx = (page - 1) * limit
        end_idx = start_idx + limit
        paged_images = [img["name"] for img in images[start_idx:end_idx]]
        
        return web.json_response({
            "images": paged_images, 
            "folders": folders,
            "total_pages": total_pages,
            "current_page": page
        })
    except Exception as e:
        return web.json_response({"error": str(e), "images": [], "folders": [], "total_pages": 1})

@PromptServer.instance.routes.get("/custom_gallery/view")
async def view_image(request):
    folder_path = request.rel_url.query.get("folder", "")
    filename = request.rel_url.query.get("file", "")
    full_path = os.path.join(folder_path, filename)
    if not os.path.exists(full_path): return web.Response(status=404)
    return web.FileResponse(full_path)


# --- 1. Folder Gallery Loader Node ---

class FolderGalleryLoader:
    # Interner Speicher für den Auto-Counter
    _auto_index = {}

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "folder_path": ("STRING", {"default": "C:/Pfad/zu/Bildern"}),
                "selection_mode": (["Multiple", "Single"], {"default": "Multiple"}),
                "min_rating": ("INT", {"default": 0, "min": 0, "max": 5, "step": 1}),
                "image_index": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
                # NEU: Der magische Schalter für automatisches Weiterblättern!
                "auto_next": ("BOOLEAN", {"default": False}),
                "preview_size": ("INT", {"default": 100, "min": 30, "max": 300, "step": 10}),
                "show_folders": ("BOOLEAN", {"default": True}),
                "selected_images": ("STRING", {"default": "[]"}), 
            },
            # Wird benötigt, damit jede Node ihren eigenen Counter hat
            "hidden": {"unique_id": "UNIQUE_ID"}
        }

    RETURN_TYPES = ("IMAGE", "MASK", "INT")
    RETURN_NAMES = ("IMAGE", "MASK", "selected_count")
    FUNCTION = "load_image"
    CATEGORY = "image/custom"

    # WICHTIG: Wenn auto_next an ist, erzwingen wir, dass ComfyUI die Node immer neu berechnet
    @classmethod
    def IS_CHANGED(s, auto_next, **kwargs):
        if auto_next:
            return float("nan")
        return ""

    def load_image(self, folder_path, selection_mode, min_rating, image_index, auto_next, preview_size, show_folders, selected_images, unique_id):
        try:
            images_list = json.loads(selected_images)
        except:
            images_list = []

        count = len(images_list)
        
        if count == 0:
            empty_img = torch.zeros((1, 64, 64, 3))
            empty_mask = torch.zeros((64, 64))
            return (empty_img, empty_mask, 0)

        # --- DIE NEUE LOGIK ---
        if auto_next:
            # 1. Aktuellen Index abrufen (Startet bei 0)
            if unique_id not in self._auto_index:
                self._auto_index[unique_id] = 0
                
            idx = self._auto_index[unique_id]
            
            # Sicherheitscheck falls Bilder abgewählt wurden
            if idx >= count:
                idx = 0
                
            # 2. Schon mal für den NÄCHSTEN Durchlauf +1 hochzählen
            next_idx = idx + 1
            if next_idx >= count:
                next_idx = 0 # Zurücksetzen, wenn am Ende
                
            self._auto_index[unique_id] = next_idx
        else:
            # Wenn auto_next aus ist, verhält sich die Node normal mit image_index
            idx = image_index % count

        # --- Bild laden ---
        selected_file = images_list[idx]
        image_path = os.path.join(folder_path, selected_file)
        
        i = Image.open(image_path)
        i = ImageOps.exif_transpose(i)
        image = i.convert("RGB")
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image)[None,]
        
        mask = torch.zeros((64,64), dtype=torch.float32)
        if 'A' in i.getbands():
            mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
            mask = 1. - mask
            mask = torch.from_numpy(mask)
            
        return (image, mask, count)

NODE_CLASS_MAPPINGS = {
    "FolderGalleryLoader": FolderGalleryLoader
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "FolderGalleryLoader": "Folder Gallery Loader"
}