"""
costmap_publisher_node.py — Publish detected objects as Nav2 costmap obstacles.

Converts SceneObjectArray detections to:
  1. sensor_msgs/PointCloud2 → consumed by Nav2's obstacle_layer
     (configure nav2_costmap_2d obstacle_layer with observation_sources pointing here)
  2. visualization_msgs/MarkerArray → RViz debug visualisation

Only publishes objects closer than max_obstacle_distance_m (default 4m) with
sufficient confidence.  Stationary objects (chairs, tables) get a larger
inflation radius than dynamic ones (persons, pets).

Configuration in nav2_costmap_2d (local_costmap):
  obstacle_layer:
    observation_sources: scene_objects
    scene_objects:
      topic: /social/scene/obstacle_cloud
      data_type: PointCloud2
      min_obstacle_height: 0.0
      max_obstacle_height: 2.0
      marking: true
      clearing: true

Subscribes:
  /social/scene/objects    SceneObjectArray

Publishes:
  /social/scene/obstacle_cloud   sensor_msgs/PointCloud2
  /social/scene/object_markers   visualization_msgs/MarkerArray
"""

import struct

import rclpy
from rclpy.node import Node
import numpy as np
from visualization_msgs.msg import Marker, MarkerArray
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header

from saltybot_scene_msgs.msg import SceneObjectArray, SceneObject
from .hazard_classifier import HAZARD_NONE, HAZARD_PET


# Hazard objects get a larger effective footprint in the costmap
_HAZARD_SCALE = 0.4    # publish extra ring of points at this radius (m)
_DEFAULT_SCALE = 0.15  # non-hazard object footprint radius


class CostmapPublisherNode(Node):

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

        self.declare_parameter('max_obstacle_distance_m', 4.0)
        self.declare_parameter('min_confidence',          0.40)
        self.declare_parameter('obstacle_frame',          'base_link')
        self.declare_parameter('publish_markers',         True)

        self._max_dist   = self.get_parameter('max_obstacle_distance_m').value
        self._min_conf   = self.get_parameter('min_confidence').value
        self._frame      = self.get_parameter('obstacle_frame').value
        self._pub_markers= self.get_parameter('publish_markers').value

        self.create_subscription(
            SceneObjectArray, '/social/scene/objects', self._on_objects, 10
        )
        self._cloud_pub   = self.create_publisher(PointCloud2,  '/social/scene/obstacle_cloud', 10)
        self._marker_pub  = self.create_publisher(MarkerArray,  '/social/scene/object_markers', 10)

        self.get_logger().info('costmap_publisher ready')

    def _on_objects(self, msg: SceneObjectArray) -> None:
        points = []
        marker_array = MarkerArray()

        for idx, obj in enumerate(msg.objects):
            if obj.confidence < self._min_conf:
                continue
            if obj.distance_m <= 0 or obj.distance_m > self._max_dist:
                continue

            p = obj.position_3d.point
            x, y, z = p.x, p.y, p.z

            # For Nav2 obstacle layer: publish a disc of points around the object
            radius = _HAZARD_SCALE if obj.hazard_type != HAZARD_NONE else _DEFAULT_SCALE
            n_ring = 8
            for k in range(n_ring):
                angle = 2 * np.pi * k / n_ring
                px = x + radius * np.cos(angle)
                py = y + radius * np.sin(angle)
                points.append((px, py, 0.1))   # z=0.1 (ground-level obstacle)
            points.append((x, y, 0.5))          # centroid at mid-height

            # ── RViz marker ─────────────────────────────────────────────────
            if self._pub_markers:
                m = Marker()
                m.header.frame_id = self._frame
                m.header.stamp    = msg.header.stamp
                m.ns              = 'scene_objects'
                m.id              = idx
                m.type            = Marker.CYLINDER
                m.action          = Marker.ADD
                m.pose.position.x = x
                m.pose.position.y = y
                m.pose.position.z = 0.5
                m.pose.orientation.w = 1.0
                m.scale.x = radius * 2
                m.scale.y = radius * 2
                m.scale.z = 1.0
                m.lifetime.sec = 1   # auto-expire after 1s

                # Colour by hazard
                if obj.hazard_type in (1, 2, 4):      # stairs / drop / glass
                    m.color.r, m.color.g, m.color.b, m.color.a = 1.0, 0.0, 0.0, 0.7
                elif obj.hazard_type == 3:             # wet floor
                    m.color.r, m.color.g, m.color.b, m.color.a = 0.0, 0.5, 1.0, 0.6
                elif obj.hazard_type == HAZARD_PET:    # pet
                    m.color.r, m.color.g, m.color.b, m.color.a = 1.0, 0.5, 0.0, 0.7
                else:
                    m.color.r, m.color.g, m.color.b, m.color.a = 0.0, 0.8, 0.0, 0.5
                marker_array.markers.append(m)

        if points:
            self._cloud_pub.publish(
                self._make_cloud(points, msg.header.stamp)
            )
        if self._pub_markers and marker_array.markers:
            self._marker_pub.publish(marker_array)

    def _make_cloud(self, points: list, stamp) -> PointCloud2:
        pts_arr = np.array(points, dtype=np.float32)
        cloud   = PointCloud2()
        cloud.header.frame_id = self._frame
        cloud.header.stamp    = stamp
        cloud.height          = 1
        cloud.width           = len(pts_arr)
        cloud.is_dense        = True
        cloud.is_bigendian    = False
        cloud.point_step      = 12   # 3 × float32
        cloud.row_step        = cloud.point_step * cloud.width
        cloud.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.data = pts_arr.tobytes()
        return cloud


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


if __name__ == '__main__':
    main()
