"""
obstacle_detect_node.py — RealSense depth obstacle detection node (Issue #611).

Pipeline per frame (30 Hz)
──────────────────────────
1. Receive D435i aligned depth image (float32, metres) + camera_info.
2. Downsample by `downsample_factor` for efficiency.
3. Back-project non-zero pixels to 3D (camera optical frame, +Z forward).
4. RANSAC ground-plane fitting on the 3D cloud.
5. Classify points: ground (inlier) vs obstacle (above ground, in height window).
6. Cluster obstacle points using 2D grid BFS on the (X, Z) plane.
7. Publish:
     /saltybot/obstacles          visualization_msgs/MarkerArray  (centroids)
     /saltybot/obstacles/cloud    sensor_msgs/PointCloud2         (non-ground pts)
     /saltybot/obstacles/alert    std_msgs/String                 (JSON alert)
8. When `depth_estop_enabled` and DANGER detected: publish zero Twist to
   `depth_estop_topic` (default /cmd_vel_input) for safety_zone integration.

Obstacle alert levels (JSON on /saltybot/obstacles/alert)
──────────────────────────────────────────────────────────
  DANGER  — closest obstacle z < danger_dist_m  (requires immediate stop)
  WARN    — closest obstacle z < warn_dist_m    (caution)
  CLEAR   — no obstacles in range

Safety zone integration
───────────────────────
When `depth_estop_enabled: true`, DANGER obstacles trigger a zero Twist
published to `depth_estop_topic` (default: /cmd_vel_input).
This feeds into the safety_zone's cmd_vel chain and triggers its e-stop,
giving independent depth-based stopping complementary to the LIDAR zones.

Subscribes
──────────
  /camera/depth/image_rect_raw    sensor_msgs/Image      30 Hz float32 (m)
  /camera/depth/camera_info       sensor_msgs/CameraInfo (once)

Publishes
─────────
  /saltybot/obstacles             visualization_msgs/MarkerArray
  /saltybot/obstacles/cloud       sensor_msgs/PointCloud2
  /saltybot/obstacles/alert       std_msgs/String   (JSON)
  [depth_estop_topic]             geometry_msgs/Twist  (zero, when DANGER)

Parameters (see config/obstacle_detect_params.yaml)
────────────────────────────────────────────────────
  downsample_factor    int    8       factor to reduce image before processing
  ransac_n_iter        int    50      RANSAC iterations for ground plane
  ransac_inlier_m      float  0.06    inlier threshold (m)
  min_obstacle_h_m     float  0.05    min height above ground to be obstacle
  max_obstacle_h_m     float  0.80    max height above ground to be obstacle
  cluster_cell_m       float  0.30    clustering grid cell size (m)
  cluster_min_pts      int    5       minimum points per cluster
  max_obstacle_dist_m  float  5.0     discard obstacles beyond this distance
  danger_dist_m        float  0.40    obstacle z < this → DANGER
  warn_dist_m          float  1.20    obstacle z < this → WARN
  depth_estop_enabled  bool   false   publish zero-vel when DANGER
  depth_estop_topic    str    /cmd_vel_input
  camera_frame         str    camera_link
  marker_lifetime_s    float  0.2     MarkerArray marker lifetime (s)
"""

from __future__ import annotations

import json
import math
import struct

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
)

import numpy as np
from cv_bridge import CvBridge

from geometry_msgs.msg import Point, Twist, Vector3
from sensor_msgs.msg import Image, CameraInfo, PointCloud2, PointField
from std_msgs.msg import Header, String, ColorRGBA
from visualization_msgs.msg import Marker, MarkerArray

from .ground_plane import fit_ground_plane, height_above_plane, obstacle_mask
from .obstacle_clusterer import ObstacleCluster, cluster_obstacles


# ── QoS ───────────────────────────────────────────────────────────────────────
_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)
_LATCHED_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)

# Alert levels
CLEAR  = 'CLEAR'
WARN   = 'WARN'
DANGER = 'DANGER'

# Marker colours per alert level
_COLOUR = {
    CLEAR:  (0.0, 1.0, 0.0, 0.8),   # green
    WARN:   (1.0, 0.8, 0.0, 0.9),   # yellow
    DANGER: (1.0, 0.1, 0.1, 1.0),   # red
}

# D435i depth scale (image is already float32 in metres when using 32FC1)
_DEPTH_ENCODING = '32FC1'


