#!/usr/bin/env python3
"""同时预览并检查两台 D405 和一台 D455 的彩色画面。"""

from __future__ import annotations

import argparse
from collections import deque
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
import threading
import time
from typing import Any


DEFAULT_WIDTH = 640
DEFAULT_HEIGHT = 480
DEFAULT_FPS = 30
DEFAULT_LEFT_SERIAL = "260322272273"
FPS_WINDOW_SECONDS = 2.0
WARMUP_SECONDS = 5.0
MAX_DROP_RATE = 0.01
MIN_FPS_RATIO = 0.9
PREVIEW_TILE_WIDTH = 640
WINDOW_NAME = "XR RM - Three RealSense Camera Test"
OUTPUT_DIR = Path(__file__).resolve().parents[1] / "test" / "camera_test_output"


@dataclass(frozen=True)
class DeviceInfo:
    name: str
    model: str
    serial: str
    usb_type: str


@dataclass(frozen=True)
class CameraAssignment:
    role: str
    name: str
    model: str
    serial: str
    usb_type: str


def assign_camera_roles(
    devices: list[DeviceInfo], left_serial: str | None
) -> list[CameraAssignment]:
    d405 = sorted(
        (device for device in devices if device.model == "D405"),
        key=lambda device: device.serial,
    )
    d455 = [device for device in devices if device.model == "D455"]
    if len(d405) != 2 or len(d455) != 1 or len(devices) != 3:
        raise ValueError(
            f"需要连接 2 台 D405 和 1 台 D455，当前识别到 "
            f"{len(d405)} 台 D405、{len(d455)} 台 D455、共 {len(devices)} 台 RealSense"
        )

    for device in devices:
        if not device.usb_type.startswith("3"):
            raise ValueError(
                f"{device.model} ({device.serial}) 当前为 USB {device.usb_type}，"
                "请检查扩展坞和数据线"
            )

    global_camera = CameraAssignment("GLOBAL", **d455[0].__dict__)
    if left_serial is None:
        arms = [
            CameraAssignment(f"D405-{suffix}", **device.__dict__)
            for suffix, device in zip(("A", "B"), d405)
        ]
    else:
        matches = [device for device in d405 if device.serial == left_serial]
        if not matches:
            raise ValueError(f"左臂序列号 {left_serial} 不属于当前连接的 D405")
        left = matches[0]
        right = next(device for device in d405 if device.serial != left_serial)
        arms = [
            CameraAssignment("LEFT", **left.__dict__),
            CameraAssignment("RIGHT", **right.__dict__),
        ]
    return [global_camera, *arms]


def snapshot_filenames(assignments: list[CameraAssignment]) -> dict[str, str]:
    role_names = {
        "GLOBAL": "global.png",
        "LEFT": "left.png",
        "RIGHT": "right.png",
    }
    return {
        camera.serial: role_names.get(camera.role, f"d405_{camera.serial}.png")
        for camera in assignments
    }


@dataclass
class FrameStats:
    target_fps: int
    start_time: float
    received: int = 0
    dropped: int = 0
    last_frame_number: int | None = None
    first_frame_time: float | None = None
    recent_times: deque[float] = field(default_factory=deque)

    def update(self, frame_number: int, now: float) -> None:
        if self.last_frame_number is not None and frame_number > self.last_frame_number:
            self.dropped += max(0, frame_number - self.last_frame_number - 1)
        self.last_frame_number = frame_number
        self.received += 1
        if self.first_frame_time is None:
            self.first_frame_time = now
        self.recent_times.append(now)
        cutoff = now - FPS_WINDOW_SECONDS
        while self.recent_times and self.recent_times[0] < cutoff:
            self.recent_times.popleft()

    @property
    def drop_rate(self) -> float:
        expected = self.received + self.dropped
        return self.dropped / expected if expected else 0.0

    @property
    def rolling_fps(self) -> float:
        if len(self.recent_times) < 2:
            return 0.0
        elapsed = self.recent_times[-1] - self.recent_times[0]
        return (len(self.recent_times) - 1) / elapsed if elapsed > 0 else 0.0

    def average_fps(self, now: float) -> float:
        if self.received < 2 or self.first_frame_time is None:
            return 0.0
        last_frame_time = self.recent_times[-1] if self.recent_times else now
        elapsed = last_frame_time - self.first_frame_time
        return (self.received - 1) / elapsed if elapsed > 0 else 0.0


@dataclass(frozen=True)
class CameraState:
    frame: Any | None
    actual_size: tuple[int, int]
    received: int
    dropped: int
    drop_rate: float
    rolling_fps: float
    average_fps: float
    elapsed: float
    error: str


