"""YOLO-based detectors for sacks and trucks.

Each detector is a single-responsibility unit (S).  New model types can be
added as new classes without touching these (O).
"""

from __future__ import annotations

import numpy as np
from ultralytics import YOLO

from src.interfaces import Detection


class SackDetector:
    """Detects sacks (and persons) using a YOLO segmentation model."""

    def __init__(self, model_path: str, conf: float = 0.40) -> None:
        self._model = YOLO(model_path)
        self._conf = conf

    def detect(self, frame: np.ndarray) -> list[Detection]:
        results = self._model.predict(
            frame, conf=self._conf, verbose=False
        )
        return self._parse(results[0])

    def _parse(self, result) -> list[Detection]:
        detections: list[Detection] = []
        masks = result.masks
        for i, box in enumerate(result.boxes):
            cls_id = int(box.cls[0])
            name = self._model.names[cls_id]
            if name != "sack":
                continue
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            mask = None
            if masks is not None and i < len(masks):
                mask = masks[i].data.cpu().numpy().squeeze()
            detections.append(
                Detection(
                    bbox=(x1, y1, x2, y2),
                    confidence=float(box.conf[0]),
                    class_id=cls_id,
                    class_name=name,
                    mask=mask,
                )
            )
        return detections


class TruckDetector:
    """Detects trucks using a YOLO detection model."""

    def __init__(self, model_path_or_model: str | YOLO, conf: float = 0.50) -> None:
        if isinstance(model_path_or_model, str):
            self._model = YOLO(model_path_or_model)
        else:
            self._model = model_path_or_model
        self._conf = conf

    def detect(self, frame: np.ndarray) -> list[Detection]:
        results = self._model.predict(
            frame, conf=self._conf, verbose=False
        )
        return self._parse(results[0])

    def _parse(self, result) -> list[Detection]:
        detections: list[Detection] = []
        for box in result.boxes:
            cls_id = int(box.cls[0])
            name = self._model.names[cls_id]
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            detections.append(
                Detection(
                    bbox=(x1, y1, x2, y2),
                    confidence=float(box.conf[0]),
                    class_id=cls_id,
                    class_name=name,
                )
            )
        return detections