"""工件二阶段检测：YOLO 定位工件 + HoughCircles 检测左右螺纹孔中心。"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import cv2
import numpy as np
import yaml

ROOT = Path(__file__).resolve().parent
DEFAULT_CONFIG = ROOT / "hole_detect_config.yaml"


def load_config(path: str | Path | None = None) -> dict[str, Any]:
    cfg_path = Path(path) if path else DEFAULT_CONFIG
    with open(cfg_path, encoding="utf-8") as f:
        data = yaml.safe_load(f) or {}
    if not isinstance(data, dict):
        raise ValueError(f"配置格式错误: {cfg_path}")
    return data


def _fail(
    msg: str,
    draw_img: np.ndarray,
    bbox: tuple[int, int, int, int] | None = None,
    gate_block: bool = False,
) -> dict[str, Any]:
    return {
        "status": False,
        "msg": msg,
        "hole_points": [],
        "hole_radii": [],
        "bbox": list(bbox) if bbox is not None else None,
        "gate_block": gate_block,
        "draw_img": draw_img,
    }


def _draw_bbox(img: np.ndarray, bbox: tuple[int, int, int, int], color=(0, 255, 0)) -> None:
    x1, y1, x2, y2 = bbox
    cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
    cv2.putText(img, "gongjian", (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)


def _draw_center_gate(img: np.ndarray, max_dx: float, max_dy: float, ok: bool) -> None:
    h, w = img.shape[:2]
    cx, cy = w * 0.5, h * 0.5
    x1 = int(round(cx - max_dx))
    y1 = int(round(cy - max_dy))
    x2 = int(round(cx + max_dx))
    y2 = int(round(cy + max_dy))
    color = (0, 255, 0) if ok else (0, 165, 255)
    cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
    cv2.drawMarker(img, (int(cx), int(cy)), color, markerType=cv2.MARKER_CROSS, markerSize=18, thickness=1)


def _bbox_center_ok(
    bbox: tuple[int, int, int, int],
    img_w: int,
    img_h: int,
    cfg: dict[str, Any],
) -> tuple[bool, float, float, float, float]:
    """返回 (是否居中, bbox_cx, bbox_cy, max_dx, max_dy)。"""
    x1, y1, x2, y2 = bbox
    bx = 0.5 * (x1 + x2)
    by = 0.5 * (y1 + y2)
    max_dx = float(cfg.get("center_max_offset_ratio_x", 0.12)) * img_w
    max_dy = float(cfg.get("center_max_offset_ratio_y", 0.12)) * img_h
    ok = abs(bx - img_w * 0.5) <= max_dx and abs(by - img_h * 0.5) <= max_dy
    return ok, bx, by, max_dx, max_dy


def _draw_holes(
    img: np.ndarray,
    points: list[list[float]],
    radii: list[float] | None = None,
) -> None:
    """画孔轮廓圆 + 中心点（十字+实心点）+ 坐标。"""
    for i, (cx, cy) in enumerate(points):
        px, py = int(round(cx)), int(round(cy))
        r = 12.0 if radii is None or i >= len(radii) else float(radii[i])
        rr = max(4, int(round(r)))
        # 轮廓
        cv2.circle(img, (px, py), rr, (0, 0, 255), 2)
        # 中心十字 + 黄点
        cv2.drawMarker(img, (px, py), (0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=16, thickness=2)
        cv2.circle(img, (px, py), 3, (0, 255, 255), -1)
        label = f"P{i + 1}({px},{py}) r={rr}"
        cv2.putText(img, label, (px + 12, py - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 255), 2)


def _parse_best_bbox(dets, conf_thr: float) -> tuple[int, int, int, int] | None:
    """最多取 1 个工件：取置信度最高且 >= conf 的框。"""
    if not dets:
        return None
    ordered = sorted(dets, key=lambda d: d.conf, reverse=True)
    for det in ordered:
        if float(det.conf) < conf_thr:
            continue
        return tuple(int(v) for v in det.xyxy)  # type: ignore[return-value]
    return None


def _pad_roi(
    img_w: int,
    img_h: int,
    bbox: tuple[int, int, int, int],
    padding: int,
) -> tuple[int, int, int, int, int, int]:
    """返回 x1,y1,x2,y2,offset_x,offset_y（裁剪左上角在原图坐标）。"""
    x1, y1, x2, y2 = bbox
    px1 = max(0, x1 - padding)
    py1 = max(0, y1 - padding)
    px2 = min(img_w, x2 + padding)
    py2 = min(img_h, y2 + padding)
    if px2 <= px1 or py2 <= py1:
        raise ValueError("ROI 无效（宽或高为 0）")
    return px1, py1, px2, py2, px1, py1


def _render_success(
    image: np.ndarray,
    bbox: tuple[int, int, int, int],
    hole_points: list[list[float]],
    hole_radii: list[float],
    dist: float,
    gate: tuple[float, float] | None = None,
) -> np.ndarray:
    draw_img = image.copy()
    if gate is not None:
        _draw_center_gate(draw_img, gate[0], gate[1], True)
    _draw_bbox(draw_img, bbox)
    _draw_holes(draw_img, hole_points, hole_radii)
    cv2.putText(
        draw_img,
        f"OK dist={dist:.1f}px",
        (16, 36),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.8,
        (0, 255, 0),
        2,
    )
    return draw_img


def detect_holes(
    image: np.ndarray,
    detector,
    config: dict[str, Any] | None = None,
    device: str | int = 0,
) -> dict[str, Any]:
    """
    二阶段检测螺纹孔中心。

    detector: 需提供 detect(image_bgr, conf=..., iou=..., imgsz=...) -> list[DetBox]
              （YOLODetector / FasterRCNNDetector 均可）

    返回:
      status, msg, hole_points [[cx,cy],...], hole_radii [r1,r2], bbox, draw_img(BGR)
    """
    try:
        if image is None or not isinstance(image, np.ndarray) or image.size == 0:
            return _fail("输入图像无效", np.zeros((480, 640, 3), dtype=np.uint8))

        cfg = config if config is not None else load_config()
        draw_img = image.copy()
        h, w = image.shape[:2]

        conf = float(cfg.get("conf", 0.25))
        iou = float(cfg.get("iou", 0.7))
        imgsz = int(cfg.get("imgsz", 640))
        padding = int(cfg.get("bbox_padding", 25))
        ksize = int(cfg.get("gaussian_ksize", 5))
        if ksize % 2 == 0:
            ksize += 1
        sigma = float(cfg.get("gaussian_sigma", 1.5))

        # 1) 检测器定位工件 bbox（最多 1 个）— 兼容 YOLO / Faster R-CNN
        dets = detector.detect(image, conf=conf, iou=iou, imgsz=imgsz)
        bbox = _parse_best_bbox(dets, conf)
        if bbox is None:
            return _fail("未检测到工件", draw_img)

        _draw_bbox(draw_img, bbox)

        # 1.5) 居中门控：bbox 中心接近画面中心才进入二阶检测
        center_enable = bool(cfg.get("center_gate_enable", True))
        centered, bx, by, max_dx, max_dy = _bbox_center_ok(bbox, w, h, cfg)
        if bool(cfg.get("center_gate_draw", True)):
            _draw_center_gate(draw_img, max_dx, max_dy, centered if center_enable else True)
        if center_enable and not centered:
            dx = abs(bx - w * 0.5)
            dy = abs(by - h * 0.5)
            cv2.putText(
                draw_img,
                "MOVE TO CENTER",
                (16, 36),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.8,
                (0, 165, 255),
                2,
            )
            return _fail(
                f"工件未居中：offset=({dx:.0f},{dy:.0f})px，允许<=({max_dx:.0f},{max_dy:.0f})px",
                draw_img,
                bbox,
                gate_block=True,
            )

        # 2) padding ROI + 边界保护
        rx1, ry1, rx2, ry2, offset_x, offset_y = _pad_roi(w, h, bbox, padding)
        roi = image[ry1:ry2, rx1:rx2]
        if roi.size == 0:
            return _fail("ROI 裁剪失败", draw_img, bbox)

        # 3) 灰度 + 高斯模糊
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if roi.ndim == 3 else roi
        blur = cv2.GaussianBlur(gray, (ksize, ksize), sigma)

        # 4) HoughCircles
        circles = cv2.HoughCircles(
            blur,
            cv2.HOUGH_GRADIENT,
            dp=float(cfg.get("hough_dp", 1.2)),
            minDist=float(cfg.get("hough_min_dist", 40)),
            param1=float(cfg.get("hough_param1", 100)),
            param2=float(cfg.get("hough_param2", 28)),
            minRadius=int(cfg.get("hough_min_radius", 6)),
            maxRadius=int(cfg.get("hough_max_radius", 45)),
        )
        if circles is None:
            return _fail("HoughCircles 未检测到圆孔", draw_img, bbox)

        # 5) ROI 坐标 -> 原图坐标
        circles = np.asarray(circles[0], dtype=np.float64)
        mapped: list[tuple[float, float, float]] = []
        for cx, cy, r in circles:
            mapped.append((float(cx + offset_x), float(cy + offset_y), float(r)))

        # 6) 工业校验
        bx1, by1, bx2, by2 = bbox
        min_r = float(cfg.get("min_radius", 6))
        max_r = float(cfg.get("max_radius", 45))
        dist_min = float(cfg.get("hole_distance_min", 80))
        dist_max = float(cfg.get("hole_distance_max", 900))

        valid: list[tuple[float, float, float]] = []
        for cx, cy, r in mapped:
            if not (bx1 <= cx <= bx2 and by1 <= cy <= by2):
                continue
            if not (min_r <= r <= max_r):
                continue
            valid.append((cx, cy, r))

        # 多于 2 个时取最左/最右一对（对应左右螺纹孔），仍须满足间距规则
        if len(valid) > 2:
            ordered = sorted(valid, key=lambda c: c[0])
            left, right = ordered[0], ordered[-1]
            d = float(np.hypot(left[0] - right[0], left[1] - right[1]))
            if dist_min <= d <= dist_max:
                valid = [left, right]
            else:
                return _fail(
                    f"候选圆孔={len(valid)}，左右端点距离={d:.1f}px 不在 [{dist_min}, {dist_max}]",
                    draw_img,
                    bbox,
                )

        if len(valid) != 2:
            return _fail(f"有效圆孔数量={len(valid)}，要求严格等于 2", draw_img, bbox)

        ordered = sorted(valid, key=lambda c: c[0])
        (x1, y1, r1), (x2, y2, r2) = ordered
        dist = float(np.hypot(x1 - x2, y1 - y2))
        if not (dist_min <= dist <= dist_max):
            return _fail(
                f"两孔距离={dist:.1f}px，不在 [{dist_min}, {dist_max}]",
                draw_img,
                bbox,
            )

        hole_points = [[x1, y1], [x2, y2]]
        hole_radii = [r1, r2]
        gate = (max_dx, max_dy) if bool(cfg.get("center_gate_draw", True)) else None
        draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate)
        return {
            "status": True,
            "msg": f"检测成功，两孔距离 {dist:.1f}px",
            "hole_points": hole_points,
            "hole_radii": hole_radii,
            "bbox": list(bbox),
            "draw_img": draw_img,
        }
    except Exception as exc:
        fallback = image.copy() if isinstance(image, np.ndarray) and image.size else np.zeros((480, 640, 3), dtype=np.uint8)
        return _fail(f"异常: {exc}", fallback)


class HoleStabilizer:
    """时序稳帧：EMA 平滑中心/半径，短暂失败时沿用上一帧结果，抑制闪烁。"""

    def __init__(self, config: dict[str, Any] | None = None) -> None:
        cfg = config or {}
        self.cfg = cfg
        self.hold_frames = int(cfg.get("stabilize_hold_frames", 15))
        self.ema_alpha = float(cfg.get("stabilize_ema_alpha", 0.35))
        self.max_jump = float(cfg.get("stabilize_max_jump", 60))
        self.reset()

    def reset(self) -> None:
        self.points: list[list[float]] | None = None
        self.radii: list[float] | None = None
        self.bbox: tuple[int, int, int, int] | None = None
        self.miss = 0

    def update(self, image: np.ndarray, result: dict[str, Any]) -> dict[str, Any]:
        try:
            # 未居中：不做二阶、不沿用旧孔，立即清空稳帧
            if result.get("gate_block"):
                self.reset()
                out = dict(result)
                out["hole_points"] = []
                out["hole_radii"] = []
                return out

            bbox = result.get("bbox")
            if bbox is not None and len(bbox) == 4:
                self.bbox = (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))

            if result.get("status") and len(result.get("hole_points") or []) == 2:
                pts = [[float(p[0]), float(p[1])] for p in result["hole_points"]]
                pts = sorted(pts, key=lambda p: p[0])
                rad = result.get("hole_radii") or [12.0, 12.0]
                rad = [float(rad[0]), float(rad[1])] if len(rad) >= 2 else [12.0, 12.0]

                if self.points is not None:
                    jump = max(
                        float(np.hypot(pts[0][0] - self.points[0][0], pts[0][1] - self.points[0][1])),
                        float(np.hypot(pts[1][0] - self.points[1][0], pts[1][1] - self.points[1][1])),
                    )
                    if jump > self.max_jump:
                        # 跳变过大，当作本帧失败，沿用稳定值
                        self.miss += 1
                    else:
                        a = self.ema_alpha
                        self.points = [
                            [
                                (1 - a) * self.points[0][0] + a * pts[0][0],
                                (1 - a) * self.points[0][1] + a * pts[0][1],
                            ],
                            [
                                (1 - a) * self.points[1][0] + a * pts[1][0],
                                (1 - a) * self.points[1][1] + a * pts[1][1],
                            ],
                        ]
                        self.radii = [
                            (1 - a) * self.radii[0] + a * rad[0],
                            (1 - a) * self.radii[1] + a * rad[1],
                        ]
                        self.miss = 0
                else:
                    self.points = pts
                    self.radii = rad
                    self.miss = 0
            else:
                self.miss += 1
                if self.miss > self.hold_frames:
                    self.points = None
                    self.radii = None

            if self.points is None or self.radii is None:
                # 彻底丢失：只保留当前帧绘制（可能仅有 bbox）
                out = dict(result)
                out["hole_points"] = []
                out["hole_radii"] = []
                return out

            dist = float(
                np.hypot(
                    self.points[0][0] - self.points[1][0],
                    self.points[0][1] - self.points[1][1],
                )
            )
            bbox_draw = self.bbox
            if bbox_draw is None and result.get("bbox") is not None:
                b = result["bbox"]
                bbox_draw = (int(b[0]), int(b[1]), int(b[2]), int(b[3]))
            if bbox_draw is None:
                # 无框时仍画孔
                draw_img = image.copy()
                _draw_holes(draw_img, self.points, self.radii)
            else:
                gate = None
                if bool(self.cfg.get("center_gate_draw", True)):
                    ih, iw = image.shape[:2]
                    max_dx = float(self.cfg.get("center_max_offset_ratio_x", 0.12)) * iw
                    max_dy = float(self.cfg.get("center_max_offset_ratio_y", 0.12)) * ih
                    gate = (max_dx, max_dy)
                draw_img = _render_success(image, bbox_draw, self.points, self.radii, dist, gate=gate)

            held = self.miss > 0
            return {
                "status": True,
                "msg": (
                    f"检测成功，两孔距离 {dist:.1f}px"
                    + (f"（稳帧 hold={self.miss}）" if held else "")
                ),
                "hole_points": [[p[0], p[1]] for p in self.points],
                "hole_radii": [float(self.radii[0]), float(self.radii[1])],
                "bbox": list(bbox_draw) if bbox_draw is not None else result.get("bbox"),
                "draw_img": draw_img,
            }
        except Exception as exc:
            fallback = image.copy() if isinstance(image, np.ndarray) and image.size else np.zeros((480, 640, 3), dtype=np.uint8)
            return _fail(f"稳帧异常: {exc}", fallback)