class CameraWorker:
    def __init__(
        self,
        assignment: CameraAssignment,
        width: int,
        height: int,
        fps: int,
        rs: Any,
        np: Any,
    ) -> None:
        self.assignment = assignment
        self.width = width
        self.height = height
        self.fps = fps
        self._rs = rs
        self._np = np
        self._pipeline = rs.pipeline()
        self._stop_event = threading.Event()
        self._lock = threading.Lock()
        self._thread: threading.Thread | None = None
        self._frame: Any | None = None
        self._actual_size = (width, height)
        self._stats = FrameStats(fps, time.monotonic())
        self._error = ""

    def start(self) -> None:
        config = self._rs.config()
        config.enable_device(self.assignment.serial)
        config.enable_stream(
            self._rs.stream.color,
            self.width,
            self.height,
            self._rs.format.bgr8,
            self.fps,
        )
        profile = self._pipeline.start(config)
        video_profile = profile.get_stream(self._rs.stream.color).as_video_stream_profile()
        with self._lock:
            self._actual_size = (video_profile.width(), video_profile.height())
            self._stats = FrameStats(self.fps, time.monotonic())
        self._thread = threading.Thread(
            target=self._capture_loop,
            name=f"camera-{self.assignment.serial}",
            daemon=True,
        )
        self._thread.start()

    def _capture_loop(self) -> None:
        while not self._stop_event.is_set():
            try:
                frames = self._pipeline.wait_for_frames(timeout_ms=1000)
            except RuntimeError as exc:
                if self._stop_event.is_set():
                    return
                with self._lock:
                    self._error = str(exc)
                return

            color_frame = frames.get_color_frame()
            if not color_frame:
                continue
            frame = self._np.asanyarray(color_frame.get_data()).copy()
            now = time.monotonic()
            with self._lock:
                self._frame = frame
                self._stats.update(color_frame.get_frame_number(), now)

    def state(self, now: float) -> CameraState:
        with self._lock:
            return CameraState(
                frame=self._frame,
                actual_size=self._actual_size,
                received=self._stats.received,
                dropped=self._stats.dropped,
                drop_rate=self._stats.drop_rate,
                rolling_fps=self._stats.rolling_fps,
                average_fps=self._stats.average_fps(now),
                elapsed=now - self._stats.start_time,
                error=self._error,
            )

    def stop(self) -> None:
        self._stop_event.set()
        if self._thread is not None:
            self._thread.join(timeout=1.2)
        try:
            self._pipeline.stop()
        except RuntimeError:
            pass
        if self._thread is not None and self._thread.is_alive():
            self._thread.join(timeout=1.0)


def load_runtime_dependencies() -> tuple[Any, Any, Any]:
    try:
        import cv2
        import numpy as np
        import pyrealsense2 as rs
    except ImportError as exc:
        raise RuntimeError(
            "缺少相机测试依赖。请使用 /home/robot/miniconda3/envs/xr/bin/python "
            "运行，并确认 xr 环境已安装 pyrealsense2、numpy 和 opencv-python。"
        ) from exc
    return cv2, np, rs


def enumerate_devices(rs: Any) -> list[DeviceInfo]:
    devices = []
    for device in rs.context().query_devices():
        name = device.get_info(rs.camera_info.name)
        if "D405" in name:
            model = "D405"
        elif "D455" in name:
            model = "D455"
        else:
            model = name
        usb_type = (
            device.get_info(rs.camera_info.usb_type_descriptor)
            if device.supports(rs.camera_info.usb_type_descriptor)
            else "unknown"
        )
        devices.append(
            DeviceInfo(
                name=name,
                model=model,
                serial=device.get_info(rs.camera_info.serial_number),
                usb_type=usb_type,
            )
        )
    return devices


def camera_status(state: CameraState, target_fps: int) -> tuple[str, tuple[int, int, int]]:
    if state.error:
        return "ERROR", (0, 0, 255)
    if state.received == 0:
        return "WAITING", (0, 215, 255)
    if state.elapsed < WARMUP_SECONDS:
        return "WARMUP", (0, 215, 255)
    if state.rolling_fps >= target_fps * MIN_FPS_RATIO and state.drop_rate <= MAX_DROP_RATE:
        return "PASS", (0, 200, 0)
    return "FAIL", (0, 0, 255)


def render_tile(
    worker: CameraWorker,
    state: CameraState,
    tile_width: int,
    tile_height: int,
    cv2: Any,
    np: Any,
) -> Any:
    if state.frame is None:
        tile = np.zeros((tile_height, tile_width, 3), dtype=np.uint8)
    else:
        tile = cv2.resize(state.frame, (tile_width, tile_height))

    status, color = camera_status(state, worker.fps)
    cv2.rectangle(tile, (0, 0), (tile_width, 100), (0, 0, 0), -1)
    width, height = state.actual_size
    lines = [
        f"{worker.assignment.role}  {worker.assignment.model}  {worker.assignment.serial}",
        f"USB {worker.assignment.usb_type}  {width}x{height}@{worker.fps}",
        f"FPS {state.rolling_fps:.1f}  Frames {state.received}  "
        f"Dropped {state.dropped} ({state.drop_rate:.2%})",
        status if not state.error else f"ERROR: {state.error[:70]}",
    ]
    for index, line in enumerate(lines):
        cv2.putText(
            tile,
            line,
            (10, 22 + index * 24),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.55,
            color if index == len(lines) - 1 else (255, 255, 255),
            1,
            cv2.LINE_AA,
        )
    return tile


