"""
Image loading utilities
"""

import os
import numpy as np
from PIL import Image
from typing import Optional

try:
    from folder_paths import get_input_directory
except ImportError:
    # Fallback for testing
    def get_input_directory():
        return "input"


def load_image_from_disk(
    filename: str,
    subfolder: str = ""
) -> Optional[np.ndarray]:
    """
    Load image from ComfyUI input directory.
    
    Args:
        filename: Image filename
        subfolder: Optional subfolder within input directory
        
    Returns:
        Image array (H, W, 3) in range [0, 1] as float32, or None if failed
    """
    if not filename:
        return None
    
    try:
        input_dir = get_input_directory()
        if subfolder:
            image_path = os.path.join(input_dir, subfolder, filename)
        else:
            image_path = os.path.join(input_dir, filename)
        
        if not os.path.exists(image_path):
            print(f"[RotoMaskNode] Image not found: {image_path}")
            return None
        
        # Load image using PIL
        img = Image.open(image_path)
        # Convert to RGB if needed
        if img.mode != "RGB":
            img = img.convert("RGB")
        
        # Convert to numpy array and normalize to [0, 1]
        img_array = np.array(img, dtype=np.float32) / 255.0
        
        return img_array
    except Exception as e:
        print(f"[RotoMaskNode] Error loading image {filename}: {e}")
        return None