"""
fleet_manager_node.py — Central coordinator for multi-robot fleet SLAM.

Run ONE instance on the "lead" robot (or a dedicated compute node).

Responsibilities:
  1. Robot registry — track active peers, detect departures via heartbeat timeout.
  2. Map merging — decompress MapChunks from all robots → merged OccupancyGrid.
  3. Transform alignment — detect shared landmarks and compute map_to_fleet_tf
     for each robot to align their local map frames.
  4. Frontier aggregation — collect FrontierArrays from all robots; de-duplicate;
     publish consolidated /fleet/exploration_frontiers.
  5. Frontier task allocation — service: RequestFrontier assigns highest-utility
     unassigned frontier nearest to the requesting robot; re-assigns stale claims.

Subscribes:
  /fleet/maps                       (MapChunk)        — robots' compressed grids
  /fleet/robots                     (RobotInfo)       — heartbeat / status
  /fleet/exploration_frontiers_raw  (FrontierArray)   — per-robot frontier reports
  /<robot_ns>/rtabmap/mapData       (rtabmap_msgs/MapData) — for landmark extraction

Publishes:
  /fleet/merged_map                 (OccupancyGrid, 2 Hz, TRANSIENT_LOCAL)
  /fleet/exploration_frontiers      (FrontierArray, 2 Hz)
  /fleet/status                     (std_msgs/String, 1 Hz JSON)

Services:
  /fleet/register_robot             (RegisterRobot)
  /fleet/request_frontier           (RequestFrontier)

Parameters:
  heartbeat_timeout_sec   float  10.0
  merge_hz                float  2.0
  frontier_timeout_sec    float  60.0   (re-assign frontier after this if robot silent)
  map_resolution          float  0.05
"""

import json
import math
import time
from typing import Dict, List, Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
)
from nav_msgs.msg import OccupancyGrid
from std_msgs.msg import String

from saltybot_fleet_msgs.msg import (
    MapChunk, RobotInfo, Frontier, FrontierArray
)
from saltybot_fleet_msgs.srv import RegisterRobot, RequestFrontier

from .map_compressor import decompress_grid
from .map_merger import MapMerger


_TRANSIENT_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)


class _RobotRecord:
    __slots__ = ('info', 'last_seen', 'chunk_seq')

    def __init__(self, info: RobotInfo, ts: float):
        self.info       = info
        self.last_seen  = ts
        self.chunk_seq  = -1


