#!/usr/bin/env python3
"""
YOLOv8n Object Detection Node with RealSense Depth Integration
Issue #468: General object detection for spatial awareness
"""

import os
import numpy as np
import cv2
from typing import Tuple, List, Optional
from pathlib import Path
import logging

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
import message_filters
from tf2_ros import TransformListener, Buffer
from tf2_geometry_msgs import PointStamped

from sensor_msgs.msg import Image, CameraInfo
from std_msgs.msg import Header
from geometry_msgs.msg import Point, PointStamped as PointStampedMsg, Quaternion
from vision_msgs.msg import BoundingBox2D, Pose2D
from cv_bridge import CvBridge

from saltybot_object_detection_msgs.msg import DetectedObject, DetectedObjectArray
from saltybot_object_detection_msgs.srv import QueryObjects

_LOGGER = logging.getLogger(__name__)

# COCO class names (0-79)
_COCO_CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
    "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
    "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
    "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis",
    "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
    "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife",
    "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
    "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed",
    "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "microwave",
    "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
    "teddy bear", "hair drier", "toothbrush"
]

_YOLO_INPUT_SIZE = 640
_CONFIDENCE_THRESHOLD = 0.5
_NMS_IOU_THRESHOLD = 0.45

_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


class _TRTBackend:
    """TensorRT inference backend (primary for Jetson)."""

    def __init__(self, engine_path: str):
        try:
            import tensorrt as trt
            import pycuda.driver as cuda
            import pycuda.autoinit  # noqa: F401
        except ImportError as e:
            raise RuntimeError(f"TensorRT/pycuda not available: {e}")

        if not Path(engine_path).exists():
            raise FileNotFoundError(f"TensorRT engine not found: {engine_path}")

        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())

        self.context = self.engine.create_execution_context()
        self.stream = cuda.Stream()

        # Allocate input/output buffers
        self.h_inputs = {}
        self.h_outputs = {}
        self.d_inputs = {}
        self.d_outputs = {}
        self.bindings = []

        for binding_idx in range(self.engine.num_bindings):
            binding_name = self.engine.get_binding_name(binding_idx)
            binding_shape = self.engine.get_binding_shape(binding_idx)
            binding_dtype = self.engine.get_binding_dtype(binding_idx)

            # Convert dtype to numpy
            if binding_dtype == trt.float32:
                np_dtype = np.float32
            elif binding_dtype == trt.float16:
                np_dtype = np.float32
            else:
                raise ValueError(f"Unsupported dtype: {binding_dtype}")

            binding_size = int(np.prod(binding_shape))

            if self.engine.binding_is_input(binding_idx):
                h_buf = cuda.pagelocked_empty(binding_size, np_dtype)
                d_buf = cuda.mem_alloc(h_buf.nbytes)
                self.h_inputs[binding_name] = h_buf
                self.d_inputs[binding_name] = d_buf
                self.bindings.append(int(d_buf))
            else:
                h_buf = cuda.pagelocked_empty(binding_size, np_dtype)
                d_buf = cuda.mem_alloc(h_buf.nbytes)
                self.h_outputs[binding_name] = h_buf.reshape(binding_shape)
                self.d_outputs[binding_name] = d_buf
                self.bindings.append(int(d_buf))

        # Get input/output names
        self.input_names = list(self.h_inputs.keys())
        self.output_names = list(self.h_outputs.keys())

    def infer(self, input_data: np.ndarray) -> List[np.ndarray]:
        """Run inference."""
        import pycuda.driver as cuda

        # Copy input to device
        input_name = self.input_names[0]
        np.copyto(self.h_inputs[input_name], input_data.ravel())
        cuda.memcpy_htod_async(self.d_inputs[input_name], self.h_inputs[input_name], self.stream)

        # Execute
        self.context.execute_async_v2(self.bindings, self.stream.handle)

        # Copy outputs back
        outputs = []
        for output_name in self.output_names:
            cuda.memcpy_dtoh_async(self.h_outputs[output_name], self.d_outputs[output_name], self.stream)
            self.stream.synchronize()
            outputs.append(self.h_outputs[output_name].copy())

        return outputs


class _ONNXBackend:
    """ONNX Runtime inference backend (fallback)."""

    def __init__(self, onnx_path: str):
        try:
            import onnxruntime as ort
        except ImportError as e:
            raise RuntimeError(f"ONNXRuntime not available: {e}")

        if not Path(onnx_path).exists():
            raise FileNotFoundError(f"ONNX model not found: {onnx_path}")

        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self.session = ort.InferenceSession(onnx_path, providers=providers)

        # Get input/output info
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [output.name for output in self.session.get_outputs()]

    def infer(self, input_data: np.ndarray) -> List[np.ndarray]:
        """Run inference."""
        outputs = self.session.run(self.output_names, {self.input_name: input_data})
        return outputs


