#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os

# 必须放在 import cv2 前，强制 RTSP 使用 TCP
os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = "rtsp_transport;tcp"

import cv2
import sys
import json
import time
import signal
import numpy as np
import hbm_runtime
import rclpy

from rclpy.node import Node
from std_msgs.msg import String
from typing import Optional, Dict, Tuple
from hobot_vio import libsrcampy as srcampy


STRIDES = np.array([8, 16, 32], dtype=np.int32)

ANCHORS = np.array([
    [10, 13], [16, 30], [33, 23],
    [30, 61], [62, 45], [59, 119],
    [116, 90], [156, 198], [373, 326]
], dtype=np.float32).reshape(3, 3, 2)


def ret_ok(ret):
    if isinstance(ret, (list, tuple)):
        return len(ret) > 0 and ret[0] == 0
    return ret == 0


def get_nalu_pos(byte_stream: bytes):
    size = len(byte_stream)
    nals = []
    retnals = []

    start_code = b"\x00\x00\x01"
    pos = 0

    while pos < size:
        is4bytes = False
        retpos = byte_stream.find(start_code, pos)

        if retpos == -1:
            break

        if retpos > 0 and byte_stream[retpos - 1] == 0:
            retpos -= 1
            is4bytes = True

        pos = retpos + (4 if is4bytes else 3)

        if pos >= size:
            break

        val = int(byte_stream[pos])
        fb = (val >> 7) & 0x1
        nri = (val >> 5) & 0x3
        nalu_type = val & 0x1F

        nals.append((pos, is4bytes, fb, nri, nalu_type))

    if not nals:
        return []

    for i in range(0, len(nals) - 1):
        start = nals[i][0]
        end = nals[i + 1][0] - (5 if nals[i + 1][1] else 4)
        retnals.append((
            start,
            end,
            nals[i][1],
            nals[i][2],
            nals[i][3],
            nals[i][4],
        ))

    start = nals[-1][0]
    end = size - 1
    retnals.append((
        start,
        end,
        nals[-1][1],
        nals[-1][2],
        nals[-1][3],
        nals[-1][4],
    ))

    return retnals


def get_h264_nalu_type(byte_stream: bytes):
    nalu_types = []
    nalu_pos = get_nalu_pos(byte_stream)

    for _, _, _, _, _, nalu_type in nalu_pos:
        nalu_types.append(nalu_type)

    return nalu_types


