from aiforge.core.registry import register
from aiforge.core.context import BaseOperator

@register("detect_yolo")
class YOLODetect(BaseOperator):
    """Run object detection on an image using Ultralytics YOLOv8."""
    def run(self, data, model_name="yolov8n.pt"):
        from ultralytics import YOLO
        import logging
        
        logging.getLogger('ultralytics').setLevel(logging.WARNING)
        model = YOLO(model_name)
        # We pass data (image/numpy array) and perform predictions
        results = model(data)
        return results

@register("yolo_extract_labels")
class YOLOExtractLabels(BaseOperator):
    """Extract class labels from YOLO detection results (Bridge: YOLO -> NLP/ML)."""
    def run(self, data):
        labels = []
        for result in data:
            names = result.names
            for box in result.boxes:
                cls_id = int(box.cls[0])
                labels.append(names[cls_id])
        return labels

@register("yolo_draw_boxes")
class YOLODrawBoxes(BaseOperator):
    """Draw YOLO bounding boxes and labels on the original image."""
    def run(self, data):
        # We can just return result.plot() for the first result if data is a list of results.
        import cv2
        import numpy as np
        if isinstance(data, list) and len(data) > 0:
            # Ultralytics Results object has a plot() method that returns a BGR numpy array
            return data[0].plot()
        return data
