"""
costmap_layer_node.py — Nav2 costmap integration for dynamic obstacles.

Converts predicted trajectories from /saltybot/moving_objects into a
PointCloud2 fed into Nav2's ObstacleLayer.  Each predicted future position
is added as a point, creating a "smeared" dynamic obstacle zone that
covers the full 2-3 s prediction horizon.

Nav2 ObstacleLayer config (in costmap_params.yaml):
  obstacle_layer:
    enabled: true
    observation_sources: dynamic_obs
    dynamic_obs:
      topic: /saltybot/dynamic_obs_cloud
      sensor_frame: odom
      data_type: PointCloud2
      obstacle_max_range: 12.0
      obstacle_min_range: 0.0
      raytrace_max_range: 12.0
      marking: true
      clearing: false   # let the tracker handle clearing

The node also clears old obstacle points when tracks are dropped, by
publishing a clearing cloud to /saltybot/dynamic_obs_clear.

Subscribes:
  /saltybot/moving_objects     saltybot_dynamic_obs_msgs/MovingObjectArray

Publishes:
  /saltybot/dynamic_obs_cloud  sensor_msgs/PointCloud2   marking cloud
  /saltybot/dynamic_obs_clear  sensor_msgs/PointCloud2   clearing cloud

Parameters:
  inflation_radius_m  float  0.35   (each predicted point inflated by this radius)
  ring_points         int    8      (polygon approximation of inflation circle)
  clear_on_empty      bool   true   (publish clear cloud when no objects tracked)
"""

from __future__ import annotations

import math
import struct
from typing import List

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

from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header

try:
    from saltybot_dynamic_obs_msgs.msg import MovingObjectArray
    _MSGS_AVAILABLE = True
except ImportError:
    _MSGS_AVAILABLE = False


_RELIABLE_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=10,
)


def _make_pc2(header: Header, points_xyz: List[tuple]) -> PointCloud2:
    """Pack a list of (x, y, z) into a PointCloud2 message."""
    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),
    ]
    point_step = 12   # 3 × float32
    data = bytearray(len(points_xyz) * point_step)
    for i, (x, y, z) in enumerate(points_xyz):
        struct.pack_into('<fff', data, i * point_step, x, y, z)

    pc = PointCloud2()
    pc.header      = header
    pc.height      = 1
    pc.width       = len(points_xyz)
    pc.fields      = fields
    pc.is_bigendian = False
    pc.point_step  = point_step
    pc.row_step    = point_step * len(points_xyz)
    pc.data        = bytes(data)
    pc.is_dense    = True
    return pc


class CostmapLayerNode(Node):

    def __init__(self):
        super().__init__('dynamic_obs_costmap')

        self.declare_parameter('inflation_radius_m', 0.35)
        self.declare_parameter('ring_points',         8)
        self.declare_parameter('clear_on_empty',      True)

        self._infl_r      = self.get_parameter('inflation_radius_m').value
        self._ring_n      = self.get_parameter('ring_points').value
        self._clear_empty = self.get_parameter('clear_on_empty').value

        # Pre-compute ring offsets for inflation
        self._ring_offsets = [
            (self._infl_r * math.cos(2 * math.pi * i / self._ring_n),
             self._infl_r * math.sin(2 * math.pi * i / self._ring_n))
            for i in range(self._ring_n)
        ]

        if _MSGS_AVAILABLE:
            self.create_subscription(
                MovingObjectArray,
                '/saltybot/moving_objects',
                self._on_objects,
                _RELIABLE_QOS,
            )
        else:
            self.get_logger().warning(
                '[costmap_layer] saltybot_dynamic_obs_msgs not built — '
                'will not subscribe to MovingObjectArray'
            )

        self._mark_pub  = self.create_publisher(
            PointCloud2, '/saltybot/dynamic_obs_cloud', 10
        )
        self._clear_pub = self.create_publisher(
            PointCloud2, '/saltybot/dynamic_obs_clear', 10
        )

        self.get_logger().info(
            f'dynamic_obs_costmap ready — '
            f'inflation={self._infl_r}m ring_pts={self._ring_n}'
        )

    # ── Callback ──────────────────────────────────────────────────────────────

    def _on_objects(self, msg: 'MovingObjectArray') -> None:
        hdr = msg.header
        mark_pts: List[tuple] = []

        for obj in msg.objects:
            if not obj.is_valid:
                continue

            # Current position
            self._add_inflated(obj.pose.pose.position.x,
                               obj.pose.pose.position.y, mark_pts)

            # Predicted future positions
            for pose in obj.predicted_path:
                self._add_inflated(pose.position.x, pose.position.y, mark_pts)

        if mark_pts:
            self._mark_pub.publish(_make_pc2(hdr, mark_pts))
        elif self._clear_empty:
            # Publish tiny clear cloud so Nav2 clears stale markings
            self._clear_pub.publish(_make_pc2(hdr, []))

    def _add_inflated(self, cx: float, cy: float, pts: List[tuple]) -> None:
        """Add the centre + ring of inflation points at height 0.5 m."""
        pts.append((cx, cy, 0.5))
        for ox, oy in self._ring_offsets:
            pts.append((cx + ox, cy + oy, 0.5))


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


if __name__ == '__main__':
    main()