class _YOLODecoder:
    """Decode YOLOv8 output to detections."""

    def __init__(self, conf_thresh: float = 0.5, nms_iou_thresh: float = 0.45):
        self.conf_thresh = conf_thresh
        self.nms_iou_thresh = nms_iou_thresh

    def decode(self, output: np.ndarray, input_size: int) -> List[Tuple[int, str, float, Tuple[int, int, int, int]]]:
        """
        Decode YOLOv8 output.
        Output shape: [1, 84, 8400]
        Returns: List[(class_id, class_name, confidence, bbox_xyxy)]
        """
        # Transpose: [1, 84, 8400] -> [8400, 84]
        output = output.squeeze(0).transpose(1, 0)

        # Extract bbox and scores
        bboxes = output[:, :4]  # [8400, 4] cx, cy, w, h
        scores = output[:, 4:]  # [8400, 80] class scores

        # Get max score and class per detection
        max_scores = scores.max(axis=1)
        class_ids = scores.argmax(axis=1)

        # Filter by confidence
        mask = max_scores >= self.conf_thresh
        bboxes = bboxes[mask]
        class_ids = class_ids[mask]
        scores = max_scores[mask]

        if len(bboxes) == 0:
            return []

        # Convert cx, cy, w, h to x1, y1, x2, y2
        bboxes_xyxy = np.zeros_like(bboxes)
        bboxes_xyxy[:, 0] = bboxes[:, 0] - bboxes[:, 2] / 2
        bboxes_xyxy[:, 1] = bboxes[:, 1] - bboxes[:, 3] / 2
        bboxes_xyxy[:, 2] = bboxes[:, 0] + bboxes[:, 2] / 2
        bboxes_xyxy[:, 3] = bboxes[:, 1] + bboxes[:, 3] / 2

        # Apply NMS
        keep_indices = self._nms(bboxes_xyxy, scores, self.nms_iou_thresh)

        # Build result
        detections = []
        for idx in keep_indices:
            x1, y1, x2, y2 = bboxes_xyxy[idx]
            class_id = int(class_ids[idx])
            conf = float(scores[idx])
            class_name = _COCO_CLASSES[class_id]
            detections.append((class_id, class_name, conf, (int(x1), int(y1), int(x2), int(y2))))

        return detections

    @staticmethod
    def _nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float) -> List[int]:
        """Non-Maximum Suppression."""
        if len(boxes) == 0:
            return []

        x1 = boxes[:, 0]
        y1 = boxes[:, 1]
        x2 = boxes[:, 2]
        y2 = boxes[:, 3]

        areas = (x2 - x1) * (y2 - y1)
        order = scores.argsort()[::-1]

        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)

            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(0.0, xx2 - xx1)
            h = np.maximum(0.0, yy2 - yy1)
            inter = w * h

            iou = inter / (areas[i] + areas[order[1:]] - inter)
            order = order[np.where(iou <= iou_threshold)[0] + 1]

        return keep


