"""
Improved LoadImageFromURL node for ComfyUI
Place this file in ComfyUI/custom_nodes/load_image_from_url.py
"""
import torch
import requests
from PIL import Image
import numpy as np
import io

class LoadImageFromURL:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "url": ("STRING", {"default": "https://example.com/image.jpg"}),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load_image"
    CATEGORY = "image"
    
    def load_image(self, url):
        try:
            # Download the image
            response = requests.get(url, timeout=30, headers={
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            })
            response.raise_for_status()
            
            # Open the image
            image = Image.open(io.BytesIO(response.content))
            
            # Convert to RGB if needed
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Convert to tensor
            image_array = np.array(image).astype(np.float32) / 255.0
            image_tensor = torch.from_numpy(image_array)[None,]
            
            # Create mask with same dimensions as image (height, width)
            h, w = image_array.shape[:2]
            mask = torch.zeros((h, w), dtype=torch.float32, device="cpu")
            
            return (image_tensor, mask)
            
        except Exception as e:
            print(f"Error loading image from URL: {e}")
            # Return a blank 512x512 image if loading fails
            blank_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32, device="cpu")
            mask = torch.zeros((512, 512), dtype=torch.float32, device="cpu")
            return (blank_image, mask)

# Node class mappings
NODE_CLASS_MAPPINGS = {
    "LoadImageFromURL": LoadImageFromURL,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadImageFromURL": "Load Image From URL",
}