import typing
import ultralytics

YOLO_V10_MODELS = {
    "nano": "jameslahm/yolov10n",
    "small": "jameslahm/yolov10s",
    "medium": "jameslahm/yolov10m",
    "base": "jameslahm/yolov10b",
    "large": "jameslahm/yolov10l",
    "xlarge": "jameslahm/yolov10x",
}


class YOLOv10Plugin:
    def __init__(
        self,
        yolo_model_name: (
            str
            | typing.Literal[
                "nano",
                "small",
                "medium",
                "base",
                "large",
                "xlarge",
            ]
        ) = "nano",
        verbose: bool = True,
    ):
        assert (
            yolo_model_name in YOLO_V10_MODELS.keys()
        ), f"`yolo_model_name` should be either one of {list(YOLO_V10_MODELS.keys())}"
        self.yolo_model_name = yolo_model_name
        self.model = ultralytics.YOLOv10.from_pretrained(
            YOLO_V10_MODELS[yolo_model_name]
        )

        self.verbose = verbose
        if self.verbose:
            print(f"YOLOv10Plugin::__init__::Model Name: {self.yolo_model_name}")

    def detect(self, image):
        results = self.model(image)
        results = results[0].summary()

        out = []
        for result in results:
            out.append(
                {
                    "name": result["name"],
                    "class": result["class"],
                    "confidence": result["confidence"],
                    "box": [
                        int(result["box"]["x1"]),
                        int(result["box"]["y1"]),
                        int(result["box"]["x2"]),
                        int(result["box"]["y2"]),
                    ],
                }
            )

        return out


if __name__ == "__main__":
    yolo = YOLOv10Plugin()
    yolo.detect("https://ultralytics.com/images/zidane.jpg")