# infra/detectors/yolo_v8.py
import cv2
from typing import Iterable, List, Optional, Set
from ultralytics import YOLO
from types.detect import Detection, Box

class YoloV8Detector:
    def __init__(self, weights: str = "yolov8n.pt", conf: float = 0.35, iou: float = 0.5,
                allowed_classes: Optional[Iterable[str]] = None):
        self.model = YOLO(weights)            # 初回は自動DL
        self.conf = conf
        self.iou = iou
        self.allowed: Optional[Set[str]] = set(allowed_classes) if allowed_classes else None
        # names は list or dict のことがある
        self.names = getattr(self.model, "names", None)

    def detect(self, image_bgr) -> List[Detection]:
        # UltralyticsはRGB想定
        res = self.model.predict(image_bgr[..., ::-1], conf=self.conf, iou=self.iou, verbose=False)[0]
        names = self.model.model.names if hasattr(self.model, "model") and hasattr(self.model.model, "names") else self.model.names
        dets: List[Detection] = []
        for b in res.boxes:
            cls_id = int(b.cls[0])
            label = names[cls_id] if isinstance(names, (list, tuple)) else names.get(cls_id, str(cls_id))
            if self.allowed and label not in self.allowed:
                continue
            x1, y1, x2, y2 = map(int, b.xyxy[0])
            score = float(b.conf[0])
            dets.append(Detection(label=label, score=score, box=Box(x1, y1, x2, y2)))
        return dets

    @staticmethod
    def draw(image_bgr, detections: Iterable[Detection]):
        out = image_bgr.copy()
        for d in detections:
            x1, y1, x2, y2 = d.box.x1, d.box.y1, d.box.x2, d.box.y2
            cv2.rectangle(out, (x1, y1), (x2, y2), (0, 200, 0), 2)
            cv2.putText(out, f"{d.label} {d.score:.2f}", (x1, max(0, y1 - 6)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 200, 0), 2)
        return out