"""Ultralytics YOLO adapter for registered detection model artifacts."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from importlib import import_module
from pathlib import Path
from typing import Any

from pydantic import BaseModel, ConfigDict, Field, field_validator

from cmvr_edge_ai.contracts import BoundingBox, Detection, ImageFrame

from ..base import DetectionModel


class YoloDependencyError(RuntimeError):
    """Raised when the optional YOLO runtime has not been installed."""


class YoloModelOptions(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    weights: str = Field(min_length=1)
    device: int | str = "cpu"
    imgsz: int = Field(default=640, ge=32, le=8192)
    iou: float = Field(default=0.7, ge=0.0, le=1.0)
    half: bool = False
    max_det: int = Field(default=300, ge=1, le=100_000)
    agnostic_nms: bool = False

    @field_validator("weights")
    @classmethod
    def normalize_weights(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("weights must not be empty")
        return value

    @field_validator("device")
    @classmethod
    def validate_device(cls, value: int | str) -> int | str:
        if isinstance(value, int):
            if value < 0:
                raise ValueError("device must be non-negative")
            return value
        value = value.strip()
        if not value:
            raise ValueError("device must not be empty")
        return value


class UltralyticsYoloModel(DetectionModel):
    """Normalize an Ultralytics detector behind the framework model API."""

    def __init__(
        self,
        params: Mapping[str, Any],
        *,
        expected_labels: tuple[str, ...],
    ) -> None:
        try:
            self._options = YoloModelOptions.model_validate(dict(params))
        except Exception as exc:
            raise ValueError(f"invalid YOLO model options: {exc}") from exc
        self._expected_labels = expected_labels
        self._label_to_id = {
            label: index for index, label in enumerate(self._expected_labels)
        }
        self._model: Any = None
        self._numpy: Any = None

    @property
    def weights_path(self) -> Path:
        return Path(self._options.weights).expanduser()

    def load(self) -> None:
        weights = self.weights_path
        if not weights.is_file():
            raise FileNotFoundError(f"YOLO weights do not exist: {weights}")
        try:
            np = import_module("numpy")
            YOLO = getattr(import_module("ultralytics"), "YOLO")
        except ImportError as exc:
            raise YoloDependencyError(
                "YOLO support is not installed; install "
                "cmvr-edge-ai[yolo]. The PPE checkpoint was created with "
                "Ultralytics 8.4.31, so validate the deployed 8.x runtime."
            ) from exc

        # Ultralytics .pt files use the PyTorch checkpoint loader.  Deploy only
        # weights obtained from a trusted build/artifact source.
        model = YOLO(str(weights), task="detect")
        actual_labels = _ordered_model_names(model.names)
        if actual_labels != self._expected_labels:
            raise ValueError(
                "YOLO checkpoint labels do not match the registered model: "
                f"expected {self._expected_labels!r}, got {actual_labels!r}"
            )
        self._numpy = np
        self._model = model

    def predict(
        self,
        frame: ImageFrame,
        labels: tuple[str, ...],
        confidence: float,
    ) -> Sequence[Detection]:
        if self._model is None or self._numpy is None:
            raise RuntimeError("YOLO model has not been loaded")
        if frame.is_encoded:
            raise ValueError("YOLO requires a decoded ImageFrame")
        pixel_format = frame.pixel_format.strip().upper()
        if pixel_format not in {"BGR8", "RGB8"}:
            raise ValueError(
                f"YOLO adapter supports BGR8 or RGB8, got {frame.pixel_format!r}"
            )
        expected_size = frame.width * frame.height * 3
        if frame.width <= 0 or frame.height <= 0 or len(frame.data) != expected_size:
            raise ValueError(
                "decoded frame buffer size does not match width*height*3: "
                f"got {len(frame.data)}, expected {expected_size}"
            )

        image = self._numpy.frombuffer(frame.data, dtype=self._numpy.uint8).reshape(
            (frame.height, frame.width, 3)
        )
        if pixel_format == "RGB8":
            # Ultralytics ndarray inputs follow the OpenCV BGR convention.
            image = image[:, :, ::-1]

        class_ids = [self._label_to_id[label] for label in labels]
        results = self._model.predict(
            source=image,
            classes=class_ids,
            conf=confidence,
            iou=self._options.iou,
            imgsz=self._options.imgsz,
            device=self._options.device,
            half=self._options.half,
            max_det=self._options.max_det,
            agnostic_nms=self._options.agnostic_nms,
            verbose=False,
        )
        if not results:
            return ()
        boxes = results[0].boxes
        if boxes is None:
            return ()

        xyxy_rows = _to_python_list(boxes.xyxy)
        confidences = _to_python_list(boxes.conf)
        class_indexes = _to_python_list(boxes.cls)
        detections: list[Detection] = []
        selected = frozenset(labels)
        for coordinates, score, class_index in zip(
            xyxy_rows, confidences, class_indexes
        ):
            index = int(class_index)
            if not 0 <= index < len(self._expected_labels):
                raise ValueError(f"YOLO returned an unknown class id: {index}")
            label = self._expected_labels[index]
            if label not in selected:
                continue
            if len(coordinates) != 4:
                raise ValueError("YOLO xyxy result must contain four coordinates")
            detections.append(
                Detection(
                    label=label,
                    confidence=float(score),
                    box=BoundingBox(*(float(value) for value in coordinates)),
                )
            )
        return tuple(detections)

    def close(self) -> None:
        self._model = None
        self._numpy = None


def _ordered_model_names(value: Any) -> tuple[str, ...]:
    if isinstance(value, Mapping):
        try:
            indexes = sorted(int(key) for key in value)
            if indexes != list(range(len(indexes))):
                raise ValueError("class IDs must be contiguous and start at zero")
            return tuple(str(value[index]) for index in indexes)
        except (KeyError, TypeError, ValueError) as exc:
            raise ValueError(f"invalid YOLO model.names mapping: {value!r}") from exc
    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
        return tuple(str(label) for label in value)
    raise ValueError(f"invalid YOLO model.names value: {value!r}")


def _to_python_list(value: Any) -> list[Any]:
    if hasattr(value, "detach"):
        value = value.detach()
    if hasattr(value, "cpu"):
        value = value.cpu()
    if hasattr(value, "tolist"):
        value = value.tolist()
    return list(value)


__all__ = [
    "UltralyticsYoloModel",
    "YoloDependencyError",
    "YoloModelOptions",
]
