import modal
from smolagents import AgentImage, Tool

from modal_apps.app import app
from modal_apps.upscaler import UpscalerModalApp


class UpscalerTool(Tool):
    name = "upscaler"
    description = """
        Perform upscaling on images.
        The "low_res_imgs" are PIL images.
        The "prompts" are strings.
        The output is a list of PIL images.
        You can upscale multiple images at once.
    """

    inputs = {
        "low_res_imgs": {
            "type": "array",
            "description": "The low resolution images to upscale",
        },
        "prompts": {
            "type": "array",
            "description": "The prompts to upscale the images",
        },
    }
    output_type = "object"

    def __init__(self):
        super().__init__()
        tool_class = modal.Cls.from_name(app.name, UpscalerModalApp.__name__)
        self.tool = tool_class()

    def forward(self, low_res_imgs: list[AgentImage], prompts: list[str]):
        upscaled_images = self.tool.forward.map(low_res_imgs, prompts)
        return list(upscaled_images)