import supervision as sv
from smolagents import AgentImage, Tool


class BBoxAnnotatorTool(Tool):
    name = "bbox_annotator"
    description = """
        Given an image and a list of detections, draw the bounding boxes on the image.
        The image is a PIL image.
        The detections are an object of type supervision.Detections. You can use the task_inference_output_converter tool to obtain the proper format for the detections.
        The output is the image with the bounding boxes drawn on it.
    """

    inputs = {
        "image": {
            "type": "image",
            "description": "The image to draw the bounding boxes on",
        },
        "detections": {
            "type": "object",
            "description": """
            The detections to annotate on the image. 
            The detections are an object of type supervision.Detections.
            You can use the task_inference_output_converter tool to obtain the proper format for the detections.
            """,
        },
        "thickness": {
            "type": "number",
            "description": "The thickness of the bounding boxes.",
            "default": 5,
            "nullable": True,
        },
    }
    output_type = "image"

    def __init__(self):
        super().__init__()

    def forward(
        self,
        image: AgentImage,
        detections: sv.Detections,
        thickness: int = 5,
    ):
        box_annotator = sv.BoxAnnotator(thickness=thickness)
        annotated_image = box_annotator.annotate(scene=image, detections=detections)
        return annotated_image