#!/usr/bin/env python3
"""方框检测与自动抓取节点

基于 YOLO 模型检测方框，并自动调用视觉抓取节点完成抓取。

功能：
1. 启动参数控制：auto_grasp=true 时自动抓取
2. 实时检测方框
3. 检测成功后停止识别
4. 发布方框中心坐标到 /vision_grasp/grasp_target
5. 等待抓取完成
"""

import math
import threading
import time

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Point
from std_srvs.srv import Trigger

try:
    import cv2
except ImportError:
    cv2 = None

try:
    import numpy as np
except ImportError:
    np = None

try:
    from ultralytics import YOLO
except ImportError:
    YOLO = None


class MjpegFrameReader:
    """持续读取 ESP32 MJPEG 流并保持最新帧"""

    def __init__(self, url, logger, reconnect_delay=1.0):
        self.url = url
        self.logger = logger
        self.reconnect_delay = reconnect_delay
        self.lock = threading.Lock()
        self.latest_frame = None
        self.latest_stamp = None
        self.running = False
        self.thread = None
        self.capture = None
        self.last_log_time = 0.0

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

    def stop(self):
        self.running = False
        if self.thread is not None:
            self.thread.join(timeout=2.0)
        if self.capture is not None:
            self.capture.release()
            self.capture = None

    def get_latest(self):
        with self.lock:
            if self.latest_frame is None:
                return None, None
            return self.latest_frame.copy(), self.latest_stamp

    def _log_periodic(self, message, period=5.0):
        now = time.monotonic()
        if now - self.last_log_time >= period:
            self.logger.warning(message)
            self.last_log_time = now

    def _open_capture(self):
        cap = cv2.VideoCapture(self.url)
        if not cap.isOpened():
            cap.release()
            return None

        cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        self.logger.info(f"已连接到 ESP32 流: {self.url}")
        return cap

    def _run(self):
        while self.running:
            if self.capture is None:
                self.capture = self._open_capture()
                if self.capture is None:
                    self._log_periodic(f"无法打开 ESP32 流 {self.url}，重试中...")
                    time.sleep(self.reconnect_delay)
                    continue

            ok, frame = self.capture.read()
            if not ok or frame is None:
                self._log_periodic("ESP32 流读取失败，重新连接中...")
                self.capture.release()
                self.capture = None
                time.sleep(self.reconnect_delay)
                continue

            with self.lock:
                self.latest_frame = frame
                self.latest_stamp = time.monotonic()


