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

#AUTHORIZED_TYPES = [
#    "string",
#    "boolean",
#    "integer",
#    "number",
#    "image",
#    "audio",
#    "array",
#    "object",
#    "any",
#    "null",
#]


class ImageLoadTool(Tool):
    name = "_my_tool_image_load"
    description = """
        Load image for the provided task id
        To invoke the tool use code as below
        <code>
        loaded_image = _my_tool_image_load(task_id="dummy")
        </code>
    """

    inputs = {
        "task_id": {
            "type": "string",
            "description": "task id to load image",
        }
    }

    output_type = "image"
    api_url = "https://agents-course-unit4-scoring.hf.space"

    def forward(self, task_id: str) -> Image:
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"
        }
        url = f"{self.api_url}/files/{task_id}"
        response = requests.get(url, headers=headers)
        image = Image.open(BytesIO(response.content)).convert("RGB")
        print(f"***KS*** Loaded image for \n\ttask id: {task_id} \n\timage: {image}")
        return image