"""
fleet_explorer_node.py — Per-robot cooperative exploration state machine.

This node runs on each robot.  It requests frontier assignments from the fleet
coordinator, drives the robot towards the assigned frontier using Nav2 NavigateToPose,
and marks the frontier complete once the robot arrives (or times out).

State machine:
  IDLE        → request a frontier from /fleet/request_frontier
  NAVIGATING  → send NavigateToPose goal; poll for completion
  ARRIVED     → update /fleet/robots status; loop back to IDLE
  STALLED     → robot stuck; wait backoff_sec, then request new frontier
  DONE        → all frontiers exhausted; switch to MAPPING status

Subscribes:
  /fleet/robots               (RobotInfo)    — monitor other robots
  /<ns>/rtabmap/odom          (Odometry)     — current robot pose

Publishes:
  /fleet/robots               (RobotInfo)    — heartbeat with status updates

Actions:
  /<ns>/navigate_to_pose      (Nav2 NavigateToPose)

Services (client):
  /fleet/request_frontier     (RequestFrontier)

Parameters:
  robot_id            str    "saltybot_1"
  robot_namespace     str    "/saltybot_1"
  exploration_hz      float  0.5   (state machine tick rate)
  goal_tolerance_m    float  0.5   (distance to frontier centroid = arrived)
  nav_timeout_sec     float  120.0 (cancel goal and stall if Nav2 doesn't finish)
  backoff_sec         float  5.0   (wait after stall before new frontier request)
  battery_min_pct     float  20.0  (stop exploration below this — return to dock)
"""

import math
import time

import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient
from rclpy.qos import (
    QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
)
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PoseStamped, Point, Quaternion
from nav2_msgs.action import NavigateToPose

from saltybot_fleet_msgs.msg import RobotInfo, Frontier
from saltybot_fleet_msgs.srv import RequestFrontier


class _State:
    IDLE       = 'IDLE'
    NAVIGATING = 'NAVIGATING'
    ARRIVED    = 'ARRIVED'
    STALLED    = 'STALLED'
    DONE       = 'DONE'