class YoloV5X:
    def __init__(
        self,
        model_path: str,
        score_thres: float,
        nms_thres: float,
        pre_utils,
        post_utils,
    ):
        self.pre_utils = pre_utils
        self.post_utils = post_utils

        self.model = hbm_runtime.HB_HBMRuntime(model_path)

        self.model_name = self.model.model_names[0]
        self.input_names = self.model.input_names[self.model_name]
        self.output_names = self.model.output_names[self.model_name]
        self.input_shapes = self.model.input_shapes[self.model_name]
        self.output_quants = self.model.output_quants[self.model_name]

        self.input_H = self.input_shapes[self.input_names[0]][1]
        self.input_W = self.input_shapes[self.input_names[0]][2]

        self.score_thres = score_thres
        self.nms_thres = nms_thres

        self.resize_type = 1
        self.classes_num = 80

    def set_scheduling_params(
        self,
        priority: Optional[int] = None,
        bpu_cores: Optional[list] = None,
    ):
        kwargs = {}

        if priority is not None:
            kwargs["priority"] = {self.model_name: priority}

        if bpu_cores is not None:
            kwargs["bpu_cores"] = {self.model_name: bpu_cores}

        if kwargs:
            self.model.set_scheduling_params(**kwargs)

    def pre_process(
        self,
        img: np.ndarray,
        width: int,
        height: int,
    ) -> Dict[str, Dict[str, np.ndarray]]:
        y, uv = self.pre_utils.split_nv12_bytes(img, width, height)

        y_resized, uv_resized = self.pre_utils.resize_nv12_yuv(
            y,
            uv,
            self.input_W,
            self.input_H,
        )

        y_input = y_resized[..., None][None, ...]
        uv_input = uv_resized[None, ...]

        return {
            self.model_name: {
                self.input_names[0]: y_input,
                self.input_names[1]: uv_input,
            }
        }

    def forward(
        self,
        input_tensor: Dict[str, Dict[str, np.ndarray]],
    ) -> Dict[str, np.ndarray]:
        outputs = self.model.run(input_tensor)
        return outputs[self.model_name]

    def post_process(
        self,
        outputs: Dict[str, np.ndarray],
        img_w: int,
        img_h: int,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        fp32_outputs = self.post_utils.dequantize_outputs(
            outputs,
            self.output_quants,
        )

        pred = self.post_utils.decode_outputs(
            self.output_names,
            fp32_outputs,
            STRIDES,
            ANCHORS,
            self.classes_num,
        )

        xyxy_boxes, score, cls = self.post_utils.filter_predictions(
            pred,
            self.score_thres,
        )

        keep = self.post_utils.NMS(
            xyxy_boxes,
            score,
            cls,
            self.nms_thres,
        )

        xyxy = self.post_utils.scale_coords_back(
            xyxy_boxes[keep],
            img_w,
            img_h,
            self.input_W,
            self.input_H,
            self.resize_type,
        )

        return xyxy, cls[keep], score[keep]


class RtspH264Decoder:
    def __init__(self, rtsp_url: str, logger=None):
        self.rtsp_url = rtsp_url
        self.logger = logger

        self.cap = None
        self.dec = None

        self.width = 0
        self.height = 0
        self.fps = 0.0

        self.dec_chn = 0
        self.dec_type = 1

        self.find_pps_sps = 0
        self.skip_count = 0

    def log_info(self, text):
        if self.logger:
            self.logger.info(text)
        else:
            print(text)

    def log_warn(self, text):
        if self.logger:
            self.logger.warn(text)
        else:
            print("[WARN]", text)

    def log_error(self, text):
        if self.logger:
            self.logger.error(text)
        else:
            print("[ERROR]", text)

    def open(self, dec_chn=0, dec_type=1):
        self.dec_chn = dec_chn
        self.dec_type = dec_type

        self.cap = cv2.VideoCapture(self.rtsp_url, cv2.CAP_FFMPEG)
        self.cap.set(cv2.CAP_PROP_FORMAT, -1)

        if not self.cap.isOpened():
            self.log_error(f"[RTSP] 打开失败: {self.rtsp_url}")
            return False

        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)

        self.log_info(
            f"[RTSP] open success, width={self.width}, height={self.height}, fps={self.fps}"
        )

        if self.width <= 0 or self.height <= 0:
            self.log_error("[RTSP] width/height 无效")
            self.close()
            return False

        self.dec = srcampy.Decoder()
        ret = self.dec.decode(
            "",
            self.dec_chn,
            self.dec_type,
            self.width,
            self.height,
        )

        self.log_info(f"[DEC] decode ret={ret}")

        if not ret_ok(ret):
            self.log_error("[DEC] 初始化失败")
            self.close()
            return False

        self.find_pps_sps = 0
        self.skip_count = 0

        return True

    def close(self):
        try:
            if self.dec is not None:
                self.dec.close()
        except Exception as e:
            self.log_warn(f"[DEC] close 异常: {e}")

        try:
            if self.cap is not None:
                self.cap.release()
        except Exception as e:
            self.log_warn(f"[RTSP] release 异常: {e}")

        self.cap = None
        self.dec = None

    def reopen(self):
        self.log_warn("[RTSP] 准备重连")
        self.close()
        time.sleep(0.5)
        return self.open(self.dec_chn, self.dec_type)

    def read_frame(self):
        if self.cap is None or self.dec is None:
            return None

        ret, stream_frame = self.cap.read()

        if not ret or stream_frame is None:
            self.log_warn("[RTSP] read failed")
            if not self.reopen():
                time.sleep(1.0)
            return None

        stream_bytes = stream_frame.tobytes()

        try:
            nalu_types = get_h264_nalu_type(stream_bytes)
        except Exception as e:
            self.log_warn(f"[RTSP] parse nalu failed: {e}")
            return None

        if not nalu_types:
            return None

        if (nalu_types[0] in [1, 5]) and self.find_pps_sps == 0:
            return None

        self.find_pps_sps = 1

        ret = self.dec.set_img(stream_bytes, self.dec_chn)
        if ret != 0:
            self.log_warn(f"[DEC] set_img failed ret={ret}")
            return None

        if self.skip_count < 8:
            self.skip_count += 1
            return None

        frame = self.dec.get_img()
        return frame


