from __future__ import annotations

import re
from typing import Any, Protocol

from .models import Detection


def parse_class_prompt(class_prompt: str | list[str] | tuple[str, ...]) -> list[str]:
    if isinstance(class_prompt, (list, tuple)):
        parts = [str(item) for item in class_prompt]
    else:
        parts = re.split(r"[,;\n]+", class_prompt)

    seen: set[str] = set()
    classes: list[str] = []
    for part in parts:
        label = " ".join(part.strip().split())
        key = label.lower()
        if label and key not in seen:
            seen.add(key)
            classes.append(label)
    return classes


class Detector(Protocol):
    class_names: list[str]

    def detect(
        self,
        frame: Any,
        *,
        frame_index: int,
        timestamp_sec: float,
        confidence: float,
        image_size: int | None = None,
        max_detections: int | None = None,
    ) -> list[Detection]:
        ...


class UltralyticsYOLOEDetector:
    def __init__(
        self,
        *,
        class_names: list[str],
        model_name: str = "yoloe-26s-seg.pt",
        device: str | None = None,
    ) -> None:
        if not class_names:
            raise ValueError("YOLOE needs at least one open-vocabulary class.")

        try:
            from ultralytics import YOLOE
        except ImportError as exc:  # pragma: no cover - optional heavy dependency
            raise RuntimeError("Install ultralytics to use the YOLOE detector.") from exc

        self.class_names = class_names
        self.model_name = model_name
        self.device = device
        self.model = YOLOE(model_name)
        self.model.set_classes(class_names)

    def detect(
        self,
        frame: Any,
        *,
        frame_index: int,
        timestamp_sec: float,
        confidence: float,
        image_size: int | None = None,
        max_detections: int | None = None,
    ) -> list[Detection]:
        kwargs: dict[str, Any] = {"conf": confidence, "verbose": False}
        if self.device:
            kwargs["device"] = self.device
        if image_size:
            kwargs["imgsz"] = image_size
        if max_detections:
            kwargs["max_det"] = max_detections

        results = self.model.predict(frame, **kwargs)
        if not results:
            return []
        return detections_from_ultralytics_result(
            results[0],
            frame_shape=frame.shape,
            frame_index=frame_index,
            timestamp_sec=timestamp_sec,
            fallback_names=self.class_names,
        )


def detections_from_ultralytics_result(
    result: Any,
    *,
    frame_shape: tuple[int, int, int],
    frame_index: int,
    timestamp_sec: float,
    fallback_names: list[str],
) -> list[Detection]:
    boxes = getattr(result, "boxes", None)
    if boxes is None or boxes.xyxy is None:
        return []

    height, width = frame_shape[:2]
    xyxy_values = boxes.xyxy.cpu().tolist()
    confidences = boxes.conf.cpu().tolist()
    class_ids = boxes.cls.cpu().tolist()
    names = getattr(result, "names", {}) or {}

    detections: list[Detection] = []
    for bbox, score, class_id in zip(xyxy_values, confidences, class_ids):
        label = _label_from_names(names, int(class_id), fallback_names)
        detections.append(
            Detection(
                frame_index=frame_index,
                timestamp_sec=timestamp_sec,
                label=label,
                confidence=float(score),
                bbox_xyxy=tuple(float(value) for value in bbox),
                bbox_xyxy_norm=_normalize_box(bbox, width, height),
            )
        )
    return detections


def _label_from_names(names: Any, class_id: int, fallback_names: list[str]) -> str:
    if isinstance(names, dict) and class_id in names:
        return str(names[class_id])
    if isinstance(names, list) and 0 <= class_id < len(names):
        return str(names[class_id])
    if 0 <= class_id < len(fallback_names):
        return fallback_names[class_id]
    return f"class_{class_id}"


def _normalize_box(bbox: list[float], width: int, height: int) -> tuple[float, float, float, float]:
    x1, y1, x2, y2 = bbox
    if width <= 0 or height <= 0:
        return (0.0, 0.0, 0.0, 0.0)
    return (
        _clamp01(float(x1) / width),
        _clamp01(float(y1) / height),
        _clamp01(float(x2) / width),
        _clamp01(float(y2) / height),
    )


def _clamp01(value: float) -> float:
    return max(0.0, min(1.0, value))