"""
person_detector_node.py — Person detection + tracking for saltybot.

Pipeline:
  1. Receive synchronized color + depth frames from RealSense D435i.
  2. Run YOLOv8n (TensorRT FP16 on Orin Nano Super, ONNX fallback elsewhere).
  3. Filter detections for class 'person' (COCO class 0).
  4. Estimate 3D position from aligned depth at bounding box centre.
  5. Track target person across frames (SimplePersonTracker).
  6. Publish:
       /person/detections  — vision_msgs/Detection2DArray (all detected persons)
       /person/target      — geometry_msgs/PoseStamped   (tracked person 3D pos)
       /person/debug_image — sensor_msgs/Image           (annotated RGB, lazy)

TensorRT engine:
  - Build once with scripts/build_trt_engine.py
  - Place engine at path specified by `engine_path` param
  - Falls back to ONNX Runtime (onnxruntime[-gpu]) if engine not found

Coordinate frame:
  /person/target is published in `base_link` frame.
  If TF unavailable, falls back to `camera_color_optical_frame`.
"""

import os
import math
import time
import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.duration import Duration
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy

import message_filters
import cv2
from cv_bridge import CvBridge

from sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import PoseStamped, PointStamped
from vision_msgs.msg import (
    Detection2D,
    Detection2DArray,
    ObjectHypothesisWithPose,
    BoundingBox2D,
)
import tf2_ros
import tf2_geometry_msgs  # noqa: F401 — registers PointStamped transform support

from .tracker import SimplePersonTracker
from .detection_utils import nms, letterbox, remap_bbox, get_depth_at, pixel_to_3d

_PERSON_CLASS_ID = 0      # COCO class index for 'person'
_YOLO_INPUT_SIZE = 640


# ── Inference backends ─────────────────────────────────────────────────────────

class _TRTBackend:
    """TensorRT inference engine (Jetson Orin)."""

    def __init__(self, engine_path: str):
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit  # noqa: F401

        self._cuda = cuda
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(logger) as runtime:
            self._engine = runtime.deserialize_cuda_engine(f.read())
        self._context = self._engine.create_execution_context()

        self._inputs = []
        self._outputs = []
        self._bindings = []
        for i in range(self._engine.num_io_tensors):
            name = self._engine.get_tensor_name(i)
            dtype = trt.nptype(self._engine.get_tensor_dtype(name))
            shape = tuple(self._engine.get_tensor_shape(name))
            nbytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
            host_mem = cuda.pagelocked_empty(shape, dtype)
            device_mem = cuda.mem_alloc(nbytes)
            self._bindings.append(int(device_mem))
            if self._engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self._inputs.append({'host': host_mem, 'device': device_mem})
            else:
                self._outputs.append({'host': host_mem, 'device': device_mem})
        self._stream = cuda.Stream()

    def infer(self, input_data: np.ndarray) -> list[np.ndarray]:
        np.copyto(self._inputs[0]['host'], input_data.ravel())
        self._cuda.memcpy_htod_async(
            self._inputs[0]['device'], self._inputs[0]['host'], self._stream)
        self._context.execute_async_v2(self._bindings, self._stream.handle)
        for out in self._outputs:
            self._cuda.memcpy_dtoh_async(out['host'], out['device'], self._stream)
        self._stream.synchronize()
        return [out['host'].copy() for out in self._outputs]


class _ONNXBackend:
    """ONNX Runtime inference (CPU / CUDA EP — fallback for non-Jetson)."""

    def __init__(self, onnx_path: str):
        import onnxruntime as ort
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        self._session = ort.InferenceSession(onnx_path, providers=providers)
        self._input_name = self._session.get_inputs()[0].name

    def infer(self, input_data: np.ndarray) -> list[np.ndarray]:
        return self._session.run(None, {self._input_name: input_data})


# ── Node ──────────────────────────────────────────────────────────────────────