class ObjectDetectionNode(Node):
    """YOLOv8n object detection with depth integration."""

    def __init__(self):
        super().__init__("object_detection")

        # Parameters
        self.declare_parameter("engine_path", "/mnt/nvme/saltybot/models/yolov8n.engine")
        self.declare_parameter("onnx_path", "/mnt/nvme/saltybot/models/yolov8n.onnx")
        self.declare_parameter("confidence_threshold", 0.5)
        self.declare_parameter("nms_iou_threshold", 0.45)
        self.declare_parameter("min_confidence_filter", 0.4)
        self.declare_parameter("enabled_classes", [0, 39, 41, 42, 43, 47, 56, 57, 61, 62, 64, 73])
        self.declare_parameter("depth_window_size", 7)
        self.declare_parameter("depth_min_range", 0.3)
        self.declare_parameter("depth_max_range", 6.0)
        self.declare_parameter("target_frame", "base_link")
        self.declare_parameter("publish_debug_image", False)

        # Load parameters
        self.engine_path = self.get_parameter("engine_path").value
        self.onnx_path = self.get_parameter("onnx_path").value
        self.confidence_threshold = self.get_parameter("confidence_threshold").value
        self.nms_iou_threshold = self.get_parameter("nms_iou_threshold").value
        self.min_confidence_filter = self.get_parameter("min_confidence_filter").value
        self.enabled_classes = self.get_parameter("enabled_classes").value
        self.depth_window_size = self.get_parameter("depth_window_size").value
        self.depth_min_range = self.get_parameter("depth_min_range").value
        self.depth_max_range = self.get_parameter("depth_max_range").value
        self.target_frame = self.get_parameter("target_frame").value
        self.publish_debug_image = self.get_parameter("publish_debug_image").value

        # Initialize backend
        self.backend = self._load_backend()
        self.decoder = _YOLODecoder(self.confidence_threshold, self.nms_iou_threshold)
        self.bridge = CvBridge()

        # TF2
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

        # Camera info
        self.camera_info: Optional[CameraInfo] = None
        self.camera_info_lock = None

        # Subscriptions
        color_sub = message_filters.Subscriber(
            self, Image, "color_image", qos_profile=_SENSOR_QOS
        )
        depth_sub = message_filters.Subscriber(
            self, Image, "depth_image", qos_profile=_SENSOR_QOS
        )
        camera_info_sub = message_filters.Subscriber(
            self, CameraInfo, "camera_info", qos_profile=_SENSOR_QOS
        )

        # Synchronize color + depth (slop = 1 frame @ 30fps)
        self.sync = message_filters.ApproximateTimeSynchronizer(
            [color_sub, depth_sub], queue_size=5, slop=0.033
        )
        self.sync.registerCallback(self._on_frame)

        # Camera info subscriber (separate, not synchronized)
        self.create_subscription(CameraInfo, "camera_info", self._on_camera_info, _SENSOR_QOS)

        # Publishers
        self.objects_pub = self.create_publisher(
            DetectedObjectArray, "/saltybot/objects", _SENSOR_QOS
        )
        if self.publish_debug_image:
            self.debug_image_pub = self.create_publisher(
                Image, "/saltybot/objects/debug_image", _SENSOR_QOS
            )
        else:
            self.debug_image_pub = None

        # Query service
        self.query_srv = self.create_service(
            QueryObjects, "/saltybot/objects/query", self._on_query
        )

        # Last detection for query service
        self.last_detections: List[DetectedObject] = []

        self.get_logger().info("ObjectDetectionNode initialized")

    def _load_backend(self):
        """Load TensorRT engine or fallback to ONNX."""
        try:
            if Path(self.engine_path).exists():
                self.get_logger().info(f"Loading TensorRT engine: {self.engine_path}")
                return _TRTBackend(self.engine_path)
            else:
                self.get_logger().warn(f"TRT engine not found: {self.engine_path}")
        except Exception as e:
            self.get_logger().error(f"TensorRT loading failed: {e}")

        # Fallback to ONNX
        self.get_logger().info(f"Loading ONNX model: {self.onnx_path}")
        return _ONNXBackend(self.onnx_path)

    def _on_camera_info(self, msg: CameraInfo):
        """Store camera intrinsics."""
        if self.camera_info is None:
            self.camera_info = msg
            self.get_logger().info(f"Camera info received: {msg.width}x{msg.height}")

    def _on_frame(self, color_msg: Image, depth_msg: Image):
        """Process synchronized color + depth frames."""
        if self.camera_info is None:
            self.get_logger().warn("Camera info not yet received, skipping frame")
            return

        # Decode images
        color_frame = self.bridge.imgmsg_to_cv2(color_msg, desired_encoding="bgr8")
        depth_frame = self.bridge.imgmsg_to_cv2(depth_msg, desired_encoding="float32")

        # Preprocess
        input_tensor = self._preprocess(color_frame)

        # Inference
        try:
            output = self.backend.infer(input_tensor)
            detections = self.decoder.decode(output[0], _YOLO_INPUT_SIZE)
        except Exception as e:
            self.get_logger().error(f"Inference error: {e}")
            return

        # Filter by enabled classes and confidence
        filtered_detections = [
            d for d in detections
            if d[0] in self.enabled_classes and d[2] >= self.min_confidence_filter
        ]

        # Project to 3D
        detected_objects = []
        for class_id, class_name, conf, (x1, y1, x2, y2) in filtered_detections:
            # Depth at bbox center
            cx = (x1 + x2) // 2
            cy = (y1 + y2) // 2
            depth_m = self._get_depth_at(depth_frame, cx, cy)

            if depth_m <= 0:
                continue  # Skip if no valid depth

            # Unproject to 3D
            pos_3d = self._pixel_to_3d(float(cx), float(cy), depth_m, self.camera_info)

            # Transform to target frame if needed
            if self.target_frame != "camera_color_optical_frame":
                try:
                    transform = self.tf_buffer.lookup_transform(
                        self.target_frame, "camera_color_optical_frame", color_msg.header.stamp
                    )
                    # Simple transform: apply rotation + translation
                    # For simplicity, use the position as-is with TF lookup
                    # In a real implementation, would use TF2 geometry helpers
                except Exception as e:
                    self.get_logger().warn(f"TF lookup failed: {e}")
                    pos_3d.header.frame_id = "camera_color_optical_frame"

            # Build DetectedObject
            obj = DetectedObject()
            obj.class_id = class_id
            obj.class_name = class_name
            obj.confidence = conf
            obj.bbox = BoundingBox2D()
            obj.bbox.center = Pose2D(x=float(cx), y=float(cy))
            obj.bbox.size_x = float(x2 - x1)
            obj.bbox.size_y = float(y2 - y1)
            obj.position_3d = pos_3d
            obj.distance_m = depth_m

            detected_objects.append(obj)

        # Publish
        self.last_detections = detected_objects
        array_msg = DetectedObjectArray()
        array_msg.header = color_msg.header
        array_msg.header.frame_id = self.target_frame
        array_msg.objects = detected_objects
        self.objects_pub.publish(array_msg)

        # Debug image
        if self.debug_image_pub is not None:
            self._publish_debug_image(color_frame, filtered_detections)

    def _preprocess(self, bgr_frame: np.ndarray) -> np.ndarray:
        """Preprocess image: letterbox, BGR->RGB, normalize, NCHW."""
        # Letterbox resize
        h, w = bgr_frame.shape[:2]
        scale = _YOLO_INPUT_SIZE / max(h, w)
        new_h, new_w = int(h * scale), int(w * scale)

        resized = cv2.resize(bgr_frame, (new_w, new_h))

        # Pad to square
        canvas = np.zeros((_YOLO_INPUT_SIZE, _YOLO_INPUT_SIZE, 3), dtype=np.uint8)
        pad_y = (_YOLO_INPUT_SIZE - new_h) // 2
        pad_x = (_YOLO_INPUT_SIZE - new_w) // 2
        canvas[pad_y : pad_y + new_h, pad_x : pad_x + new_w] = resized

        # BGR -> RGB
        rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)

        # Normalize 0-1
        tensor = rgb.astype(np.float32) / 255.0

        # HWC -> CHW -> NCHW
        tensor = tensor.transpose(2, 0, 1)
        tensor = np.ascontiguousarray(tensor[np.newaxis])

        return tensor

    def _get_depth_at(self, depth_frame: np.ndarray, u: int, v: int) -> float:
        """Get depth at pixel with median filtering."""
        h, w = depth_frame.shape
        half = self.depth_window_size // 2

        u1 = max(0, u - half)
        u2 = min(w, u + half + 1)
        v1 = max(0, v - half)
        v2 = min(h, v + half + 1)

        patch = depth_frame[v1:v2, u1:u2]
        valid = patch[(patch > self.depth_min_range) & (patch < self.depth_max_range)]

        if len(valid) == 0:
            return 0.0

        return float(np.median(valid))

    def _pixel_to_3d(self, u: float, v: float, depth_m: float, cam_info: CameraInfo) -> PointStampedMsg:
        """Unproject pixel to 3D point in camera frame."""
        K = cam_info.K
        fx, fy = K[0], K[4]
        cx, cy = K[2], K[5]

        X = (u - cx) * depth_m / fx
        Y = (v - cy) * depth_m / fy
        Z = depth_m

        point_msg = PointStampedMsg()
        point_msg.header.frame_id = "camera_color_optical_frame"
        point_msg.header.stamp = self.get_clock().now().to_msg()
        point_msg.point = Point(x=X, y=Y, z=Z)

        return point_msg

    def _publish_debug_image(self, frame: np.ndarray, detections: List):
        """Publish annotated debug image."""
        debug_frame = frame.copy()

        for class_id, class_name, conf, (x1, y1, x2, y2) in detections:
            # Draw bbox
            cv2.rectangle(debug_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # Draw label
            label = f"{class_name} {conf:.2f}"
            cv2.putText(
                debug_frame, label, (x1, y1 - 5),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2
            )

        msg = self.bridge.cv2_to_imgmsg(debug_frame, encoding="bgr8")
        self.debug_image_pub.publish(msg)

    def _on_query(self, request, response) -> QueryObjects.Response:
        """Handle query service."""
        if not self.last_detections:
            response.description = "No objects detected."
            response.success = False
            return response

        # Format description
        descriptions = []
        for obj in self.last_detections:
            if obj.distance_m > 0:
                descriptions.append(f"{obj.class_name} at {obj.distance_m:.1f}m")
            else:
                descriptions.append(obj.class_name)

        response.description = f"I see {', '.join(descriptions)}."
        response.success = True
        return response


def main(args=None):
    rclpy.init(args=args)
    node = ObjectDetectionNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