class PersonApproachNode(Node):
    def __init__(self):
        super().__init__("person_approach_node")

        self.declare_params()
        self.load_params()

        self.load_pydev_utils()

        self.event_pub = self.create_publisher(
            String,
            self.event_topic,
            10,
        )

        self.running = True
        self.decoder = None
        self.yolo = None
        self.class_names = []

        self.stable_count = 0
        self.last_event_time = 0.0
        self.frame_count = 0
        self.fps_start_time = time.time()

        self.get_logger().info("[人物靠近] 节点初始化完成")

    def declare_params(self):
        self.declare_parameter(
            "rtsp_url",
            "rtsp://admin:hhzx2025@192.168.0.36:554/Streaming/Channels/201",
        )
        self.declare_parameter(
            "model_path",
            "/opt/hobot/model/s100/basic/yolov5x_672x672_nv12.hbm",
        )
        self.declare_parameter(
            "label_file",
            "/app/res/labels/coco_classes.names",
        )
        self.declare_parameter(
            "pydev_demo_path",
            "/app/pydev_demo",
        )

        self.declare_parameter("priority", 0)
        self.declare_parameter("bpu_cores", [0])

        self.declare_parameter("score_thres", 0.25)
        self.declare_parameter("nms_thres", 0.45)

        self.declare_parameter("person_score_thres", 0.45)
        self.declare_parameter("area_ratio_thres", 0.15)
        self.declare_parameter("center_offset_thres", 0.25)
        self.declare_parameter("stable_frames", 5)
        self.declare_parameter("cooldown_sec", 8.0)
        self.declare_parameter("loop_sleep_sec", 0.01)

        self.declare_parameter("debug", True)
        self.declare_parameter("event_topic", "/greeting/person_approach")

    def load_params(self):
        self.rtsp_url = self.get_parameter("rtsp_url").value
        self.model_path = self.get_parameter("model_path").value
        self.label_file = self.get_parameter("label_file").value
        self.pydev_demo_path = self.get_parameter("pydev_demo_path").value

        self.priority = int(self.get_parameter("priority").value)
        self.bpu_cores = list(self.get_parameter("bpu_cores").value)

        self.score_thres = float(self.get_parameter("score_thres").value)
        self.nms_thres = float(self.get_parameter("nms_thres").value)

        self.person_score_thres = float(
            self.get_parameter("person_score_thres").value
        )
        self.area_ratio_thres = float(
            self.get_parameter("area_ratio_thres").value
        )
        self.center_offset_thres = float(
            self.get_parameter("center_offset_thres").value
        )
        self.stable_frames = int(
            self.get_parameter("stable_frames").value
        )
        self.cooldown_sec = float(
            self.get_parameter("cooldown_sec").value
        )
        self.loop_sleep_sec = float(
            self.get_parameter("loop_sleep_sec").value
        )

        self.debug = bool(self.get_parameter("debug").value)
        self.event_topic = self.get_parameter("event_topic").value

    def load_pydev_utils(self):
        """
        加载 YOLO demo 的 utils。
        优先使用 person_approach_skill 包内置 utils。
        """
        try:
            from person_approach_skill.utils import preprocess_utils as pre_utils
            from person_approach_skill.utils import postprocess_utils as post_utils
            from person_approach_skill.utils import common_utils as common

            self.pre_utils = pre_utils
            self.post_utils = post_utils
            self.common = common

            self.get_logger().info("[人物靠近] 使用 person_approach_skill 内置 utils")
            return

        except Exception as e:
            self.get_logger().warn(f"[人物靠近] 内置 utils 加载失败，尝试外部 utils: {e}")

        candidate_paths = [
            self.pydev_demo_path,
            "/app/pydev_demo",
            "/app/pydev_demo/12_rtsp_yolov5x_display_sample/..",
        ]

        for path in candidate_paths:
            path = os.path.abspath(path)
            if path not in sys.path:
                sys.path.insert(0, path)

        try:
            import skill_utils.preprocess_utils as pre_utils
            import skill_utils.postprocess_utils as post_utils
            import skill_utils.common_utils as common

            self.pre_utils = pre_utils
            self.post_utils = post_utils
            self.common = common

            self.get_logger().info("[人物靠近] 使用外部 /app/pydev_demo/utils")

        except Exception as e:
            self.get_logger().error(
                f"[人物靠近] utils 加载失败: {e}，请检查 pydev_demo_path={self.pydev_demo_path}"
            )
            raise

    def init_detector(self):
        if not os.path.exists(self.model_path):
            self.get_logger().error(f"[人物靠近] 模型不存在: {self.model_path}")
            return False

        if not os.path.exists(self.label_file):
            self.get_logger().error(f"[人物靠近] label 文件不存在: {self.label_file}")
            return False

        self.class_names = self.common.load_class_names(self.label_file)

        self.yolo = YoloV5X(
            model_path=self.model_path,
            score_thres=self.score_thres,
            nms_thres=self.nms_thres,
            pre_utils=self.pre_utils,
            post_utils=self.post_utils,
        )

        self.yolo.set_scheduling_params(
            priority=self.priority,
            bpu_cores=self.bpu_cores,
        )

        self.get_logger().info(
            f"[人物靠近] 模型加载完成: {self.model_path}"
        )

        self.decoder = RtspH264Decoder(
            rtsp_url=self.rtsp_url,
            logger=self.get_logger(),
        )

        if not self.decoder.open(dec_chn=0, dec_type=1):
            self.get_logger().error("[人物靠近] RTSP/Decoder 初始化失败")
            return False

        self.get_logger().info(
            f"[人物靠近] RTSP 初始化完成: {self.decoder.width}x{self.decoder.height}"
        )

        return True

    def get_person_class_id(self):
        for idx, name in enumerate(self.class_names):
            if name.strip().lower() == "person":
                return idx
        return 0

    def detect_person_close(self, boxes, cls_ids, scores, width, height):
        person_cls_id = self.get_person_class_id()

        image_area = float(width * height)
        img_cx = width / 2.0

        best = None

        for box, cls_id, score in zip(boxes, cls_ids, scores):
            cls_id = int(cls_id)
            score = float(score)

            if cls_id != person_cls_id:
                continue

            if score < self.person_score_thres:
                continue

            x1, y1, x2, y2 = box.astype(float)

            x1 = max(0.0, min(float(width - 1), x1))
            y1 = max(0.0, min(float(height - 1), y1))
            x2 = max(0.0, min(float(width - 1), x2))
            y2 = max(0.0, min(float(height - 1), y2))

            box_w = max(0.0, x2 - x1)
            box_h = max(0.0, y2 - y1)

            if box_w <= 0 or box_h <= 0:
                continue

            area_ratio = (box_w * box_h) / image_area
            cx = (x1 + x2) / 2.0
            center_offset = abs(cx - img_cx) / float(width)

            is_close = (
                area_ratio >= self.area_ratio_thres
                and center_offset <= self.center_offset_thres
            )

            candidate = {
                "is_close": is_close,
                "score": score,
                "bbox": [
                    int(x1),
                    int(y1),
                    int(x2),
                    int(y2),
                ],
                "area_ratio": float(area_ratio),
                "center_offset": float(center_offset),
                "class_id": int(cls_id),
                "class_name": "person",
            }

            if best is None:
                best = candidate
            elif candidate["area_ratio"] > best["area_ratio"]:
                best = candidate

        return best

    def publish_person_approach(self, result):
        now = time.time()

        if now - self.last_event_time < self.cooldown_sec:
            return

        self.last_event_time = now

        event = {
            "event": "person_approach",
            "score": float(result["score"]),
            "bbox": result["bbox"],
            "area_ratio": float(result["area_ratio"]),
            "center_offset": float(result["center_offset"]),
            "timestamp": now,
            "source": "person_approach_node",
        }

        self.event_pub.publish(
            String(data=json.dumps(event, ensure_ascii=False))
        )

        self.get_logger().info(
            f"[人物靠近] 发布人员靠近事件: {event}"
        )

    def process_once(self):
        frame = self.decoder.read_frame()

        if frame is None:
            return

        width = self.decoder.width
        height = self.decoder.height

        expected_len = width * height * 3 // 2
        if len(frame) != expected_len:
            self.get_logger().warn(
                f"[人物靠近] NV12 长度不匹配 len={len(frame)}, expected={expected_len}"
            )
            return

        img_np = np.frombuffer(frame, dtype=np.uint8)

        input_tensor = self.yolo.pre_process(
            img_np,
            width,
            height,
        )

        outputs = self.yolo.forward(input_tensor)

        boxes, cls_ids, scores = self.yolo.post_process(
            outputs,
            width,
            height,
        )

        result = self.detect_person_close(
            boxes,
            cls_ids,
            scores,
            width,
            height,
        )

        if result and result["is_close"]:
            self.stable_count += 1

            if self.debug:
                self.get_logger().info(
                    "[人物靠近] person close candidate "
                    f"stable={self.stable_count}/{self.stable_frames}, "
                    f"score={result['score']:.2f}, "
                    f"area={result['area_ratio']:.3f}, "
                    f"center_offset={result['center_offset']:.3f}, "
                    f"bbox={result['bbox']}"
                )

            if self.stable_count >= self.stable_frames:
                self.publish_person_approach(result)
                self.stable_count = 0
        else:
            self.stable_count = 0

        self.frame_count += 1
        now = time.time()

        if now - self.fps_start_time > 3.0:
            fps = self.frame_count / (now - self.fps_start_time)
            self.get_logger().info(f"[人物靠近] FPS: {fps:.2f}")
            self.frame_count = 0
            self.fps_start_time = now

    def run_loop(self):
        self.get_logger().info("[人物靠近] 开始检测循环")

        if not self.init_detector():
            self.get_logger().error("[人物靠近] 初始化失败，节点退出")
            return

        while rclpy.ok() and self.running:
            rclpy.spin_once(self, timeout_sec=0.0)

            try:
                self.process_once()
            except Exception as e:
                self.get_logger().error(f"[人物靠近] process_once 异常: {e}")
                time.sleep(0.1)

            if self.loop_sleep_sec > 0:
                time.sleep(self.loop_sleep_sec)

        self.close()

    def close(self):
        self.get_logger().info("[人物靠近] 正在释放资源")

        try:
            if self.decoder is not None:
                self.decoder.close()
        except Exception as e:
            self.get_logger().warn(f"[人物靠近] decoder close 异常: {e}")

        self.get_logger().info("[人物靠近] 节点退出")


def main(args=None):
    rclpy.init(args=args)

    node = PersonApproachNode()

    def _signal_handler(sig, frame):
        node.get_logger().info("[人物靠近] 收到退出信号")
        node.running = False

    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    try:
        node.run_loop()
    except KeyboardInterrupt:
        node.get_logger().info("[人物靠近] KeyboardInterrupt")
    finally:
        node.close()
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()