class PersonDetectorNode(Node):

    def __init__(self):
        super().__init__('person_detector')
        self._bridge = CvBridge()
        self._camera_info: CameraInfo | None = None
        self._backend = None

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('engine_path', '')
        self.declare_parameter('onnx_path', '')
        self.declare_parameter('confidence_threshold', 0.4)
        self.declare_parameter('nms_iou_threshold', 0.45)
        self.declare_parameter('min_depth', 0.5)
        self.declare_parameter('max_depth', 5.0)
        self.declare_parameter('track_hold_duration', 2.0)
        self.declare_parameter('track_iou_threshold', 0.25)
        self.declare_parameter('target_frame', 'base_link')
        self.declare_parameter('publish_debug_image', False)

        self._conf_thresh = self.get_parameter('confidence_threshold').value
        self._nms_thresh = self.get_parameter('nms_iou_threshold').value
        self._min_depth = self.get_parameter('min_depth').value
        self._max_depth = self.get_parameter('max_depth').value
        self._target_frame = self.get_parameter('target_frame').value
        self._pub_debug = self.get_parameter('publish_debug_image').value

        hold_dur = self.get_parameter('track_hold_duration').value
        track_iou = self.get_parameter('track_iou_threshold').value
        self._tracker = SimplePersonTracker(
            hold_duration=hold_dur,
            iou_threshold=track_iou,
            min_depth=self._min_depth,
            max_depth=self._max_depth,
        )

        # Letterbox state (set during preprocessing)
        self._scale = 1.0
        self._pad_w = 0
        self._pad_h = 0
        self._orig_w = 0
        self._orig_h = 0

        # ── TF ──────────────────────────────────────────────────────────────
        self._tf_buffer = tf2_ros.Buffer()
        self._tf_listener = tf2_ros.TransformListener(self._tf_buffer, self)

        # ── Publishers ──────────────────────────────────────────────────────
        best_effort_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )
        self._pub_detections = self.create_publisher(
            Detection2DArray, '/person/detections', best_effort_qos)
        self._pub_target = self.create_publisher(
            PoseStamped, '/person/target', best_effort_qos)
        if self._pub_debug:
            self._pub_debug_img = self.create_publisher(
                Image, '/person/debug_image', best_effort_qos)

        # ── Camera info subscriber ───────────────────────────────────────────
        self.create_subscription(
            CameraInfo,
            '/camera/color/camera_info',
            self._on_camera_info,
            QoSProfile(reliability=ReliabilityPolicy.BEST_EFFORT,
                       history=HistoryPolicy.KEEP_LAST, depth=1),
        )

        # ── Synchronized color + depth subscribers ───────────────────────────
        color_sub = message_filters.Subscriber(
            self, Image, '/camera/color/image_raw',
            qos_profile=QoSProfile(
                reliability=ReliabilityPolicy.BEST_EFFORT,
                history=HistoryPolicy.KEEP_LAST, depth=4))
        depth_sub = message_filters.Subscriber(
            self, Image, '/camera/depth/image_rect_raw',
            qos_profile=QoSProfile(
                reliability=ReliabilityPolicy.BEST_EFFORT,
                history=HistoryPolicy.KEEP_LAST, depth=4))
        self._sync = message_filters.ApproximateTimeSynchronizer(
            [color_sub, depth_sub], queue_size=4, slop=0.05)
        self._sync.registerCallback(self._on_frame)

        # ── Load model ───────────────────────────────────────────────────────
        self._load_backend()
        self.get_logger().info('PersonDetectorNode ready.')

    # ── Model loading ─────────────────────────────────────────────────────────

    def _load_backend(self):
        engine_path = self.get_parameter('engine_path').value
        onnx_path = self.get_parameter('onnx_path').value

        if engine_path and os.path.isfile(engine_path):
            try:
                self._backend = _TRTBackend(engine_path)
                self.get_logger().info(f'TensorRT backend loaded: {engine_path}')
                return
            except Exception as e:
                self.get_logger().warn(f'TRT load failed ({e}), falling back to ONNX')

        if onnx_path and os.path.isfile(onnx_path):
            try:
                self._backend = _ONNXBackend(onnx_path)
                self.get_logger().info(f'ONNX backend loaded: {onnx_path}')
                return
            except Exception as e:
                self.get_logger().error(f'ONNX load failed: {e}')

        self.get_logger().error(
            'No model found. Set engine_path or onnx_path parameter. '
            'Detection disabled — node spinning without publishing.'
        )

    # ── Callbacks ─────────────────────────────────────────────────────────────

    def _on_camera_info(self, msg: CameraInfo):
        self._camera_info = msg

    def _on_frame(self, color_msg: Image, depth_msg: Image):
        if self._backend is None or self._camera_info is None:
            return

        t0 = time.monotonic()

        # Decode images
        try:
            bgr = self._bridge.imgmsg_to_cv2(color_msg, desired_encoding='bgr8')
            depth = self._bridge.imgmsg_to_cv2(depth_msg, desired_encoding='passthrough')
        except Exception as e:
            self.get_logger().error(f'Image decode error: {e}', throttle_duration_sec=5.0)
            return

        # Depth image should be float32 metres (realsense2_camera default)
        if depth.dtype != np.float32:
            depth = depth.astype(np.float32)
            if depth.max() > 100.0:       # uint16 mm → float32 m
                depth /= 1000.0

        # Run detection
        tensor = self._preprocess(bgr)
        try:
            raw_outputs = self._backend.infer(tensor)
        except Exception as e:
            self.get_logger().error(f'Inference error: {e}', throttle_duration_sec=5.0)
            return

        detections_px = self._postprocess_yolov8(raw_outputs[0])

        # Get depth for each detection
        detections = []
        for x1, y1, x2, y2, conf in detections_px:
            cx = (x1 + x2) / 2.0
            cy = (y1 + y2) / 2.0
            d = self._get_depth_at(depth, cx, cy)
            detections.append(((x1, y1, x2, y2), d, conf))

        # Update tracker
        track = self._tracker.update(detections)

        # Publish Detection2DArray
        det_array = Detection2DArray()
        det_array.header = color_msg.header
        for (x1, y1, x2, y2), d, conf in detections:
            det = self._make_detection2d(
                color_msg.header, x1, y1, x2, y2, conf)
            det_array.detections.append(det)
        self._pub_detections.publish(det_array)

        # Publish target PoseStamped
        if track is not None:
            x1, y1, x2, y2 = track.bbox
            cx = (x1 + x2) / 2.0
            cy = (y1 + y2) / 2.0
            d = track.depth

            if d > 0:
                X, Y, Z = self._pixel_to_3d(cx, cy, d)
                track.position_3d = (X, Y, Z)

                pose = PoseStamped()
                pose.header = color_msg.header
                pose.header.frame_id = 'camera_color_optical_frame'
                pose.pose.position.x = X
                pose.pose.position.y = Y
                pose.pose.position.z = Z
                pose.pose.orientation.w = 1.0

                # Transform to target_frame
                if self._target_frame != 'camera_color_optical_frame':
                    pose = self._transform_pose(pose)

                self._pub_target.publish(pose)

        # Debug image
        if self._pub_debug and hasattr(self, '_pub_debug_img'):
            debug = self._draw_debug(bgr, detections, track)
            self._pub_debug_img.publish(
                self._bridge.cv2_to_imgmsg(debug, encoding='bgr8'))

        dt = (time.monotonic() - t0) * 1000
        self.get_logger().debug(
            f'Frame: {len(detections)} persons, track={track is not None}, {dt:.1f}ms',
            throttle_duration_sec=1.0,
        )

    # ── Preprocessing ─────────────────────────────────────────────────────────

    def _preprocess(self, bgr: np.ndarray) -> np.ndarray:
        """Letterbox resize to 640×640, normalise, HWC→CHW, add batch dim."""
        h, w = bgr.shape[:2]
        canvas, scale, pad_w, pad_h = letterbox(bgr, _YOLO_INPUT_SIZE)
        self._scale = scale
        self._pad_w = pad_w
        self._pad_h = pad_h
        self._orig_w = w
        self._orig_h = h

        rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
        tensor = rgb.astype(np.float32) / 255.0
        tensor = tensor.transpose(2, 0, 1)            # HWC → CHW
        return np.ascontiguousarray(tensor[np.newaxis])  # [1, 3, H, W]

    def _remap_bbox(self, x1, y1, x2, y2):
        """Map bbox from 640×640 space back to original image space."""
        return remap_bbox(x1, y1, x2, y2,
                          self._scale, self._pad_w, self._pad_h,
                          self._orig_w, self._orig_h)

    # ── Post-processing ───────────────────────────────────────────────────────

    def _postprocess_yolov8(self, raw: np.ndarray) -> list:
        """
        Parse YOLOv8n output tensor and return person detections.

        YOLOv8n output shape: [1, 84, 8400] or [84, 8400]
          rows 0-3:  cx, cy, w, h  (in 640×640 input space)
          rows 4-83: class scores (no objectness score in v8)

        Returns:
            list of (x1, y1, x2, y2, confidence) in original image space
        """
        pred = raw.squeeze()     # [84, 8400]
        if pred.ndim != 2 or pred.shape[0] < 5:
            return []

        person_scores = pred[4 + _PERSON_CLASS_ID, :]   # [8400]
        mask = person_scores > self._conf_thresh
        if not mask.any():
            return []

        scores = person_scores[mask]
        boxes_raw = pred[:4, mask]    # cx, cy, w, h — [4, N]

        # cx,cy,w,h → x1,y1,x2,y2
        cx, cy = boxes_raw[0], boxes_raw[1]
        w2, h2 = boxes_raw[2] / 2.0, boxes_raw[3] / 2.0
        x1, y1 = cx - w2, cy - h2
        x2, y2 = cx + w2, cy + h2

        boxes = np.stack([x1, y1, x2, y2], axis=1)  # [N, 4]
        keep = nms(boxes, scores, self._nms_thresh)

        results = []
        for i in keep:
            rx1, ry1, rx2, ry2 = self._remap_bbox(
                boxes[i, 0], boxes[i, 1], boxes[i, 2], boxes[i, 3])
            # Skip degenerate boxes
            if rx2 - rx1 < 4 or ry2 - ry1 < 4:
                continue
            results.append((rx1, ry1, rx2, ry2, float(scores[i])))
        return results

    # ── Depth & 3D ────────────────────────────────────────────────────────────

    def _get_depth_at(self, depth_img: np.ndarray, u: float, v: float,
                      window: int = 7) -> float:
        """Median depth in a window around pixel (u, v). Returns 0 if invalid."""
        return get_depth_at(depth_img, u, v, window,
                            self._min_depth, self._max_depth)

    def _pixel_to_3d(self, u: float, v: float, depth_m: float):
        """Back-project pixel (u, v) at depth_m to 3D point in camera frame."""
        return pixel_to_3d(u, v, depth_m, self._camera_info.k)

    # ── TF transform ──────────────────────────────────────────────────────────

    def _transform_pose(self, pose_in: PoseStamped) -> PoseStamped:
        try:
            return self._tf_buffer.transform(
                pose_in, self._target_frame,
                timeout=Duration(seconds=0.05))
        except Exception as e:
            self.get_logger().warn(
                f'TF {pose_in.header.frame_id}→{self._target_frame} failed: {e}',
                throttle_duration_sec=5.0)
            return pose_in  # publish in camera frame as fallback

    # ── Message builders ──────────────────────────────────────────────────────

    def _make_detection2d(self, header, x1, y1, x2, y2, conf) -> Detection2D:
        det = Detection2D()
        det.header = header

        hyp = ObjectHypothesisWithPose()
        hyp.hypothesis.class_id = 'person'
        hyp.hypothesis.score = conf
        det.results.append(hyp)

        det.bbox.center.position.x = (x1 + x2) / 2.0
        det.bbox.center.position.y = (y1 + y2) / 2.0
        det.bbox.center.theta = 0.0
        det.bbox.size_x = x2 - x1
        det.bbox.size_y = y2 - y1
        return det

    # ── Debug visualisation ───────────────────────────────────────────────────

    def _draw_debug(self, bgr, detections, track):
        vis = bgr.copy()
        for (x1, y1, x2, y2), d, conf in detections:
            cv2.rectangle(vis, (int(x1), int(y1)), (int(x2), int(y2)),
                          (0, 255, 0), 2)
            cv2.putText(vis, f'{conf:.2f} {d:.1f}m',
                        (int(x1), int(y1) - 6),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
        if track is not None:
            x1, y1, x2, y2 = track.bbox
            cv2.rectangle(vis, (int(x1), int(y1)), (int(x2), int(y2)),
                          (0, 0, 255), 3)
            label = f'ID:{track.track_id} {track.depth:.1f}m'
            if track.age > 0.05:
                label += f' (held {track.age:.1f}s)'
            cv2.putText(vis, label, (int(x1), int(y1) - 6),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
        return vis


# ── Entry point ───────────────────────────────────────────────────────────────

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


if __name__ == '__main__':
    main()
