"""Ultralytics YOLO detector.

Wraps an Ultralytics YOLO model (detection or segmentation) behind the
:class:`Detector` interface. Designed for the LADA mosaic-detection weights
(https://huggingface.co/ladaapp/lada), which are YOLO *segmentation* models with
a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose class
names map onto :class:`CensorType`.

Heavy imports (``ultralytics``/``torch``) happen lazily in ``__init__`` so the
rest of the app never pulls them in until detection actually runs.

Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README.
"""

from __future__ import annotations

import os

from ...config import DetectionConfig
from ..video.frame import Frame
from .base import Detector
from .types import CensorType, Detection


def _name_to_type(name: str) -> CensorType:
    n = name.lower()
    if "mosaic" in n or "pixel" in n:
        return CensorType.MOSAIC
    if "blur" in n:
        return CensorType.BLUR
    if "bar" in n or "black" in n:
        return CensorType.BLACK_BAR
    return CensorType.UNKNOWN


class YoloDetector(Detector):
    def __init__(
        self, model_path: str, config: DetectionConfig | None = None, label: str = ""
    ) -> None:
        self.cfg = config or DetectionConfig()
        self._label = label  # category (models/yolo/<label>) tagged onto every detection
        if not os.path.isfile(model_path):
            raise FileNotFoundError(
                f"Файл весов не найден: {model_path}\n"
                "Скачайте модель детекции мозаики LADA, например:\n"
                "  curl.exe -L -o models\\lada_mosaic_detection_model_v4_accurate.pt "
                '"https://huggingface.co/ladaapp/lada/resolve/main/'
                'lada_mosaic_detection_model_v4_accurate.pt?download=true"'
            )
        try:
            from ultralytics import YOLO
        except ImportError as exc:  # pragma: no cover - environment dependent
            raise ImportError(
                "Не установлен ultralytics. Установите: pip install -e \".[yolo]\" "
                "(и PyTorch с CUDA отдельно — см. README)."
            ) from exc

        # Resolve the device: CUDA only when actually available, else CPU. A CPU-only
        # torch build raises "Torch not compiled with CUDA enabled" if asked for cuda,
        # so we never request it without a working GPU (even on an explicit override).
        try:
            import torch

            cuda_ok = torch.cuda.is_available()
        except Exception:
            cuda_ok = False
        device = self.cfg.yolo_device
        if device is None:
            device = "cuda" if cuda_ok else "cpu"
        elif "cuda" in str(device) and not cuda_ok:
            device = "cpu"
        self._device = device
        self._model = YOLO(model_path)

    @property
    def name(self) -> str:
        tag = f", {self._label}" if self._label else ""
        return f"YoloDetector(device={self._device}{tag})"

    def detect(self, frame: Frame) -> list[Detection]:
        results = self._model.predict(
            source=frame.image,           # BGR ndarray; ultralytics handles it
            conf=self.cfg.yolo_conf,
            imgsz=self.cfg.yolo_imgsz,
            device=self._device,
            verbose=False,
        )
        if not results:
            return []
        res = results[0]
        boxes = getattr(res, "boxes", None)
        if boxes is None or len(boxes) == 0:
            return []

        names = res.names  # {class_index: class_name}
        xyxy = boxes.xyxy.cpu().numpy()
        confs = boxes.conf.cpu().numpy()
        classes = boxes.cls.cpu().numpy().astype(int)
        # Segmentation polygons in source-pixel coords, one per detection (if any).
        polygons = res.masks.xy if getattr(res, "masks", None) is not None else None

        out: list[Detection] = []
        for i in range(len(xyxy)):
            x1, y1, x2, y2 = xyxy[i]
            bbox = (int(x1), int(y1), int(x2 - x1), int(y2 - y1))
            poly: list[tuple[int, int]] = []
            if polygons is not None and i < len(polygons):
                poly = [(int(px), int(py)) for px, py in polygons[i]]
            ctype = _name_to_type(names.get(int(classes[i]), ""))
            out.append(Detection(
                type=ctype, score=float(confs[i]), bbox=bbox, polygon=poly, label=self._label
            ))
        return out