"""
obstacle_size_node.py — Depth-based obstacle size estimator (Issue #348).

Fuses 2-D LIDAR clusters with the D435i depth image to estimate the full
3-D width and height (metres) of each obstacle.

Subscribes
----------
  /scan                         sensor_msgs/LaserScan    (BEST_EFFORT, ~10 Hz)
  /camera/depth/image_rect_raw  sensor_msgs/Image        (BEST_EFFORT, ~15 Hz)
  /camera/depth/camera_info     sensor_msgs/CameraInfo   (RELIABLE)

Publishes
---------
  /saltybot/obstacle_sizes      saltybot_scene_msgs/ObstacleSizeArray

Parameters
----------
  Clustering
    distance_threshold_m  float  0.20   Max gap between consecutive scan pts (m)
    min_points            int    3      Min LIDAR points per cluster
    range_min_m           float  0.05   Discard ranges below this (m)
    range_max_m           float  12.0   Discard ranges above this (m)

  Depth sampling
    depth_window_px       int    5      Half-side of median sampling window
    search_rows_px        int    120    Half-height of height search strip
    col_hw_px             int    10     Half-width  of height search strip
    z_tol_m               float  0.30   Depth tolerance for height collection

  Extrinsics (LIDAR origin in camera frame — metres)
    lidar_ex              float  0.0    Camera-frame X (right) offset
    lidar_ey              float  0.05   Camera-frame Y (down) offset
    lidar_ez              float  0.0    Camera-frame Z (forward) offset

Notes
-----
Camera intrinsics are read from /camera/depth/camera_info and override the
defaults in CameraParams once the first CameraInfo message arrives.

If the depth image is unavailable, the node publishes with depth_z=0 and
height_m=0, confidence=0 (LIDAR-only width estimate).
"""

from __future__ import annotations

from typing import Optional

import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from sensor_msgs.msg import CameraInfo, Image, LaserScan
from saltybot_scene_msgs.msg import ObstacleSize, ObstacleSizeArray

from ._lidar_clustering import scan_to_cartesian, cluster_points
from ._obstacle_size import CameraParams, estimate_cluster_size


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