class FleetManagerNode(Node):

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('heartbeat_timeout_sec', 10.0)
        self.declare_parameter('merge_hz',              2.0)
        self.declare_parameter('frontier_timeout_sec',  60.0)
        self.declare_parameter('map_resolution',        0.05)

        self._hb_timeout     = self.get_parameter('heartbeat_timeout_sec').value
        self._merge_hz       = self.get_parameter('merge_hz').value
        self._frontier_to    = self.get_parameter('frontier_timeout_sec').value
        self._resolution     = self.get_parameter('map_resolution').value

        # ── State ────────────────────────────────────────────────────────────
        self._robots:    Dict[str, _RobotRecord] = {}
        self._frontiers: Dict[str, Frontier]     = {}   # frontier_id → Frontier
        self._merger = MapMerger(resolution=self._resolution)
        self._merged_grid: Optional[OccupancyGrid] = None

        # ── Subscriptions ────────────────────────────────────────────────────
        self.create_subscription(MapChunk,    '/fleet/maps',                   self._on_map_chunk,   10)
        self.create_subscription(RobotInfo,   '/fleet/robots',                 self._on_robot_info,  10)
        self.create_subscription(FrontierArray, '/fleet/exploration_frontiers_raw',
                                 self._on_frontiers_raw, 10)

        # ── Publishers ───────────────────────────────────────────────────────
        self._merged_pub   = self.create_publisher(OccupancyGrid, '/fleet/merged_map', _TRANSIENT_QOS)
        self._frontier_pub = self.create_publisher(FrontierArray,  '/fleet/exploration_frontiers', 10)
        self._status_pub   = self.create_publisher(String,          '/fleet/status',  10)

        # ── Services ─────────────────────────────────────────────────────────
        self.create_service(RegisterRobot,   '/fleet/register_robot',   self._svc_register)
        self.create_service(RequestFrontier, '/fleet/request_frontier', self._svc_request_frontier)

        # ── Timers ────────────────────────────────────────────────────────────
        self.create_timer(1.0 / self._merge_hz, self._merge_tick)
        self.create_timer(1.0,                  self._status_tick)

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

    # ── Subscription callbacks ────────────────────────────────────────────────

    def _on_map_chunk(self, msg: MapChunk) -> None:
        rid = msg.robot_id
        rec = self._robots.get(rid)
        if rec is None:
            # Auto-register on first chunk (robot may not have called service)
            rec = _RobotRecord(RobotInfo(), time.monotonic())
            rec.info.robot_id  = rid
            rec.info.namespace = f'/{rid}'
            self._robots[rid]  = rec
            self.get_logger().info(f'[fleet] auto-registered robot {rid} from map chunk')

        if msg.seq <= rec.chunk_seq:
            return   # duplicate or stale
        rec.chunk_seq = msg.seq

        try:
            grid = decompress_grid(bytes(msg.compressed_map), frame_id=msg.map_frame)
        except Exception as exc:
            self.get_logger().warning(f'[fleet] decompress failed for {rid}: {exc}')
            return

        tf = msg.map_to_fleet_tf if msg.tf_valid else None
        ok = self._merger.update(rid, grid, tf_valid=msg.tf_valid, tf=tf)
        if not ok:
            self.get_logger().warning(
                f'[fleet] resolution mismatch from {rid} '
                f'({grid.info.resolution:.4f} vs {self._resolution:.4f}) — skipped'
            )

    def _on_robot_info(self, msg: RobotInfo) -> None:
        rid = msg.robot_id
        now = time.monotonic()
        if rid not in self._robots:
            self._robots[rid] = _RobotRecord(msg, now)
            self.get_logger().info(f'[fleet] discovered robot {rid}')
        else:
            self._robots[rid].info      = msg
            self._robots[rid].last_seen = now

    def _on_frontiers_raw(self, msg: FrontierArray) -> None:
        """Merge per-robot frontier reports into the global frontier pool."""
        now_stamp = self.get_clock().now().to_msg()
        for f in msg.frontiers:
            if f.frontier_id not in self._frontiers:
                self._frontiers[f.frontier_id] = f
            else:
                # Update utility (robot may have moved closer/farther)
                existing = self._frontiers[f.frontier_id]
                if f.assigned_to == '':
                    existing.utility = f.utility

    # ── Timer callbacks ───────────────────────────────────────────────────────

    def _merge_tick(self) -> None:
        # ── Expire timed-out robots ──────────────────────────────────────────
        now = time.monotonic()
        departed = [
            rid for rid, rec in self._robots.items()
            if now - rec.last_seen > self._hb_timeout
        ]
        for rid in departed:
            self.get_logger().warning(f'[fleet] robot {rid} timed out — removing')
            del self._robots[rid]
            self._merger.remove(rid)
            # Release frontiers assigned to departed robot
            for f in self._frontiers.values():
                if f.assigned_to == rid:
                    f.assigned_to = ''

        # ── Re-assign stale frontier claims ──────────────────────────────────
        # (robot claimed a frontier but hasn't updated its RobotInfo recently)
        for f in self._frontiers.values():
            if not f.assigned_to:
                continue
            rec = self._robots.get(f.assigned_to)
            if rec is None or now - rec.last_seen > self._frontier_to:
                self.get_logger().info(
                    f'[fleet] frontier {f.frontier_id} re-assigned '
                    f'(robot {f.assigned_to} stale)'
                )
                f.assigned_to = ''

        # ── Build merged map ─────────────────────────────────────────────────
        merged = self._merger.build_merged_grid()
        if merged:
            merged.header.stamp = self.get_clock().now().to_msg()
            self._merged_grid   = merged
            self._merged_pub.publish(merged)

        # ── Publish frontier list ─────────────────────────────────────────────
        if self._frontiers:
            fa = FrontierArray()
            fa.stamp       = self.get_clock().now().to_msg()
            fa.source_robot = 'fleet_manager'
            fa.frontiers   = list(self._frontiers.values())
            self._frontier_pub.publish(fa)

    def _status_tick(self) -> None:
        status = {
            'robots':           list(self._robots.keys()),
            'robot_count':      len(self._robots),
            'frontier_count':   len(self._frontiers),
            'assigned_frontiers': sum(
                1 for f in self._frontiers.values() if f.assigned_to
            ),
            'map_ready':        self._merged_grid is not None,
        }
        msg = String()
        msg.data = json.dumps(status)
        self._status_pub.publish(msg)

    # ── Service handlers ──────────────────────────────────────────────────────

    def _svc_register(
        self, req: RegisterRobot.Request, resp: RegisterRobot.Response
    ) -> RegisterRobot.Response:
        rid = req.robot_id
        # Detect namespace collision — assign suffix if needed
        if rid in self._robots:
            # Check if it's the same robot re-registering (same namespace)
            if self._robots[rid].info.namespace != req.namespace:
                suffix = 2
                while f'{rid}_{suffix}' in self._robots:
                    suffix += 1
                rid = f'{rid}_{suffix}'
        self._robots[rid] = _RobotRecord(RobotInfo(), time.monotonic())
        self._robots[rid].info.robot_id  = rid
        self._robots[rid].info.namespace = req.namespace
        resp.accepted     = True
        resp.assigned_id  = rid
        resp.message      = f'Registered as {rid}'
        self.get_logger().info(f'[fleet] registered robot {rid} ns={req.namespace}')
        return resp

    def _svc_request_frontier(
        self, req: RequestFrontier.Request, resp: RequestFrontier.Response
    ) -> RequestFrontier.Response:
        rid  = req.robot_id
        rx, ry = req.robot_pose_xy.x, req.robot_pose_xy.y

        # Pick highest-utility unassigned frontier (re-score by distance to this robot)
        best: Optional[Frontier] = None
        best_score = -1.0
        for f in self._frontiers.values():
            if f.assigned_to:
                continue
            dist  = math.hypot(f.centroid.x - rx, f.centroid.y - ry) + 1.0
            score = f.area_m2 / dist
            if score > best_score:
                best_score = score
                best       = f

        if best is None:
            resp.assigned = False
            resp.message  = 'No unassigned frontiers available'
            return resp

        best.assigned_to = rid
        best.assigned_at = self.get_clock().now().to_msg()
        resp.assigned    = True
        resp.frontier    = best
        resp.message     = f'Assigned frontier {best.frontier_id}'
        self.get_logger().info(
            f'[fleet] assigned frontier {best.frontier_id} to {rid} '
            f'(utility={best.utility:.2f})'
        )
        return resp


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


if __name__ == '__main__':
    main()
