"""TensorRT YOLO output adapters for multiple runtime output styles."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import numpy as np
import torch


if TYPE_CHECKING:
    from .schemas import DetectionPrediction


@dataclass(frozen=True)
class TensorRTYOLOAdapterConfig:
    """Configuration for adapting TensorRT YOLO detections."""

    conf_threshold: float = 0.25
    iou_threshold: float = 0.45
    max_detections: int = 300
    boxes_normalized: bool | None = None
    apply_nms: bool = True
    strict: bool = False


def adapt_tensorrt_yolo_output(
    output: Any,
    *,
    image_width: int,
    image_height: int,
    config: TensorRTYOLOAdapterConfig | None = None,
) -> list[DetectionPrediction]:
    """Normalize YOLO detections from TensorRT/trtexec/Ultralytics outputs.

    Supported styles:
    - Ultralytics results list (result.boxes.xyxy/conf/cls)
    - TensorRT NMS plugin output (`num_dets`, `boxes`, `scores`, `classes`)
    - Dense detection matrix shaped [N, 6+] with `[x1,y1,x2,y2,score,cls,...]`
    - Raw head-like matrix [N, 5 + num_classes] or [N, 4 + num_classes]

    Args:
        output: Runtime-specific detection output object.
        image_width: Source image width in pixels.
        image_height: Source image height in pixels.
        config: Optional adapter configuration overrides.

    Returns:
        List of normalized detection prediction dictionaries.
    """
    adapter_config = config or TensorRTYOLOAdapterConfig()
    try:
        detections = _coerce_to_detection_matrix(output)
    except Exception:
        if adapter_config.strict:
            raise
        return []

    if detections.numel() == 0:
        return []

    detections = _filter_by_confidence(detections, threshold=adapter_config.conf_threshold)
    if detections.numel() == 0:
        return []

    boxes = detections[:, 0:4].clone()
    scores = detections[:, 4].clone()
    classes = detections[:, 5].clone()

    normalized = (
        adapter_config.boxes_normalized
        if adapter_config.boxes_normalized is not None
        else _looks_normalized(boxes)
    )
    if normalized:
        scale = torch.tensor(
            [image_width, image_height, image_width, image_height],
            dtype=boxes.dtype,
            device=boxes.device,
        )
        boxes = boxes * scale

    boxes[:, 0::2] = boxes[:, 0::2].clamp(0, max(0, image_width - 1))
    boxes[:, 1::2] = boxes[:, 1::2].clamp(0, max(0, image_height - 1))

    if adapter_config.apply_nms:
        keep = _class_aware_nms(
            boxes,
            scores,
            classes,
            iou_threshold=adapter_config.iou_threshold,
            max_detections=adapter_config.max_detections,
        )
        boxes = boxes[keep]
        scores = scores[keep]
        classes = classes[keep]

    return [
        {
            "category": str(int(cls.item())),
            "confidence": float(score.item()),
            "bounding_box": {
                "xmin": int(box[0].item()),
                "ymin": int(box[1].item()),
                "xmax": int(box[2].item()),
                "ymax": int(box[3].item()),
            },
        }
        for box, score, cls in zip(boxes, scores, classes, strict=True)
    ]


def _to_tensor(value: Any) -> torch.Tensor:
    """Convert tensor-like object to CPU torch tensor.

    Args:
        value: Tensor-like input value.

    Returns:
        CPU tensor representation of ``value``.
    """
    if isinstance(value, torch.Tensor):
        return value.detach().to("cpu")
    if hasattr(value, "cpu") and hasattr(value, "numpy"):
        return torch.from_numpy(value.cpu().numpy())
    if isinstance(value, np.ndarray):
        return torch.from_numpy(value)
    if isinstance(value, (list, tuple)):
        return torch.as_tensor(value)
    raise TypeError(f"Unsupported tensor-like value type: {type(value).__name__}")


def _from_ultralytics_results(output: Any) -> torch.Tensor | None:
    """Extract ``[N, 6]`` detections from Ultralytics results-style objects.

    Args:
        output: Candidate Ultralytics results sequence.

    Returns:
        Detection matrix tensor, empty tensor, or ``None`` when unsupported.
    """
    if not isinstance(output, Sequence) or isinstance(output, (str, bytes)):
        return None
    if not output:
        return torch.zeros((0, 6), dtype=torch.float32)

    rows: list[torch.Tensor] = []
    for item in output:
        boxes_obj = getattr(item, "boxes", None)
        if boxes_obj is None:
            return None
        try:
            xyxy = _to_tensor(boxes_obj.xyxy)
            conf = _to_tensor(boxes_obj.conf).reshape(-1, 1)
            cls = _to_tensor(boxes_obj.cls).reshape(-1, 1)
        except AttributeError:
            return None
        if xyxy.numel() == 0:
            continue
        rows.append(torch.cat([xyxy.float(), conf.float(), cls.float()], dim=1))

    if not rows:
        return torch.zeros((0, 6), dtype=torch.float32)
    return torch.cat(rows, dim=0)


def _from_trt_nms_mapping(data: Mapping[str, Any]) -> torch.Tensor | None:
    """Extract ``[N, 6]`` matrix from TensorRT NMS mapping output.

    Args:
        data: Mapping containing TensorRT plugin outputs.

    Returns:
        Detection matrix tensor or ``None`` when required keys are missing.
    """
    key_sets = [
        ("num_dets", "boxes", "scores", "classes"),
        ("num", "boxes", "scores", "classes"),
        ("dets", "boxes", "scores", "labels"),
    ]
    keys = next((candidate for candidate in key_sets if all(k in data for k in candidate)), None)
    if keys is None:
        return None

    num_key, boxes_key, scores_key, cls_key = keys
    num_dets_tensor = _to_tensor(data[num_key]).reshape(-1)
    det_count = int(num_dets_tensor[0].item()) if num_dets_tensor.numel() else 0
    boxes = _to_tensor(data[boxes_key]).float().reshape(-1, 4)
    scores = _to_tensor(data[scores_key]).float().reshape(-1, 1)
    classes = _to_tensor(data[cls_key]).float().reshape(-1, 1)
    matrix = torch.cat([boxes, scores, classes], dim=1)
    if det_count > 0:
        return matrix[:det_count]
    return matrix


def _from_trt_nms_sequence(output: Sequence[Any]) -> torch.Tensor | None:
    """Extract ``[N, 6]`` matrix from TensorRT NMS tuple/list output.

    Args:
        output: Ordered sequence containing TensorRT NMS outputs.

    Returns:
        Detection matrix tensor or ``None`` when coercion fails.
    """
    if len(output) < 4:
        return None

    try:
        num_dets = int(_to_tensor(output[0]).reshape(-1)[0].item())
        boxes = _to_tensor(output[1]).float().reshape(-1, 4)
        scores = _to_tensor(output[2]).float().reshape(-1, 1)
        classes = _to_tensor(output[3]).float().reshape(-1, 1)
    except Exception:
        return None

    matrix = torch.cat([boxes, scores, classes], dim=1)
    return matrix[:num_dets] if num_dets > 0 else matrix


def _coerce_dense_matrix(raw: torch.Tensor) -> torch.Tensor:
    """Convert dense detection tensors to unified ``[N, 6]`` format.

    Args:
        raw: Raw dense detection tensor.

    Returns:
        Detection matrix tensor with columns ``[x1, y1, x2, y2, conf, cls]``.
    """
    matrix = raw.float()
    if matrix.ndim == 3 and matrix.shape[0] == 1:
        matrix = matrix.squeeze(0)
    if matrix.ndim != 2:
        raise ValueError(f"Expected 2D detection matrix, got shape {tuple(matrix.shape)}")

    if matrix.shape[0] < matrix.shape[1] and matrix.shape[0] <= 16 and matrix.shape[1] >= 64:
        matrix = matrix.transpose(0, 1)

    if matrix.shape[1] == 6:
        return matrix
    if matrix.shape[1] > 6:
        # Style A: [x1,y1,x2,y2,obj,cls_0,...,cls_n]
        # Style B: [x1,y1,x2,y2,cls_0,...,cls_n]
        boxes = matrix[:, 0:4]
        remainder = matrix[:, 4:]
        if remainder.shape[1] >= 2:
            obj = remainder[:, 0:1]
            class_scores = remainder[:, 1:]
            class_conf, class_idx = torch.max(class_scores, dim=1, keepdim=True)
            conf = obj * class_conf
            return torch.cat([boxes, conf, class_idx.float()], dim=1)

        class_conf, class_idx = torch.max(remainder, dim=1, keepdim=True)
        return torch.cat([boxes, class_conf, class_idx.float()], dim=1)

    raise ValueError(f"Unsupported dense detection matrix shape: {tuple(matrix.shape)}")


def _coerce_to_detection_matrix(output: Any) -> torch.Tensor:
    """Adapt supported output types to unified ``[N, 6]`` tensor.

    Args:
        output: Runtime-specific detection output object.

    Returns:
        Detection matrix tensor with unified column layout.
    """
    ultra = _from_ultralytics_results(output)
    if ultra is not None:
        return ultra

    if isinstance(output, Mapping):
        mapping_matrix = _from_trt_nms_mapping(output)
        if mapping_matrix is not None:
            return mapping_matrix

    if isinstance(output, Sequence) and not isinstance(output, (str, bytes)):
        seq_matrix = _from_trt_nms_sequence(output)
        if seq_matrix is not None:
            return seq_matrix

    return _coerce_dense_matrix(_to_tensor(output))


def _filter_by_confidence(detections: torch.Tensor, *, threshold: float) -> torch.Tensor:
    """Filter detection rows by confidence threshold.

    Args:
        detections: Detection matrix tensor with confidence at column index 4.
        threshold: Minimum confidence required to keep a row.

    Returns:
        Filtered detection matrix tensor.
    """
    if detections.numel() == 0:
        return detections.reshape(0, 6)
    mask = detections[:, 4] >= float(threshold)
    return detections[mask]


def _looks_normalized(boxes: torch.Tensor) -> bool:
    """Infer whether box coordinates appear normalized.

    Args:
        boxes: Box tensor expected in ``xyxy`` format.

    Returns:
        ``True`` when values look normalized to ``[0, 1]`` scale.
    """
    if boxes.numel() == 0:
        return True
    max_abs = float(torch.max(torch.abs(boxes)).item())
    return max_abs <= 1.5


def _box_iou(box: torch.Tensor, boxes: torch.Tensor) -> torch.Tensor:
    """Compute IoU between one box and many boxes.

    Args:
        box: Single ``xyxy`` box tensor.
        boxes: Tensor of ``xyxy`` boxes to compare against.

    Returns:
        IoU tensor for ``box`` against each row in ``boxes``.
    """
    xx1 = torch.maximum(box[0], boxes[:, 0])
    yy1 = torch.maximum(box[1], boxes[:, 1])
    xx2 = torch.minimum(box[2], boxes[:, 2])
    yy2 = torch.minimum(box[3], boxes[:, 3])

    inter_w = (xx2 - xx1).clamp(min=0)
    inter_h = (yy2 - yy1).clamp(min=0)
    inter = inter_w * inter_h

    area1 = (box[2] - box[0]).clamp(min=0) * (box[3] - box[1]).clamp(min=0)
    area2 = (boxes[:, 2] - boxes[:, 0]).clamp(min=0) * (boxes[:, 3] - boxes[:, 1]).clamp(min=0)
    union = area1 + area2 - inter
    return inter / union.clamp(min=1e-9)


def _nms_indices(boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor:
    """Apply pure-torch NMS and return kept indices.

    Args:
        boxes: Tensor of candidate ``xyxy`` boxes.
        scores: Confidence scores aligned with ``boxes``.
        iou_threshold: IoU suppression threshold.

    Returns:
        Tensor of kept box indices.
    """
    if boxes.numel() == 0:
        return torch.empty((0,), dtype=torch.long)
    order = torch.argsort(scores, descending=True)
    keep: list[int] = []
    while order.numel() > 0:
        i = int(order[0].item())
        keep.append(i)
        if order.numel() == 1:
            break
        rest = order[1:]
        iou = _box_iou(boxes[i], boxes[rest])
        order = rest[iou <= iou_threshold]
    return torch.as_tensor(keep, dtype=torch.long)


def _class_aware_nms(
    boxes: torch.Tensor,
    scores: torch.Tensor,
    classes: torch.Tensor,
    *,
    iou_threshold: float,
    max_detections: int,
) -> torch.Tensor:
    """Apply class-aware NMS and return global kept indices.

    Args:
        boxes: Tensor of candidate ``xyxy`` boxes.
        scores: Confidence scores aligned with ``boxes``.
        classes: Class ids aligned with ``boxes``.
        iou_threshold: IoU suppression threshold for each class.
        max_detections: Maximum detections retained after sorting.

    Returns:
        Tensor of global kept indices ordered by descending score.
    """
    keep_chunks: list[torch.Tensor] = []
    for class_id in torch.unique(classes):
        class_mask = classes == class_id
        class_indices = torch.where(class_mask)[0]
        class_keep_local = _nms_indices(boxes[class_mask], scores[class_mask], iou_threshold)
        keep_chunks.append(class_indices[class_keep_local])

    if not keep_chunks:
        return torch.empty((0,), dtype=torch.long)

    keep = torch.cat(keep_chunks, dim=0)
    keep = keep[torch.argsort(scores[keep], descending=True)]
    return keep[:max_detections]