class FleetExplorerNode(Node):

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('robot_id',         'saltybot_1')
        self.declare_parameter('robot_namespace',  '/saltybot_1')
        self.declare_parameter('exploration_hz',   0.5)
        self.declare_parameter('goal_tolerance_m', 0.5)
        self.declare_parameter('nav_timeout_sec',  120.0)
        self.declare_parameter('backoff_sec',      5.0)
        self.declare_parameter('battery_min_pct',  20.0)

        self._robot_id    = self.get_parameter('robot_id').value
        self._ns          = self.get_parameter('robot_namespace').value
        hz                = self.get_parameter('exploration_hz').value
        self._goal_tol    = self.get_parameter('goal_tolerance_m').value
        self._nav_timeout = self.get_parameter('nav_timeout_sec').value
        self._backoff     = self.get_parameter('backoff_sec').value
        self._bat_min     = self.get_parameter('battery_min_pct').value

        # ── State ────────────────────────────────────────────────────────────
        self._state            = _State.IDLE
        self._current_frontier: Frontier | None      = None
        self._nav_goal_handle                        = None
        self._nav_start_time:  float                 = 0.0
        self._stall_until:     float                 = 0.0
        self._robot_x:         float                 = 0.0
        self._robot_y:         float                 = 0.0
        self._battery_pct:     float                 = -1.0

        # ── Subscriptions ────────────────────────────────────────────────────
        odom_topic = f'{self._ns}/rtabmap/odom'.replace('//', '/')
        self.create_subscription(Odometry, odom_topic, self._on_odom, 10)
        self.create_subscription(RobotInfo, '/fleet/robots', self._on_fleet_info, 10)

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

        # ── Service clients ──────────────────────────────────────────────────
        self._frontier_client = self.create_client(
            RequestFrontier, '/fleet/request_frontier'
        )

        # ── Nav2 action client ────────────────────────────────────────────────
        nav_action = f'{self._ns}/navigate_to_pose'.replace('//', '/')
        self._nav_client = ActionClient(self, NavigateToPose, nav_action)

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

        self.get_logger().info(
            f'[{self._robot_id}] fleet_explorer ready (battery_min={self._bat_min}%)'
        )

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

    def _on_odom(self, msg: Odometry) -> None:
        p = msg.pose.pose.position
        self._robot_x = p.x
        self._robot_y = p.y

    def _on_fleet_info(self, msg: RobotInfo) -> None:
        # Track own battery from heartbeats of other nodes (e.g. battery_manager)
        if msg.robot_id == self._robot_id:
            self._battery_pct = msg.battery_pct

    # ── State machine ─────────────────────────────────────────────────────────

    def _explore_tick(self) -> None:
        now = time.monotonic()

        # Low battery — abort
        if self._battery_pct > 0 and self._battery_pct < self._bat_min:
            if self._state != _State.DONE:
                self.get_logger().warning(
                    f'[{self._robot_id}] battery {self._battery_pct:.0f}% < '
                    f'{self._bat_min:.0f}% — stopping exploration'
                )
                self._cancel_nav()
                self._state = _State.DONE
            self._publish_info(RobotInfo.STATUS_RETURNING)
            return

        if self._state == _State.IDLE:
            self._request_frontier()

        elif self._state == _State.NAVIGATING:
            self._check_navigation(now)

        elif self._state == _State.ARRIVED:
            self._on_arrived()

        elif self._state == _State.STALLED:
            if now >= self._stall_until:
                self.get_logger().info(f'[{self._robot_id}] backoff over — requesting new frontier')
                self._state = _State.IDLE

        elif self._state == _State.DONE:
            self._publish_info(RobotInfo.STATUS_IDLE)

    def _request_frontier(self) -> None:
        if not self._frontier_client.wait_for_service(timeout_sec=2.0):
            self.get_logger().warning(
                f'[{self._robot_id}] /fleet/request_frontier unavailable — retrying'
            )
            return
        req = RequestFrontier.Request()
        req.robot_id         = self._robot_id
        req.robot_pose_xy    = Point(x=self._robot_x, y=self._robot_y, z=0.0)
        future = self._frontier_client.call_async(req)
        future.add_done_callback(self._on_frontier_response)

    def _on_frontier_response(self, future) -> None:
        try:
            resp = future.result()
        except Exception as exc:
            self.get_logger().error(f'[{self._robot_id}] frontier request failed: {exc}')
            return
        if not resp.assigned:
            self.get_logger().info(f'[{self._robot_id}] {resp.message} — exploration complete')
            self._state = _State.DONE
            return
        self._current_frontier = resp.frontier
        self.get_logger().info(
            f'[{self._robot_id}] assigned frontier {resp.frontier.frontier_id} '
            f'at ({resp.frontier.centroid.x:.2f}, {resp.frontier.centroid.y:.2f})'
        )
        self._send_nav_goal(resp.frontier)

    def _send_nav_goal(self, frontier: Frontier) -> None:
        if not self._nav_client.wait_for_server(timeout_sec=5.0):
            self.get_logger().error(
                f'[{self._robot_id}] Nav2 action server not available'
            )
            self._stall(reason='nav_server_missing')
            return

        goal = NavigateToPose.Goal()
        goal.pose.header.frame_id      = 'fleet_map'
        goal.pose.header.stamp         = self.get_clock().now().to_msg()
        goal.pose.pose.position.x      = frontier.centroid.x
        goal.pose.pose.position.y      = frontier.centroid.y
        goal.pose.pose.orientation.w   = 1.0

        send_fut = self._nav_client.send_goal_async(goal)
        send_fut.add_done_callback(self._on_nav_goal_accepted)
        self._nav_start_time = time.monotonic()
        self._state = _State.NAVIGATING
        self._publish_info(RobotInfo.STATUS_EXPLORING)

    def _on_nav_goal_accepted(self, future) -> None:
        goal_handle = future.result()
        if not goal_handle.accepted:
            self.get_logger().warning(f'[{self._robot_id}] Nav2 rejected goal')
            self._stall(reason='goal_rejected')
            return
        self._nav_goal_handle = goal_handle
        result_fut = goal_handle.get_result_async()
        result_fut.add_done_callback(self._on_nav_result)

    def _on_nav_result(self, future) -> None:
        try:
            result = future.result()
            self.get_logger().info(
                f'[{self._robot_id}] Nav2 completed frontier '
                f'{self._current_frontier.frontier_id if self._current_frontier else "?"}'
            )
            self._state = _State.ARRIVED
        except Exception as exc:
            self.get_logger().warning(f'[{self._robot_id}] Nav2 result error: {exc}')
            self._stall(reason='nav_error')

    def _check_navigation(self, now: float) -> None:
        # ── Timeout guard ────────────────────────────────────────────────────
        if now - self._nav_start_time > self._nav_timeout:
            self.get_logger().warning(
                f'[{self._robot_id}] navigation timeout after {self._nav_timeout:.0f}s'
            )
            self._cancel_nav()
            self._stall(reason='timeout')
            return

        # ── Arrival check (proximity) ────────────────────────────────────────
        if self._current_frontier is not None:
            dist = math.hypot(
                self._current_frontier.centroid.x - self._robot_x,
                self._current_frontier.centroid.y - self._robot_y,
            )
            if dist < self._goal_tol:
                self._state = _State.ARRIVED

    def _on_arrived(self) -> None:
        fid = self._current_frontier.frontier_id if self._current_frontier else '?'
        self.get_logger().info(f'[{self._robot_id}] arrived at frontier {fid}')
        self._current_frontier = None
        self._nav_goal_handle  = None
        self._state = _State.IDLE
        self._publish_info(RobotInfo.STATUS_EXPLORING)

    def _stall(self, reason: str = '') -> None:
        self.get_logger().warning(f'[{self._robot_id}] stalled ({reason}), backing off {self._backoff}s')
        self._state       = _State.STALLED
        self._stall_until = time.monotonic() + self._backoff
        self._current_frontier = None

    def _cancel_nav(self) -> None:
        if self._nav_goal_handle is not None:
            self._nav_goal_handle.cancel_goal_async()
            self._nav_goal_handle = None

    def _publish_info(self, status: int) -> None:
        info = RobotInfo()
        info.stamp           = self.get_clock().now().to_msg()
        info.robot_id        = self._robot_id
        info.namespace       = self._ns
        info.status          = status
        info.battery_pct     = self._battery_pct
        info.active_frontier = (
            self._current_frontier.frontier_id
            if self._current_frontier else ''
        )
        self._info_pub.publish(info)


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


if __name__ == '__main__':
    main()
