"""
通用推理后端

为 UltralyticsTracker 提供 RKNN / ONNX 模型的统一检测接口，
与安全检测（安全帽/反光衣）解耦。
"""

import os
import cv2
import numpy as np
from typing import List, Tuple, Dict, Any
from dataclasses import dataclass


@dataclass
class Detection:
    """检测结果 (用于 RKNN/ONNX 模型)"""
    class_id: int
    class_name: str
    confidence: float
    bbox: Tuple[int, int, int, int]


def nms(dets, iou_threshold=0.45):
    """非极大值抑制"""
    if len(dets) == 0:
        return []

    boxes = np.array([[d.bbox[0], d.bbox[1], d.bbox[2], d.bbox[3], d.confidence] for d in dets])
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    scores = boxes[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = 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:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= iou_threshold)[0]
        order = order[inds + 1]

    return [dets[i] for i in keep]


class BaseDetector:
    """检测器基类 (用于 RKNN/ONNX 模型)"""

    # 默认 COCO 类别映射；子类可覆盖
    LABEL_MAP = {0: 'person'}

    def __init__(self, label_map: Dict[int, str] = None):
        self.input_size = (640, 640)
        self.num_classes = len(label_map) if label_map else max(self.LABEL_MAP.keys()) + 1
        if label_map:
            self.LABEL_MAP = label_map

    def letterbox(self, image):
        """Letterbox 预处理，保持宽高比"""
        h0, w0 = image.shape[:2]
        ih, iw = self.input_size
        scale = min(iw / w0, ih / h0)
        new_w, new_h = int(w0 * scale), int(h0 * scale)
        pad_w = (iw - new_w) // 2
        pad_h = (ih - new_h) // 2
        resized = cv2.resize(image, (new_w, new_h))
        canvas = np.full((ih, iw, 3), 114, dtype=np.uint8)
        canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
        return canvas, scale, pad_w, pad_h, h0, w0

    def postprocess(self, outputs, scale, pad_w, pad_h, h0, w0, conf_threshold_map):
        """后处理"""
        dets = []

        if not outputs:
            return dets

        output = outputs[0]

        if len(output.shape) == 3:
            output = output[0]

        num_boxes = output.shape[1]

        for i in range(num_boxes):
            x_center = float(output[0, i])
            y_center = float(output[1, i])
            width = float(output[2, i])
            height = float(output[3, i])

            class_probs = output[4:4+self.num_classes, i]
            best_class = int(np.argmax(class_probs))
            confidence = float(class_probs[best_class])

            if best_class not in self.LABEL_MAP:
                continue

            conf_threshold = conf_threshold_map.get(best_class, 0.5)

            if confidence < conf_threshold:
                continue

            # 移除 padding 并缩放到原始图像尺寸
            x1 = int(((x_center - width / 2) - pad_w) / scale)
            y1 = int(((y_center - height / 2) - pad_h) / scale)
            x2 = int(((x_center + width / 2) - pad_w) / scale)
            y2 = int(((y_center + height / 2) - pad_h) / scale)

            x1 = max(0, min(w0, x1))
            y1 = max(0, min(h0, y1))
            x2 = max(0, min(w0, x2))
            y2 = max(0, min(h0, y2))

            det = Detection(
                class_id=best_class,
                class_name=self.LABEL_MAP[best_class],
                confidence=confidence,
                bbox=(x1, y1, x2, y2)
            )
            dets.append(det)

        dets = nms(dets, iou_threshold=0.45)
        return dets

    def detect(self, image, conf_threshold_map):
        raise NotImplementedError

    def release(self):
        pass


class RKNNDetector(BaseDetector):
    """RKNN 检测器 - 使用 NHWC 输入格式 (1, H, W, C)"""

    def __init__(self, model_path: str, label_map: Dict[int, str] = None):
        super().__init__(label_map=label_map)
        self.model_path = model_path
        self.rknn = None

        try:
            from rknnlite.api import RKNNLite
            self.rknn = RKNNLite()
        except ImportError:
            raise ImportError("未安装 rknnlite，请运行: pip install rknnlite2 或参考 testrk3588/setup_rknn.sh")

        ret = self.rknn.load_rknn(model_path)
        if ret != 0:
            raise RuntimeError(f"加载 RKNN 模型失败: {model_path}")

        ret = self.rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
        if ret != 0:
            raise RuntimeError("初始化 RKNN 运行时失败")

        print(f"RKNN 模型加载成功: {model_path}")

    def detect(self, image, conf_threshold_map):
        canvas, scale, pad_w, pad_h, h0, w0 = self.letterbox(image)
        # RKNN 期望 NHWC (1, H, W, C), RGB, 归一化 0-1
        img = canvas[..., ::-1].astype(np.float32) / 255.0
        blob = img[None, ...]  # (1, 640, 640, 3)
        outs = self.rknn.inference(inputs=[blob])
        return self.postprocess(outs, scale, pad_w, pad_h, h0, w0, conf_threshold_map)

    def release(self):
        if self.rknn:
            self.rknn.release()
            self.rknn = None


class ONNXDetector(BaseDetector):
    """ONNX 检测器 - 使用 NCHW 输入格式 (1, C, H, W)"""

    def __init__(self, model_path: str, label_map: Dict[int, str] = None):
        super().__init__(label_map=label_map)
        self.model_path = model_path

        try:
            import onnxruntime as ort
            self.session = ort.InferenceSession(model_path)
            self.input_name = self.session.get_inputs()[0].name
            self.output_name = self.session.get_outputs()[0].name
            print(f"ONNX 模型加载成功: {model_path}")
        except ImportError:
            raise ImportError("未安装 onnxruntime，请运行: pip install onnxruntime")
        except Exception as e:
            raise RuntimeError(f"加载 ONNX 模型失败: {e}")

    def detect(self, image, conf_threshold_map):
        canvas, scale, pad_w, pad_h, h0, w0 = self.letterbox(image)
        # ONNX 期望 NCHW (1, C, H, W), RGB, 归一化 0-1
        img = canvas[..., ::-1].astype(np.float32) / 255.0
        img = img.transpose(2, 0, 1)
        blob = img[None, ...]  # (1, 3, 640, 640)
        outs = self.session.run([self.output_name], {self.input_name: blob})
        return self.postprocess(outs, scale, pad_w, pad_h, h0, w0, conf_threshold_map)

    def release(self):
        self.session = None