class BoxDetectionGraspNode(Node):
    """方框检测与自动抓取节点"""

    def __init__(self):
        super().__init__('box_detection_grasp')

        # 检查依赖
        if cv2 is None:
            raise RuntimeError("需要 OpenCV。安装: pip install opencv-python")
        if np is None:
            raise RuntimeError("需要 NumPy。安装: pip install numpy")
        if YOLO is None:
            raise RuntimeError("需要 ultralytics。安装: pip install ultralytics")

        # 声明参数
        self.declare_parameter('stream_url', 'http://192.168.4.1/stream')
        self.declare_parameter('model_path', '/home/fallensigh/Dev/craic/tmp/best.pt')
        self.declare_parameter('confidence', 0.35)
        self.declare_parameter('imgsz', 768)
        self.declare_parameter('device', '')
        self.declare_parameter('detection_rate', 10.0)
        self.declare_parameter('box_size_m', 0.03)
        self.declare_parameter('horizontal_fov_deg', 66.0)
        self.declare_parameter('max_frame_age', 0.5)
        self.declare_parameter('show_debug_window', False)
        self.declare_parameter('auto_grasp', False)  # 是否自动抓取

        # 获取参数
        self.stream_url = self.get_parameter('stream_url').value
        self.model_path = self.get_parameter('model_path').value
        self.confidence = float(self.get_parameter('confidence').value)
        self.imgsz = int(self.get_parameter('imgsz').value)
        self.device = str(self.get_parameter('device').value)
        self.detection_rate = float(self.get_parameter('detection_rate').value)
        self.box_size_m = float(self.get_parameter('box_size_m').value)
        self.horizontal_fov_deg = float(self.get_parameter('horizontal_fov_deg').value)
        self.max_frame_age = float(self.get_parameter('max_frame_age').value)
        self.show_debug_window = bool(self.get_parameter('show_debug_window').value)
        self.auto_grasp = bool(self.get_parameter('auto_grasp').value)

        # 状态标志
        self.detection_active = self.auto_grasp  # 如果 auto_grasp=true，自动开始检测
        self.grasp_in_progress = False

        # 发布器：发布到视觉抓取节点
        self.grasp_target_pub = self.create_publisher(Point, '/vision_grasp/grasp_target', 10)

        # 服务：外部触发检测
        self.start_detection_srv = self.create_service(
            Trigger,
            'box_detection/start',
            self.handle_start_detection
        )
        self.stop_detection_srv = self.create_service(
            Trigger,
            'box_detection/stop',
            self.handle_stop_detection
        )

        # 启动相机流读取器
        self.reader = MjpegFrameReader(self.stream_url, self.get_logger())
        self.reader.start()

        # 加载 YOLO 模型
        self.get_logger().info(f'加载 YOLO 模型: {self.model_path}')
        self.model = YOLO(self.model_path)
        self.last_no_detection_log = 0.0

        # 创建定时器
        period = 1.0 / max(self.detection_rate, 0.1)
        self.timer = self.create_timer(period, self.process_frame)

        if self.auto_grasp:
            self.get_logger().info('自动抓取模式已启动，开始实时检测方框...')
        else:
            self.get_logger().info('手动模式，等待外部触发检测')

    def process_frame(self):
        """处理最新帧"""
        if not self.detection_active or self.grasp_in_progress:
            return

        frame, frame_stamp = self.reader.get_latest()
        if frame is None:
            return

        # 检查帧是否过期
        if frame_stamp is not None and time.monotonic() - frame_stamp > self.max_frame_age:
            return

        # YOLO 检测
        predict_kwargs = {
            'source': frame,
            'imgsz': self.imgsz,
            'conf': self.confidence,
            'verbose': False,
        }
        if self.device:
            predict_kwargs['device'] = self.device

        results = self.model.predict(**predict_kwargs)
        detection = self.select_best_box(results[0], frame)

        if detection is None:
            self.log_no_detection()
            if self.show_debug_window:
                cv2.imshow('box_detection', frame)
                cv2.waitKey(1)
            return

        # 估计 3D 坐标
        point = self.estimate_camera_point(detection, frame.shape)
        if point is None:
            return

        # 检测成功！停止检测并触发抓取
        self.get_logger().info(f'✓ 检测到方框: 相机坐标 ({point[0]*1000:.1f}, {point[1]*1000:.1f}, {point[2]*1000:.1f}) mm')
        self.detection_active = False  # 停止检测
        self.grasp_in_progress = True

        if self.show_debug_window:
            self.show_debug_frame(frame, detection, point)

        # 发布抓取目标（转换为 mm）
        msg = Point()
        msg.x = point[0] * 1000.0  # m -> mm
        msg.y = point[1] * 1000.0
        msg.z = point[2] * 1000.0
        self.grasp_target_pub.publish(msg)
        self.get_logger().info('已发布抓取目标到 /vision_grasp/grasp_target')

        # 在独立线程中等待抓取完成
        threading.Thread(target=self._wait_for_grasp_completion, daemon=True).start()

    def select_best_box(self, result, frame):
        """选择最佳检测框"""
        if result.boxes is None or len(result.boxes) == 0:
            return None

        best = None
        best_conf = -1.0

        for box in result.boxes:
            conf = float(box.conf[0].item()) if box.conf is not None else 0.0
            xyxy = box.xyxy[0].detach().cpu().numpy().astype(float)
            x1, y1, x2, y2 = xyxy
            w = max(0.0, x2 - x1)
            h = max(0.0, y2 - y1)

            if w < 2.0 or h < 2.0:
                continue

            if conf > best_conf:
                best_conf = conf
                best = {
                    'xyxy': xyxy,
                    'confidence': conf,
                }

        return best

    def estimate_camera_point(self, detection, frame_shape):
        """估计相机坐标系 3D 坐标（单位：米）"""
        frame_h, frame_w = frame_shape[:2]
        x1, y1, x2, y2 = detection['xyxy']
        bbox_w = max(1.0, x2 - x1)
        bbox_h = max(1.0, y2 - y1)
        pixel_side = (bbox_w + bbox_h) * 0.5

        # 相机内参
        cx = frame_w * 0.5
        cy = frame_h * 0.5
        fov_rad = math.radians(max(1.0, min(self.horizontal_fov_deg, 179.0)))
        fx = frame_w / (2.0 * math.tan(fov_rad * 0.5))
        fy = fx

        if fx <= 0.0 or fy <= 0.0:
            self.get_logger().error('无效的相机内参')
            return None

        # 估计深度
        focal = (fx + fy) * 0.5
        z = self.box_size_m * focal / pixel_side

        # 计算 3D 坐标
        u = (x1 + x2) * 0.5
        v = (y1 + y2) * 0.5
        x = (u - cx) * z / fx
        y = (v - cy) * z / fy

        return x, y, z

    def show_debug_frame(self, frame, detection, point):
        """显示调试窗口"""
        x1, y1, x2, y2 = detection['xyxy'].astype(int)
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        text = f"x={point[0]*1000:.1f} y={point[1]*1000:.1f} z={point[2]*1000:.1f} mm"
        cv2.putText(
            frame, text, (x1, max(20, y1 - 8)),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA
        )
        cv2.imshow('box_detection', frame)
        cv2.waitKey(1)

    def log_no_detection(self):
        """周期性记录未检测到目标"""
        now = time.monotonic()
        if now - self.last_no_detection_log >= 5.0:
            self.get_logger().info('当前帧未检测到方框')
            self.last_no_detection_log = now

    def _wait_for_grasp_completion(self):
        """等待抓取完成（独立线程）"""
        self.get_logger().info('等待抓取完成...')
        time.sleep(10.0)  # 等待抓取流程完成
        self.grasp_in_progress = False
        self.get_logger().info('抓取流程完成，节点已停止检测')

    def handle_start_detection(self, request, response):
        """处理启动检测服务"""
        if self.detection_active:
            response.success = False
            response.message = '检测已在运行中'
            return response

        if self.grasp_in_progress:
            response.success = False
            response.message = '抓取流程正在进行中，请等待完成'
            return response

        self.detection_active = True
        self.get_logger().info('✓ 启动方框检测')
        response.success = True
        response.message = '已启动检测'
        return response

    def handle_stop_detection(self, request, response):
        """处理停止检测服务"""
        if not self.detection_active:
            response.success = False
            response.message = '检测未运行'
            return response

        self.detection_active = False
        self.get_logger().info('✓ 停止方框检测')
        response.success = True
        response.message = '已停止检测'
        return response

    def destroy_node(self):
        """节点销毁时的清理"""
        self.reader.stop()
        if self.show_debug_window:
            cv2.destroyAllWindows()
        super().destroy_node()


def main(args=None):
    rclpy.init(args=args)
    node = BoxDetectionGraspNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