class ObstacleSizeNode(Node):

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

        # ── Parameters ──────────────────────────────────────────────────────
        # Clustering
        self.declare_parameter('distance_threshold_m', 0.20)
        self.declare_parameter('min_points',           3)
        self.declare_parameter('range_min_m',          0.05)
        self.declare_parameter('range_max_m',          12.0)
        # Depth sampling
        self.declare_parameter('depth_window_px',  5)
        self.declare_parameter('search_rows_px',   120)
        self.declare_parameter('col_hw_px',        10)
        self.declare_parameter('z_tol_m',          0.30)
        # Extrinsics
        self.declare_parameter('lidar_ex', 0.0)
        self.declare_parameter('lidar_ey', 0.05)
        self.declare_parameter('lidar_ez', 0.0)

        p = self.get_parameter
        self._dist_thresh = float(p('distance_threshold_m').value)
        self._min_pts     = int(p('min_points').value)
        self._range_min   = float(p('range_min_m').value)
        self._range_max   = float(p('range_max_m').value)
        self._win_px      = int(p('depth_window_px').value)
        self._rows_px     = int(p('search_rows_px').value)
        self._col_hw      = int(p('col_hw_px').value)
        self._z_tol       = float(p('z_tol_m').value)

        # Build default CameraParams (overridden when CameraInfo arrives)
        self._cam = CameraParams(
            ex=float(p('lidar_ex').value),
            ey=float(p('lidar_ey').value),
            ez=float(p('lidar_ez').value),
        )
        self._cam_info_received = False

        # Latest depth frame
        self._depth_image: Optional[np.ndarray] = None

        # ── Subscribers ──────────────────────────────────────────────────────
        self._scan_sub = self.create_subscription(
            LaserScan, '/scan', self._on_scan, _SENSOR_QOS
        )
        self._depth_sub = self.create_subscription(
            Image, '/camera/depth/image_rect_raw',
            self._on_depth, _SENSOR_QOS
        )
        self._info_sub = self.create_subscription(
            CameraInfo, '/camera/depth/camera_info',
            self._on_camera_info, _RELIABLE_QOS
        )

        # ── Publisher ────────────────────────────────────────────────────────
        self._pub = self.create_publisher(
            ObstacleSizeArray, '/saltybot/obstacle_sizes', 10
        )

        self.get_logger().info(
            f'obstacle_size_node ready — '
            f'dist_thresh={self._dist_thresh}m, z_tol={self._z_tol}m'
        )

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

    def _on_camera_info(self, msg: CameraInfo) -> None:
        """Latch camera intrinsics on first message."""
        if self._cam_info_received:
            return
        ex = self._cam.ex
        ey = self._cam.ey
        ez = self._cam.ez
        self._cam = CameraParams(
            fx=float(msg.k[0]),
            fy=float(msg.k[4]),
            cx=float(msg.k[2]),
            cy=float(msg.k[5]),
            width=int(msg.width),
            height=int(msg.height),
            depth_scale=self._cam.depth_scale,
            ex=ex, ey=ey, ez=ez,
        )
        self._cam_info_received = True
        self.get_logger().info(
            f'Camera intrinsics loaded: '
            f'fx={self._cam.fx:.1f} fy={self._cam.fy:.1f} '
            f'cx={self._cam.cx:.1f} cy={self._cam.cy:.1f} '
            f'{self._cam.width}×{self._cam.height}'
        )

    def _on_depth(self, msg: Image) -> None:
        """Cache the latest depth frame as a numpy uint16 array."""
        enc = msg.encoding.lower()
        if enc not in ('16uc1', 'mono16'):
            self.get_logger().warn(
                f'Unexpected depth encoding {msg.encoding!r} — expected 16UC1',
                throttle_duration_sec=10.0,
            )
            return
        data = np.frombuffer(msg.data, dtype=np.uint16)
        try:
            self._depth_image = data.reshape((msg.height, msg.width))
        except ValueError as exc:
            self.get_logger().warn(f'Depth reshape failed: {exc}')

    def _on_scan(self, msg: LaserScan) -> None:
        """Cluster LIDAR scan, project clusters to depth image, publish sizes."""
        points = scan_to_cartesian(
            list(msg.ranges),
            msg.angle_min,
            msg.angle_increment,
            range_min=self._range_min,
            range_max=self._range_max,
        )
        clusters = cluster_points(
            points,
            distance_threshold_m=self._dist_thresh,
            min_points=self._min_pts,
        )

        depth = self._depth_image

        # Build output message
        out              = ObstacleSizeArray()
        out.header.stamp = msg.header.stamp
        out.header.frame_id = msg.header.frame_id or 'laser'

        for idx, cluster in enumerate(clusters):
            if depth is not None:
                est = estimate_cluster_size(
                    cluster,
                    depth,
                    self._cam,
                    depth_window = self._win_px,
                    search_rows  = self._rows_px,
                    col_hw       = self._col_hw,
                    z_tol        = self._z_tol,
                    obstacle_id  = idx,
                )
            else:
                # No depth image available — publish LIDAR-only width
                from ._obstacle_size import ObstacleSizeEstimate
                import math as _math
                cx = float(cluster.centroid[0])
                cy = float(cluster.centroid[1])
                est = ObstacleSizeEstimate(
                    obstacle_id = idx,
                    centroid_x  = cx,
                    centroid_y  = cy,
                    depth_z     = 0.0,
                    width_m     = float(cluster.width_m),
                    height_m    = 0.0,
                    pixel_u     = -1,
                    pixel_v     = -1,
                    lidar_range = float(_math.sqrt(cx * cx + cy * cy)),
                    confidence  = 0.0,
                )

            obs = ObstacleSize()
            obs.header      = out.header
            obs.obstacle_id = idx
            obs.centroid_x  = float(est.centroid_x)
            obs.centroid_y  = float(est.centroid_y)
            obs.depth_z     = float(est.depth_z)
            obs.width_m     = float(est.width_m)
            obs.height_m    = float(est.height_m)
            obs.pixel_u     = int(est.pixel_u)
            obs.pixel_v     = int(est.pixel_v)
            obs.lidar_range = float(est.lidar_range)
            obs.confidence  = float(est.confidence)
            out.obstacles.append(obs)

        self._pub.publish(out)


def main(args=None) -> None:
    rclpy.init(args=args)
    node = ObstacleSizeNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
