import io
import os
from abc import ABC, abstractmethod

from PIL import Image


def _label_set(labels) -> set[str]:
    return {label.strip().lower() for label in labels if label.strip()}


def abnormal_probability(detections: list[dict], labels: tuple[str, ...] | list[str]) -> float:
    """单帧异常概率：取异常类检测的最高置信度，避免多框求和造成虚高。"""
    abnormal = _label_set(labels) - {"healthy"}
    best = 0.0
    for detection in detections:
        class_name = str(detection.get("class_name", "")).strip().lower()
        if class_name in ("", "healthy", "unknown"):
            continue
        if abnormal and class_name not in abnormal:
            continue
        try:
            confidence = float(detection.get("confidence", 0))
        except (TypeError, ValueError):
            confidence = 0.0
        best = max(best, min(1.0, max(0.0, confidence)))
    return best


def detection_status(detections: list[dict], labels: tuple[str, ...] | list[str]) -> str:
    """空检测或 unknown 类不再被当作 healthy。"""
    if not detections:
        return "unknown"
    abnormal = _label_set(labels) - {"healthy"}
    healthy_seen = False
    unknown_seen = False
    for detection in detections:
        class_name = str(detection.get("class_name", "")).strip().lower()
        if abnormal and class_name in abnormal:
            return "abnormal"
        if not abnormal and class_name not in ("", "healthy", "unknown"):
            return "abnormal"
        if class_name == "healthy":
            healthy_seen = True
        elif class_name in ("", "unknown"):
            unknown_seen = True
        else:
            unknown_seen = True
    if healthy_seen and not unknown_seen:
        return "healthy"
    return "unknown"


class Detector(ABC):
    @abstractmethod
    def detect(self, image_bytes: bytes) -> list[dict]:
        """返回 [{"bbox": {x,y,w,h}, "class_name", "confidence"}]"""


class MockDetector(Detector):
    """开发联调用：校验图片可解析后返回固定检测结果。"""

    def __init__(self, class_name: str = "healthy", confidence: float = 0.95):
        self.class_name = class_name
        self.confidence = confidence

    def detect(self, image_bytes: bytes) -> list[dict]:
        try:
            with Image.open(io.BytesIO(image_bytes)) as img:
                img.verify()
        except Exception as exc:
            raise ValueError("无法解析图片") from exc
        return [
            {
                "bbox": {"x": 10, "y": 10, "w": 100, "h": 100},
                "class_name": self.class_name,
                "confidence": self.confidence,
            }
        ]


class ONNXDetector(Detector):
    """YOLOv8 导出的 best.onnx 推理（CPU）。

    YOLO 训练恢复后放置模型文件，MODEL_MODE=onnx 生效。
    输出张量按 YOLOv8 常见格式 [1, 4+nc, 8400]（或转置）解析，含 NMS。
    """

    def __init__(
        self,
        model_path: str,
        labels: tuple[str, ...],
        input_size: int = 640,
        conf_threshold: float = 0.25,
        iou_threshold: float = 0.45,
    ):
        import numpy as np
        import onnxruntime as ort

        if not os.path.exists(model_path):
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        self.session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
        self.labels = labels
        self.input_size = input_size
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.np = np

    def detect(self, image_bytes: bytes) -> list[dict]:
        np = self.np
        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        orig_w, orig_h = img.size
        scale = min(self.input_size / orig_w, self.input_size / orig_h)
        new_w, new_h = int(round(orig_w * scale)), int(round(orig_h * scale))
        resized = img.resize((new_w, new_h))
        canvas = Image.new("RGB", (self.input_size, self.input_size), (114, 114, 114))
        pad_x = (self.input_size - new_w) // 2
        pad_y = (self.input_size - new_h) // 2
        canvas.paste(resized, (pad_x, pad_y))
        blob = np.asarray(canvas, dtype=np.float32) / 255.0
        blob = blob.transpose(2, 0, 1)[None, ...]
        out = self.session.run(None, {self.session.get_inputs()[0].name: blob})[0]

        nc = len(self.labels)
        if out.shape[1] != 4 + nc:
            out = out.transpose(0, 2, 1)
        out = out[0]
        if out.shape[0] == 4 + nc:
            out = out.T

        boxes, scores, class_ids = [], [], []
        for row in out:
            cx, cy, w, h = row[:4]
            cls_scores = row[4 : 4 + nc]
            cls_id = int(cls_scores.argmax())
            score = float(cls_scores[cls_id])
            if score < self.conf_threshold:
                continue
            boxes.append(
                [
                    (cx - w / 2 - pad_x) / scale,
                    (cy - h / 2 - pad_y) / scale,
                    w / scale,
                    h / scale,
                ]
            )
            scores.append(score)
            class_ids.append(cls_id)

        keep = self._nms(np.asarray(boxes), np.asarray(scores))
        return [
            {
                "bbox": {
                    "x": float(boxes[i][0]),
                    "y": float(boxes[i][1]),
                    "w": float(boxes[i][2]),
                    "h": float(boxes[i][3]),
                },
                "class_name": self.labels[class_ids[i]],
                "confidence": float(scores[i]),
            }
            for i in keep
        ]

    def _nms(self, boxes, scores):
        np = self.np
        if len(boxes) == 0:
            return []
        x1 = boxes[:, 0]
        y1 = boxes[:, 1]
        x2 = boxes[:, 0] + boxes[:, 2]
        y2 = boxes[:, 1] + boxes[:, 3]
        areas = boxes[:, 2] * boxes[:, 3]
        order = scores.argsort()[::-1]
        keep = []
        while order.size > 0:
            i = int(order[0])
            keep.append(i)
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
            order = order[1:][iou <= self.iou_threshold]
        return keep