"""
Image Processing Tool
Provides basic image analysis capabilities for GAIA benchmark questions.
"""

from smolagents import Tool
from PIL import Image
import requests
from io import BytesIO

class ImageProcessingTool(Tool):
    name = "image_processing"
    description = "Analyzes images from URLs. Can extract basic information about images including dimensions, format, and basic properties."
    inputs = {
        "image_url": {
            "type": "string",
            "description": "URL of the image to analyze"
        },
        "task": {
            "type": "string",
            "description": "What to do with the image: 'info' (get basic info), 'describe' (describe the image)",
            "default": "info",
            "nullable": True
        }
    }
    output_type = "string"

    def forward(self, image_url: str, task: str = "info") -> str:
        """
        Process an image from a URL.
        
        Args:
            image_url: URL of the image
            task: What task to perform ('info' or 'describe')
            
        Returns:
            Information about the image
        """
        try:
            # Download the image
            response = requests.get(image_url, timeout=10)
            response.raise_for_status()
            
            # Open image
            image = Image.open(BytesIO(response.content))
            
            if task == "info":
                # Return basic image information
                info = {
                    "format": image.format,
                    "mode": image.mode,
                    "size": image.size,
                    "width": image.width,
                    "height": image.height,
                }
                
                return (
                    f"Image Information:\n"
                    f"Format: {info['format']}\n"
                    f"Mode: {info['mode']}\n"
                    f"Dimensions: {info['width']}x{info['height']} pixels\n"
                    f"Size: {info['size']}"
                )
            elif task == "describe":
                # Basic description based on image properties
                return (
                    f"This is a {image.format} image with dimensions {image.width}x{image.height} pixels. "
                    f"Color mode: {image.mode}. "
                    f"For detailed visual description, use a vision model."
                )
            else:
                return f"Unknown task: {task}. Use 'info' or 'describe'."
                
        except requests.exceptions.RequestException as e:
            return f"Error downloading image: {str(e)}"
        except Exception as e:
            return f"Error processing image: {str(e)}"