"""
map_broadcaster_node.py — Publish this robot's compressed occupancy grid to the fleet.

Subscribes:
  /<ns>/rtabmap/map            (nav_msgs/OccupancyGrid, TRANSIENT_LOCAL)

Publishes:
  /fleet/maps                  (saltybot_fleet_msgs/MapChunk, 1 Hz)
  /fleet/robots                (saltybot_fleet_msgs/RobotInfo, 1 Hz)

Parameters:
  robot_id          str    "saltybot_1"
  robot_namespace   str    "/saltybot_1"
  broadcast_hz      float  1.0   (map publish rate — compress-on-change, rate-limited)
  battery_topic     str    ""    (optional: std_msgs/Float32 battery percentage)

Map is only re-compressed and published when it changes (seq advances), but the
RobotInfo heartbeat is always published at broadcast_hz so peers detect presence.
"""

import time

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
)
from std_msgs.msg import Float32
from nav_msgs.msg import OccupancyGrid
from builtin_interfaces.msg import Time as RosTime
from geometry_msgs.msg import PoseStamped, Pose, Point, Quaternion

from saltybot_fleet_msgs.msg import MapChunk, RobotInfo
from .map_compressor import compress_grid


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


class MapBroadcasterNode(Node):

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('robot_id',        'saltybot_1')
        self.declare_parameter('robot_namespace', '/saltybot_1')
        self.declare_parameter('broadcast_hz',    1.0)
        self.declare_parameter('battery_topic',   '')

        self._robot_id  = self.get_parameter('robot_id').value
        self._ns        = self.get_parameter('robot_namespace').value
        hz              = self.get_parameter('broadcast_hz').value
        bat_topic       = self.get_parameter('battery_topic').value

        # ── State ────────────────────────────────────────────────────────────
        self._last_grid:       OccupancyGrid | None = None
        self._last_compressed: bytes                = b''
        self._last_uncomp_sz:  int                  = 0
        self._map_seq:         int                  = 0
        self._battery_pct:     float                = -1.0
        self._keyframe_count:  int                  = 0

        # ── Subscribers ──────────────────────────────────────────────────────
        map_topic = f'{self._ns}/rtabmap/map'.replace('//', '/')
        self.create_subscription(
            OccupancyGrid, map_topic, self._on_map, _LATCHED_QOS
        )
        if bat_topic:
            self.create_subscription(
                Float32, bat_topic, self._on_battery, _BEST_EFFORT_QOS
            )

        # ── Publishers ───────────────────────────────────────────────────────
        self._chunk_pub = self.create_publisher(MapChunk, '/fleet/maps', 10)
        self._info_pub  = self.create_publisher(RobotInfo, '/fleet/robots', 10)

        # ── Timer ────────────────────────────────────────────────────────────
        self.create_timer(1.0 / hz, self._broadcast_tick)

        self.get_logger().info(
            f'[{self._robot_id}] map_broadcaster ready — '
            f'subscribing to {map_topic}, publishing to /fleet/maps @ {hz:.1f} Hz'
        )

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

    def _on_map(self, msg: OccupancyGrid) -> None:
        self._last_grid = msg
        # Re-compress immediately; keep bytes ready for next publish tick
        compressed, sz = compress_grid(msg)
        self._last_compressed = compressed
        self._last_uncomp_sz  = sz
        self._map_seq        += 1

    def _on_battery(self, msg: Float32) -> None:
        self._battery_pct = float(msg.data)

    def _broadcast_tick(self) -> None:
        now = self.get_clock().now().to_msg()

        # ── RobotInfo heartbeat ──────────────────────────────────────────────
        info = RobotInfo()
        info.stamp          = now
        info.robot_id       = self._robot_id
        info.namespace      = self._ns
        info.status         = RobotInfo.STATUS_MAPPING
        info.battery_pct    = self._battery_pct
        info.keyframe_count = self._keyframe_count
        if self._last_grid:
            cells = self._last_grid.info.width * self._last_grid.info.height
            res   = self._last_grid.info.resolution
            info.map_coverage_m2 = float(cells * res * res)
        self._info_pub.publish(info)

        # ── MapChunk — only publish when we have data ────────────────────────
        if not self._last_compressed:
            return
        grid = self._last_grid
        chunk = MapChunk()
        chunk.stamp            = now
        chunk.robot_id         = self._robot_id
        chunk.map_frame        = grid.header.frame_id if grid else 'map'
        chunk.seq              = self._map_seq
        chunk.compressed_map   = list(self._last_compressed)
        chunk.uncompressed_size = self._last_uncomp_sz
        chunk.resolution       = grid.info.resolution if grid else 0.05
        chunk.origin_x         = grid.info.origin.position.x if grid else 0.0
        chunk.origin_y         = grid.info.origin.position.y if grid else 0.0
        chunk.width            = grid.info.width if grid else 0
        chunk.height           = grid.info.height if grid else 0
        chunk.tf_valid         = False   # coordinator sets this after alignment
        self._chunk_pub.publish(chunk)


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


if __name__ == '__main__':
    main()
