"""统一检测器接口：YOLO / Faster R-CNN，供推理与二阶段脚本共用。"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import torch

ROOT = Path(__file__).resolve().parent


@dataclass
class DetBox:
    xyxy: tuple[int, int, int, int]
    conf: float
    cls_id: int = 0


def import_yolo():
    """导入 YOLO。项目下 ultralytics/ 源码目录会挡住包，需优先指向仓库根。"""
    import sys

    repo = ROOT / "ultralytics"
    if (repo / "ultralytics" / "__init__.py").exists():
        repo_s = str(repo)
        while repo_s in sys.path:
            sys.path.remove(repo_s)
        sys.path.insert(0, repo_s)

    mod = sys.modules.get("ultralytics")
    if mod is not None and getattr(mod, "__file__", None) is None:
        for key in list(sys.modules):
            if key == "ultralytics" or key.startswith("ultralytics."):
                del sys.modules[key]

    from ultralytics import YOLO

    return YOLO


def resolve_device(device: str | int) -> torch.device:
    text = str(device).strip().lower()
    if text in {"cpu", "-1"}:
        return torch.device("cpu")
    if torch.cuda.is_available():
        try:
            return torch.device(f"cuda:{int(text)}")
        except Exception:
            return torch.device("cuda:0")
    return torch.device("cpu")


def build_faster_rcnn(num_classes: int, pretrained_backbone: bool = True):
    """num_classes = 前景类数 + 1（含背景）。"""
    from torchvision.models.detection import (
        FasterRCNN_ResNet50_FPN_Weights,
        fasterrcnn_resnet50_fpn,
    )
    from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

    weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT if pretrained_backbone else None
    model = fasterrcnn_resnet50_fpn(weights=weights)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model


class YOLODetector:
    backend = "yolo"

    def __init__(self, weights: str | Path, device: str | int = "0") -> None:
        YOLO = import_yolo()
        self.model = YOLO(str(weights))
        self.device = str(device)

    def detect(self, image_bgr: np.ndarray, conf: float = 0.25, iou: float = 0.7, imgsz: int = 640) -> list[DetBox]:
        results = self.model.predict(
            source=image_bgr,
            imgsz=imgsz,
            conf=conf,
            iou=iou,
            device=self.device,
            verbose=False,
            max_det=50,
        )
        if not results:
            return []
        boxes = results[0].boxes
        if boxes is None or len(boxes) == 0:
            return []
        out: list[DetBox] = []
        xyxy = boxes.xyxy.cpu().numpy()
        confs = boxes.conf.cpu().numpy()
        clss = boxes.cls.cpu().numpy() if boxes.cls is not None else np.zeros(len(boxes))
        for i in range(len(boxes)):
            x1, y1, x2, y2 = xyxy[i]
            out.append(
                DetBox(
                    xyxy=(int(round(x1)), int(round(y1)), int(round(x2)), int(round(y2))),
                    conf=float(confs[i]),
                    cls_id=int(clss[i]),
                )
            )
        out.sort(key=lambda b: b.conf, reverse=True)
        return out

    def plot(self, image_bgr: np.ndarray, conf: float = 0.25, iou: float = 0.7, imgsz: int = 640) -> np.ndarray:
        results = self.model.predict(
            source=image_bgr,
            imgsz=imgsz,
            conf=conf,
            iou=iou,
            device=self.device,
            verbose=False,
        )
        if not results:
            return image_bgr.copy()
        return results[0].plot()


class FasterRCNNDetector:
    backend = "faster_rcnn"

    def __init__(
        self,
        weights: str | Path | None = None,
        device: str | int = "0",
        num_classes: int = 2,
        score_thresh: float = 0.25,
    ) -> None:
        self.device = resolve_device(device)
        self.score_thresh = float(score_thresh)
        self.num_classes = int(num_classes)
        self.model = build_faster_rcnn(self.num_classes, pretrained_backbone=weights is None)
        if weights is not None:
            ckpt = torch.load(str(weights), map_location="cpu", weights_only=False)
            state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
            self.model.load_state_dict(state)
        self.model.to(self.device)
        self.model.eval()

    @torch.inference_mode()
    def detect(self, image_bgr: np.ndarray, conf: float | None = None, **_: Any) -> list[DetBox]:
        thr = self.score_thresh if conf is None else float(conf)
        rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        tensor = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0
        tensor = tensor.to(self.device)
        out = self.model([tensor])[0]
        boxes = out["boxes"].detach().cpu().numpy()
        scores = out["scores"].detach().cpu().numpy()
        labels = out["labels"].detach().cpu().numpy()
        dets: list[DetBox] = []
        for box, score, label in zip(boxes, scores, labels):
            if float(score) < thr:
                continue
            # Faster R-CNN: 0=背景，前景从 1 开始；对外统一成 0-based 类别
            cls_id = int(label) - 1 if int(label) > 0 else 0
            x1, y1, x2, y2 = box
            dets.append(
                DetBox(
                    xyxy=(int(round(x1)), int(round(y1)), int(round(x2)), int(round(y2))),
                    conf=float(score),
                    cls_id=cls_id,
                )
            )
        dets.sort(key=lambda b: b.conf, reverse=True)
        return dets

    def plot(self, image_bgr: np.ndarray, conf: float | None = None, **kwargs: Any) -> np.ndarray:
        dets = self.detect(image_bgr, conf=conf, **kwargs)
        vis = image_bgr.copy()
        for det in dets:
            x1, y1, x2, y2 = det.xyxy
            cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(
                vis,
                f"id{det.cls_id} {det.conf:.2f}",
                (x1, max(20, y1 - 6)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (0, 255, 0),
                2,
            )
        return vis


def guess_backend(weights: str | Path) -> str:
    path = Path(weights)
    name = path.name.lower()
    suffix = path.suffix.lower()
    if "faster" in name or "rcnn" in name or "frcnn" in name:
        return "faster_rcnn"
    if suffix == ".pt":
        # Ultralytics YOLO 常用 .pt；torchvision 训练脚本默认 .pth
        return "yolo"
    if suffix in {".pth", ".pkl"}:
        return "faster_rcnn"
    return "yolo"


def load_detector(
    backend: str,
    weights: str | Path,
    device: str | int = "0",
    num_classes: int = 2,
    conf: float = 0.25,
):
    """
    backend: yolo | faster_rcnn | auto
    Faster R-CNN 的 num_classes = 前景类数 + 1（背景）
    """
    kind = backend.strip().lower()
    if kind == "auto":
        kind = guess_backend(weights)
    if kind in {"yolo", "ultralytics"}:
        return YOLODetector(weights, device=device)
    if kind in {"faster_rcnn", "rcnn", "faster-rcnn", "frcnn"}:
        return FasterRCNNDetector(weights, device=device, num_classes=num_classes, score_thresh=conf)
    raise ValueError(f"不支持的检测后端: {backend}")


def best_box(dets: list[DetBox], conf_thr: float = 0.25) -> DetBox | None:
    for det in dets:
        if det.conf >= conf_thr:
            return det
    return None