"""视频流工作线程：读取视频帧、调用检测器、绘制 OSD 并缓存 JPEG。"""

from __future__ import annotations

import threading
import time
from typing import Any, Callable

import cv2


class StreamWorker:
    def __init__(
        self,
        stream_url: str | Callable[[], str],
        detector: Any,
        frame_skip: int = 3,
        jpeg_quality: int = 80,
        resize_width: int | None = None,
    ):
        self.stream_url = stream_url
        self.detector = detector
        self.frame_skip = max(1, frame_skip)
        self.jpeg_quality = jpeg_quality
        self.resize_width = resize_width

        self.lock = threading.Lock()
        self.latest_jpeg: bytes | None = None
        self.latest_detections: list[dict[str, Any]] = []
        self.frame_id = 0
        self.updated_at = 0.0
        self.running = False
        self.connected = False
        self.error = "尚未启动"
        self.fps = 0.0
        self.thread: threading.Thread | None = None
        self.reconnect_requested = False
        self.reconnect_version = 0
        self.resolve_ms = 0.0
        self.open_ms = 0.0
        self.first_frame_ms = 0.0

    def start(self) -> None:
        if self.running:
            return

        self.running = True
        self.thread = threading.Thread(target=self._run, daemon=True)
        self.thread.start()

    def stop(self) -> None:
        self.running = False
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=2)

    def reconnect(self) -> None:
        with self.lock:
            self.latest_jpeg = None
            self.latest_detections = []
            self.frame_id = 0
            self.fps = 0.0
            self.reconnect_requested = True
            self.reconnect_version += 1
            self.connected = False
            self.error = "正在切换视频源"
            self.resolve_ms = 0.0
            self.open_ms = 0.0
            self.first_frame_ms = 0.0

    def get_jpeg(self) -> bytes | None:
        with self.lock:
            return self.latest_jpeg

    def get_snapshot(self) -> dict[str, Any]:
        with self.lock:
            return {
                "frame_id": self.frame_id,
                "updated_at": self.updated_at,
                "detections": list(self.latest_detections),
                "running": self.running,
                "connected": self.connected,
                "error": self.error,
                "fps": round(self.fps, 2),
                "timings": {
                    "resolve_ms": self.resolve_ms,
                    "open_ms": self.open_ms,
                    "first_frame_ms": self.first_frame_ms,
                },
            }

    def _run(self) -> None:
        cap: cv2.VideoCapture | None = None
        last_detections: list[dict[str, Any]] = []
        fps_window_start = time.monotonic()
        fps_frames = 0

        while self.running:
            with self.lock:
                should_reconnect = self.reconnect_requested
                run_version = self.reconnect_version
                self.reconnect_requested = False
            if should_reconnect:
                # 切换摄像头时必须释放旧连接，否则 OpenCV 会继续阻塞读旧流。
                if cap is not None:
                    cap.release()
                    cap = None

            if cap is None or not cap.isOpened():
                started = time.monotonic()
                stream_url = self.stream_url() if callable(self.stream_url) else self.stream_url
                resolve_ms = round((time.monotonic() - started) * 1000, 2)
                started = time.monotonic()
                cap = cv2.VideoCapture(stream_url)
                open_ms = round((time.monotonic() - started) * 1000, 2)
                with self.lock:
                    self.open_ms = open_ms
                    self.resolve_ms = resolve_ms
                    self.first_frame_ms = 0.0
                if not cap.isOpened():
                    self._set_connection_state(False, "无法打开视频流，2 秒后重试")
                    cap.release()
                    cap = None
                    time.sleep(2)
                    continue
                self._set_connection_state(True, "已连接")

            started = time.monotonic()
            ok, frame = cap.read()
            with self.lock:
                current_version = self.reconnect_version
            # 丢弃切换期间从旧连接读到的帧，避免前端画面回跳。
            if current_version != run_version:
                continue
            if self.first_frame_ms == 0.0:
                with self.lock:
                    self.first_frame_ms = round((time.monotonic() - started) * 1000, 2)
            if not ok:
                self._set_connection_state(False, "读取视频帧失败，正在重连")
                cap.release()
                cap = None
                time.sleep(2)
                continue

            frame = self._resize(frame)
            self.frame_id += 1

            if self.frame_id % self.frame_skip == 0:
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                last_detections = self.detector.detect(frame_rgb)

            annotated = self._draw(frame, last_detections)
            ok, jpeg = cv2.imencode(
                ".jpg",
                annotated,
                [cv2.IMWRITE_JPEG_QUALITY, self.jpeg_quality],
            )
            if not ok:
                continue

            fps_frames += 1
            now = time.monotonic()
            if now - fps_window_start >= 1:
                fps = fps_frames / (now - fps_window_start)
                fps_window_start = now
                fps_frames = 0
            else:
                fps = self.fps

            with self.lock:
                current_version = self.reconnect_version
            if current_version != run_version:
                continue

            with self.lock:
                self.latest_jpeg = jpeg.tobytes()
                self.latest_detections = list(last_detections)
                self.updated_at = time.time()
                self.connected = True
                self.error = ""
                self.fps = fps

        if cap is not None:
            cap.release()
        self._set_connection_state(False, "已停止")

    def _resize(self, frame: Any) -> Any:
        if not self.resize_width:
            return frame

        height, width = frame.shape[:2]
        if width <= self.resize_width:
            return frame

        scale = self.resize_width / width
        return cv2.resize(frame, (self.resize_width, int(height * scale)))

    def _draw(self, frame: Any, detections: list[dict[str, Any]]) -> Any:
        for detection in detections:
            x1, y1, x2, y2 = detection["box"]
            label = detection["label"]
            score = detection["score"]

            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 220, 80), 2)
            text = f"{label} {score:.2f}"
            cv2.rectangle(frame, (x1, max(0, y1 - 26)), (x1 + 150, y1), (0, 220, 80), -1)
            cv2.putText(
                frame,
                text,
                (x1 + 5, max(18, y1 - 7)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.55,
                (0, 0, 0),
                2,
                cv2.LINE_AA,
            )

        return frame

    def _set_connection_state(self, connected: bool, error: str) -> None:
        with self.lock:
            self.connected = connected
            self.error = error
            self.updated_at = time.time()