class ObstacleDetectNode(Node):
    """RealSense depth-based obstacle detection — see module docstring."""

    def __init__(self) -> None:
        super().__init__('obstacle_detect')

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter('downsample_factor',   8)
        self.declare_parameter('ransac_n_iter',       50)
        self.declare_parameter('ransac_inlier_m',     0.06)
        self.declare_parameter('min_obstacle_h_m',    0.05)
        self.declare_parameter('max_obstacle_h_m',    0.80)
        self.declare_parameter('cluster_cell_m',      0.30)
        self.declare_parameter('cluster_min_pts',     5)
        self.declare_parameter('max_obstacle_dist_m', 5.0)
        self.declare_parameter('danger_dist_m',       0.40)
        self.declare_parameter('warn_dist_m',         1.20)
        self.declare_parameter('depth_estop_enabled', False)
        self.declare_parameter('depth_estop_topic',   '/cmd_vel_input')
        self.declare_parameter('camera_frame',        'camera_link')
        self.declare_parameter('marker_lifetime_s',   0.2)

        self._ds          = self.get_parameter('downsample_factor').value
        self._n_iter      = self.get_parameter('ransac_n_iter').value
        self._inlier_m    = self.get_parameter('ransac_inlier_m').value
        self._min_h       = self.get_parameter('min_obstacle_h_m').value
        self._max_h       = self.get_parameter('max_obstacle_h_m').value
        self._cell_m      = self.get_parameter('cluster_cell_m').value
        self._min_pts     = self.get_parameter('cluster_min_pts').value
        self._max_dist    = self.get_parameter('max_obstacle_dist_m').value
        self._danger_m    = self.get_parameter('danger_dist_m').value
        self._warn_m      = self.get_parameter('warn_dist_m').value
        self._estop_en    = self.get_parameter('depth_estop_enabled').value
        self._cam_frame   = self.get_parameter('camera_frame').value
        self._marker_life = self.get_parameter('marker_lifetime_s').value

        estop_topic = self.get_parameter('depth_estop_topic').value

        # ── Internal state ────────────────────────────────────────────────────
        self._bridge   = CvBridge()
        self._fx: float | None = None
        self._fy: float | None = None
        self._cx: float | None = None
        self._cy: float | None = None

        # Ground plane stability: blend current with previous (exponential)
        self._plane_normal: np.ndarray | None = None
        self._plane_d:      float | None       = None
        self._plane_alpha   = 0.3   # blend weight for new plane estimate

        # ── Publishers ─────────────────────────────────────────────────────────
        self._marker_pub = self.create_publisher(
            MarkerArray, '/saltybot/obstacles', 10
        )
        self._cloud_pub = self.create_publisher(
            PointCloud2, '/saltybot/obstacles/cloud', 10
        )
        self._alert_pub = self.create_publisher(
            String, '/saltybot/obstacles/alert', 10
        )
        if self._estop_en:
            self._estop_pub = self.create_publisher(
                Twist, estop_topic, 10
            )
        else:
            self._estop_pub = None

        # ── Subscriptions ──────────────────────────────────────────────────────
        self.create_subscription(
            CameraInfo,
            '/camera/depth/camera_info',
            self._on_camera_info,
            _LATCHED_QOS,
        )
        self.create_subscription(
            Image,
            '/camera/depth/image_rect_raw',
            self._on_depth,
            _SENSOR_QOS,
        )

        self.get_logger().info(
            f'obstacle_detect ready — '
            f'ds={self._ds} ransac={self._n_iter}it '
            f'danger={self._danger_m}m warn={self._warn_m}m '
            f'estop={"ON → " + estop_topic if self._estop_en else "OFF"}'
        )

    # ── Camera info ───────────────────────────────────────────────────────────

    def _on_camera_info(self, msg: CameraInfo) -> None:
        if self._fx is not None:
            return
        self._fx = float(msg.k[0])
        self._fy = float(msg.k[4])
        self._cx = float(msg.k[2])
        self._cy = float(msg.k[5])
        self.get_logger().info(
            f'camera_info: fx={self._fx:.1f} fy={self._fy:.1f} '
            f'cx={self._cx:.1f} cy={self._cy:.1f}'
        )

    # ── Main depth callback ───────────────────────────────────────────────────

    def _on_depth(self, msg: Image) -> None:
        if self._fx is None:
            return  # wait for camera_info

        # ── Decode depth ──────────────────────────────────────────────────────
        try:
            depth = self._bridge.imgmsg_to_cv2(msg, desired_encoding=_DEPTH_ENCODING)
        except Exception as exc:
            self.get_logger().warning(f'depth decode error: {exc}', throttle_duration_sec=5.0)
            return

        h_img, w_img = depth.shape
        stamp  = msg.header.stamp

        # ── Downsample ────────────────────────────────────────────────────────
        ds = self._ds
        depth_ds = depth[::ds, ::ds]   # stride-based downsample
        h_ds, w_ds = depth_ds.shape

        # Back-projection intrinsics at downsampled resolution
        fx = self._fx / ds
        fy = self._fy / ds
        cx = self._cx / ds
        cy = self._cy / ds

        # ── Back-project to 3D ────────────────────────────────────────────────
        points = _backproject(depth_ds, fx, fy, cx, cy, max_dist=self._max_dist)

        if len(points) < 30:
            self._publish_empty(stamp)
            return

        # ── RANSAC ground plane ───────────────────────────────────────────────
        result = fit_ground_plane(
            points,
            n_iter          = self._n_iter,
            inlier_thresh_m = self._inlier_m,
        )

        if result is None:
            self._publish_empty(stamp)
            return

        new_normal, new_d = result

        # Blend with previous plane for temporal stability
        if self._plane_normal is None:
            self._plane_normal = new_normal
            self._plane_d      = new_d
        else:
            alpha = self._plane_alpha
            blended = (1.0 - alpha) * self._plane_normal + alpha * new_normal
            norm_len = float(np.linalg.norm(blended))
            if norm_len > 1e-8:
                self._plane_normal = blended / norm_len
                self._plane_d      = (1.0 - alpha) * self._plane_d + alpha * new_d

        plane = (self._plane_normal, self._plane_d)

        # ── Obstacle mask ─────────────────────────────────────────────────────
        obs_mask = obstacle_mask(points, plane, self._min_h, self._max_h)
        obs_pts  = points[obs_mask]

        heights_all = height_above_plane(points, plane)
        obs_heights = heights_all[obs_mask]

        # ── Cluster obstacles ─────────────────────────────────────────────────
        clusters = cluster_obstacles(
            obs_pts, obs_heights,
            cell_m  = self._cell_m,
            min_pts = self._min_pts,
            z_range = (0.15, self._max_dist),
        )

        # ── Alert level ───────────────────────────────────────────────────────
        alert_level  = CLEAR
        closest_dist = float('inf')

        for cluster in clusters:
            d = cluster.distance_m
            if d < closest_dist:
                closest_dist = d
            if d < self._danger_m:
                alert_level = DANGER
                break
            if d < self._warn_m and alert_level != DANGER:
                alert_level = WARN

        if closest_dist == float('inf'):
            closest_dist = -1.0

        # ── Publish ───────────────────────────────────────────────────────────
        self._publish_markers(clusters, stamp)
        self._publish_cloud(obs_pts, stamp)
        self._publish_alert(alert_level, clusters, closest_dist, stamp)

        # Depth e-stop: feed zero vel to safety zone chain
        if self._estop_en and self._estop_pub is not None and alert_level == DANGER:
            self._estop_pub.publish(Twist())

    # ── Publishers ────────────────────────────────────────────────────────────

    def _publish_markers(self, clusters: list, stamp) -> None:
        """Publish obstacle centroids as sphere MarkerArray."""
        ma      = MarkerArray()
        lifetime = _ros_duration(self._marker_life)

        # Delete-all marker first for clean slate
        del_marker         = Marker()
        del_marker.action  = Marker.DELETEALL
        del_marker.header.stamp    = stamp
        del_marker.header.frame_id = self._cam_frame
        ma.markers.append(del_marker)

        for idx, cluster in enumerate(clusters):
            level  = self._cluster_alert_level(cluster)
            r, g, b, a = _COLOUR[level]
            cx, cy, cz = cluster.centroid

            m               = Marker()
            m.header.stamp    = stamp
            m.header.frame_id = self._cam_frame
            m.ns              = 'obstacles'
            m.id              = idx + 1
            m.type            = Marker.SPHERE
            m.action          = Marker.ADD
            m.pose.position   = Point(x=float(cx), y=float(cy), z=float(cz))
            m.pose.orientation.w = 1.0
            rad = max(0.05, cluster.radius_m)
            m.scale           = Vector3(x=rad * 2.0, y=rad * 2.0, z=rad * 2.0)
            m.color           = ColorRGBA(r=r, g=g, b=b, a=a)
            m.lifetime        = lifetime
            ma.markers.append(m)

            # Text label above sphere
            lbl               = Marker()
            lbl.header          = m.header
            lbl.ns              = 'obstacle_labels'
            lbl.id              = idx + 1
            lbl.type            = Marker.TEXT_VIEW_FACING
            lbl.action          = Marker.ADD
            lbl.pose.position   = Point(x=float(cx), y=float(cy) - 0.15, z=float(cz))
            lbl.pose.orientation.w = 1.0
            lbl.scale.z         = 0.12
            lbl.color           = ColorRGBA(r=1.0, g=1.0, b=1.0, a=0.9)
            lbl.text            = f'{cluster.distance_m:.2f}m'
            lbl.lifetime        = lifetime
            ma.markers.append(lbl)

        self._marker_pub.publish(ma)

    def _publish_cloud(self, points: np.ndarray, stamp) -> None:
        """Publish non-ground obstacle points as PointCloud2 (xyz float32)."""
        if len(points) == 0:
            # Publish empty cloud
            cloud              = PointCloud2()
            cloud.header.stamp    = stamp
            cloud.header.frame_id = self._cam_frame
            self._cloud_pub.publish(cloud)
            return

        pts_f32 = points.astype(np.float32)

        fields = [
            PointField(name='x', offset=0,  datatype=PointField.FLOAT32, count=1),
            PointField(name='y', offset=4,  datatype=PointField.FLOAT32, count=1),
            PointField(name='z', offset=8,  datatype=PointField.FLOAT32, count=1),
        ]
        cloud                 = PointCloud2()
        cloud.header.stamp    = stamp
        cloud.header.frame_id = self._cam_frame
        cloud.height          = 1
        cloud.width           = len(pts_f32)
        cloud.fields          = fields
        cloud.is_bigendian    = False
        cloud.point_step      = 12   # 3 × float32
        cloud.row_step        = cloud.point_step * cloud.width
        cloud.data            = pts_f32.tobytes()
        cloud.is_dense        = True
        self._cloud_pub.publish(cloud)

    def _publish_alert(
        self,
        level:       str,
        clusters:    list,
        closest_m:   float,
        stamp,
    ) -> None:
        obstacles_json = [
            {
                'x':        round(float(c.centroid[0]), 3),
                'y':        round(float(c.centroid[1]), 3),
                'z':        round(float(c.centroid[2]), 3),
                'radius_m': round(c.radius_m, 3),
                'height_m': round(c.height_m, 3),
                'n_pts':    c.n_pts,
                'level':    self._cluster_alert_level(c),
            }
            for c in clusters
        ]
        alert = {
            'alert_level':      level,
            'closest_m':        round(closest_m, 3),
            'obstacle_count':   len(clusters),
            'obstacles':        obstacles_json,
        }
        msg      = String()
        msg.data = json.dumps(alert)
        self._alert_pub.publish(msg)

    def _publish_empty(self, stamp) -> None:
        """Publish empty state (no obstacles detected / ground plane failed)."""
        self._publish_markers([], stamp)
        self._publish_cloud(np.zeros((0, 3), np.float32), stamp)
        alert = {
            'alert_level':    CLEAR,
            'closest_m':      -1.0,
            'obstacle_count': 0,
            'obstacles':      [],
        }
        msg      = String()
        msg.data = json.dumps(alert)
        self._alert_pub.publish(msg)

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _cluster_alert_level(self, cluster: ObstacleCluster) -> str:
        d = cluster.distance_m
        if d < self._danger_m:
            return DANGER
        if d < self._warn_m:
            return WARN
        return CLEAR