def compose_preview(
    workers: list[CameraWorker],
    states: dict[str, CameraState],
    capture_width: int,
    capture_height: int,
    cv2: Any,
    np: Any,
) -> Any:
    tile_width = min(capture_width, PREVIEW_TILE_WIDTH)
    tile_height = round(tile_width * capture_height / capture_width)
    tiles = {
        worker.assignment.serial: render_tile(
            worker,
            states[worker.assignment.serial],
            tile_width,
            tile_height,
            cv2,
            np,
        )
        for worker in workers
    }
    global_worker = next(worker for worker in workers if worker.assignment.role == "GLOBAL")
    arm_workers = [worker for worker in workers if worker.assignment.role != "GLOBAL"]

    top = np.zeros((tile_height, tile_width * 2, 3), dtype=np.uint8)
    offset = tile_width // 2
    top[:, offset : offset + tile_width] = tiles[global_worker.assignment.serial]
    bottom = np.hstack([tiles[worker.assignment.serial] for worker in arm_workers])
    return np.vstack((top, bottom))


def save_snapshots(
    assignments: list[CameraAssignment],
    states: dict[str, CameraState],
    cv2: Any,
) -> Path:
    missing = [camera.role for camera in assignments if states[camera.serial].frame is None]
    if missing:
        raise RuntimeError(f"以下相机尚无有效画面，不能保存快照: {', '.join(missing)}")

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
    snapshot_dir = OUTPUT_DIR / timestamp
    snapshot_dir.mkdir(parents=True, exist_ok=False)
    filenames = snapshot_filenames(assignments)
    for camera in assignments:
        path = snapshot_dir / filenames[camera.serial]
        if not cv2.imwrite(str(path), states[camera.serial].frame):
            raise RuntimeError(f"保存快照失败: {path}")
    return snapshot_dir


def print_summary(workers: list[CameraWorker]) -> bool:
    now = time.monotonic()
    print("\n相机测试汇总:")
    passed = True
    for worker in workers:
        state = worker.state(now)
        status, _color = camera_status(state, worker.fps)
        passed = passed and status == "PASS"
        print(
            f"  {worker.assignment.role:<7} {worker.assignment.serial}: "
            f"平均 {state.average_fps:.1f} FPS, 接收 {state.received}, "
            f"掉帧 {state.dropped} ({state.drop_rate:.2%}), {status}"
        )
    print("结论: " + ("三路链路满足当前阈值" if passed else "至少一路未满足当前阈值"))
    return passed


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--width", type=int, default=DEFAULT_WIDTH, help="采集宽度")
    parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT, help="采集高度")
    parser.add_argument("--fps", type=int, default=DEFAULT_FPS, help="目标帧率")
    parser.add_argument(
        "--left-serial",
        default=DEFAULT_LEFT_SERIAL,
        help=f"左臂 D405 的 RealSense 序列号（默认: {DEFAULT_LEFT_SERIAL}）",
    )
    args = parser.parse_args()
    if args.width <= 0 or args.height <= 0 or args.fps <= 0:
        parser.error("width、height 和 fps 必须为正数")
    return args


def main() -> int:
    args = parse_args()
    try:
        cv2, np, rs = load_runtime_dependencies()
        assignments = assign_camera_roles(enumerate_devices(rs), args.left_serial)
    except RuntimeError as exc:
        print(f"错误: {exc}")
        return 2
    except ValueError as exc:
        print(f"设备检查失败: {exc}")
        return 2

    print("相机分配:")
    for camera in assignments:
        print(
            f"  {camera.role:<7} {camera.model}  serial={camera.serial}  "
            f"USB={camera.usb_type}"
        )

    workers: list[CameraWorker] = []
    passed = False
    try:
        for assignment in assignments:
            worker = CameraWorker(
                assignment,
                args.width,
                args.height,
                args.fps,
                rs,
                np,
            )
            workers.append(worker)
            worker.start()

        cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
        print("按 S 保存三路快照，按 Q 或 Esc 退出。")
        while True:
            now = time.monotonic()
            states = {worker.assignment.serial: worker.state(now) for worker in workers}
            preview = compose_preview(
                workers,
                states,
                args.width,
                args.height,
                cv2,
                np,
            )
            cv2.imshow(WINDOW_NAME, preview)
            key = cv2.waitKey(1) & 0xFF
            if key in (ord("q"), ord("Q"), 27):
                break
            if key in (ord("s"), ord("S")):
                try:
                    output = save_snapshots(assignments, states, cv2)
                    print(f"快照已保存: {output}")
                except RuntimeError as exc:
                    print(f"快照失败: {exc}")
            if cv2.getWindowProperty(WINDOW_NAME, cv2.WND_PROP_VISIBLE) < 1:
                break
    except KeyboardInterrupt:
        print("\n收到 Ctrl+C，正在停止相机。")
    except Exception as exc:
        print(f"相机启动或显示失败: {exc}")
    finally:
        for worker in reversed(workers):
            worker.stop()
        try:
            cv2.destroyAllWindows()
        except Exception:
            pass
        if workers:
            passed = print_summary(workers)
    return 0 if passed else 1


if __name__ == "__main__":
    raise SystemExit(main())