import os
from PIL import Image
import numpy as np
import torch

class LoadWebPAnimation:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": ("STRING", {"default": ""}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "load_animation"

    def load_animation(self, path):
        # Check if the file exists
        if not os.path.exists(path):
            raise FileNotFoundError(f"File not found: {path}")

        # Open the .webp file
        with Image.open(path) as img:
            # Check if it's an animation
            if hasattr(img, 'is_animated') and img.is_animated:
                frames = []
                # Extract each frame
                for i in range(img.n_frames):
                    img.seek(i)
                    frame = img.copy().convert('RGB')  # Ensure 3 channels (RGB)
                    frames.append(np.array(frame))
                batch = np.stack(frames, axis=0)  # Shape: (num_frames, height, width, 3)
            else:
                # Single image case
                frame = img.convert('RGB')
                batch = np.expand_dims(np.array(frame), axis=0)  # Shape: (1, height, width, 3)

        # Convert to float32 and normalize to [0, 1]
        batch = batch.astype(np.float32) / 255.0
        
        # Convert to PyTorch tensor without permuting
        batch = torch.from_numpy(batch)  # Shape: (num_frames, height, width, 3)
        
        return (batch,)  # Return as a tuple

# Register the node with ComfyUI
NODE_CLASS_MAPPINGS = {
    "LoadWebPAnimation": LoadWebPAnimation
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadWebPAnimation": "Load WebP Animation"
}