# ── Pure-function helpers ──────────────────────────────────────────────────────

def _backproject(
    depth:    np.ndarray,
    fx: float, fy: float,
    cx: float, cy: float,
    min_dist: float = 0.15,
    max_dist: float = 6.0,
) -> np.ndarray:
    """
    Back-project a float32 depth image (metres) to a (N, 3) 3D point array.

    Camera optical frame: +X right, +Y down, +Z forward.
    Only returns pixels where depth is in [min_dist, max_dist].
    """
    h, w = depth.shape

    # Build pixel-grid meshgrid
    u_arr = np.arange(w, dtype=np.float32)
    v_arr = np.arange(h, dtype=np.float32)
    uu, vv = np.meshgrid(u_arr, v_arr)

    z = depth.astype(np.float32).ravel()
    u = uu.ravel()
    v = vv.ravel()

    valid = (z > min_dist) & (z < max_dist) & np.isfinite(z)
    z = z[valid]
    u = u[valid]
    v = v[valid]

    x = (u - cx) * z / fx
    y = (v - cy) * z / fy

    return np.stack([x, y, z], axis=1).astype(np.float64)


def _ros_duration(seconds: float):
    """Build a rclpy Duration-compatible value for Marker.lifetime."""
    from rclpy.duration import Duration
    sec  = int(seconds)
    nsec = int((seconds - sec) * 1e9)
    return Duration(seconds=sec, nanoseconds=nsec).to_msg()


def main(args=None) -> None:
    rclpy.init(args=args)
    node = ObstacleDetectNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == '__main__':
    main()
