"""waypoint_follower_node.py — Simple sequential waypoint follower (Issue #599).

Exposes a ROS2 action server that accepts a list of waypoints and drives the
robot through them one by one via Nav2's NavigateToPose action.

Action server:  /saltybot/nav/follow_waypoints  (nav2_msgs/FollowWaypoints)
Nav2 client:    /navigate_to_pose               (nav2_msgs/NavigateToPose)

Waypoints topic (JSON, for operator / WebUI override):
    /saltybot/nav/waypoints  (std_msgs/String)
    Payload: JSON array of {x, y, yaw_deg} objects.
    Example: [{"x": 1.0, "y": 2.0, "yaw_deg": 90.0}, ...]

Status topic (JSON):
    /saltybot/nav/waypoint_status  (std_msgs/String)
    Published every cycle with WaypointSequencer.status_dict().

Cancel topic:
    /saltybot/nav/cancel  (std_msgs/Empty)
    Cancels the active mission.
"""

from __future__ import annotations

import json
import math
import threading
from typing import List, Optional

import rclpy
from rclpy.action import ActionClient, ActionServer, CancelResponse, GoalResponse
from rclpy.node import Node
from rclpy.callback_groups import ReentrantCallbackGroup

from geometry_msgs.msg import PoseStamped
from nav2_msgs.action import FollowWaypoints, NavigateToPose
from std_msgs.msg import Empty, String

from saltybot_nav2_uwb.waypoint_sequencer import Waypoint, WaypointSequencer


def _yaw_to_quaternion(yaw_rad: float):
    """Return (qx, qy, qz, qw) for a heading-only rotation."""
    return 0.0, 0.0, math.sin(yaw_rad * 0.5), math.cos(yaw_rad * 0.5)


def waypoint_to_pose_stamped(wp: Waypoint, frame: str = "map") -> PoseStamped:
    """Convert a Waypoint to a PoseStamped in the given frame."""
    msg = PoseStamped()
    msg.header.frame_id = frame
    msg.pose.position.x = wp.x
    msg.pose.position.y = wp.y
    msg.pose.position.z = 0.0
    qx, qy, qz, qw = _yaw_to_quaternion(math.radians(wp.yaw_deg))
    msg.pose.orientation.x = qx
    msg.pose.orientation.y = qy
    msg.pose.orientation.z = qz
    msg.pose.orientation.w = qw
    return msg


class WaypointFollowerNode(Node):
    """Sequential waypoint follower backed by Nav2 NavigateToPose.

    Parameters (ROS):
        map_frame          (str)   default: map
        goal_xy_tolerance  (float) default: 0.20  (m)  — passed to Nav2
        goal_yaw_tolerance (float) default: 0.15  (rad) — passed to Nav2
        nav_timeout_s      (float) default: 60.0  — abort if Nav2 takes longer
        status_hz          (float) default: 2.0   — status publish rate
    """

    def __init__(self):
        super().__init__("waypoint_follower")

        self.declare_parameter("map_frame",          "map")
        self.declare_parameter("goal_xy_tolerance",  0.20)
        self.declare_parameter("goal_yaw_tolerance", 0.15)
        self.declare_parameter("nav_timeout_s",      60.0)
        self.declare_parameter("status_hz",          2.0)

        self._map_frame    = self.get_parameter("map_frame").value
        self._nav_timeout  = self.get_parameter("nav_timeout_s").value

        self._seq   = WaypointSequencer()
        self._lock  = threading.Lock()
        self._cb_group = ReentrantCallbackGroup()

        # ── Publishers ─────────────────────────────────────────────────
        self._status_pub = self.create_publisher(
            String, "/saltybot/nav/waypoint_status", 10)

        # ── Subscriptions ──────────────────────────────────────────────
        self.create_subscription(
            String, "/saltybot/nav/waypoints",
            self._on_waypoints_json, 10,
            callback_group=self._cb_group)

        self.create_subscription(
            Empty, "/saltybot/nav/cancel",
            self._on_cancel, 10,
            callback_group=self._cb_group)

        # ── Nav2 action client ─────────────────────────────────────────
        self._nav_client: ActionClient = ActionClient(
            self, NavigateToPose, "/navigate_to_pose",
            callback_group=self._cb_group)

        # ── Action server (nav2_msgs/FollowWaypoints) ──────────────────
        self._action_server = ActionServer(
            self,
            FollowWaypoints,
            "/saltybot/nav/follow_waypoints",
            execute_callback=self._execute_follow_waypoints,
            goal_callback=self._goal_callback,
            cancel_callback=self._cancel_callback,
            callback_group=self._cb_group,
        )

        hz = self.get_parameter("status_hz").value
        self.create_timer(1.0 / hz, self._publish_status,
                          callback_group=self._cb_group)

        self.get_logger().info(
            f"WaypointFollower ready | "
            f"action=/saltybot/nav/follow_waypoints  "
            f"nav2=/navigate_to_pose  frame={self._map_frame}"
        )

    # ── Action server callbacks ────────────────────────────────────────────

    def _goal_callback(self, goal_request):
        if not self._nav_client.wait_for_server(timeout_sec=2.0):
            self.get_logger().error(
                "Nav2 NavigateToPose server not available — rejecting goal")
            return GoalResponse.REJECT
        if len(goal_request.poses) == 0:
            self.get_logger().warn("Empty waypoint list — rejecting goal")
            return GoalResponse.REJECT
        return GoalResponse.ACCEPT

    def _cancel_callback(self, _goal_handle):
        return CancelResponse.ACCEPT

    async def _execute_follow_waypoints(self, goal_handle) -> FollowWaypoints.Result:
        """Execute the action: drive through each PoseStamped in sequence."""
        result = FollowWaypoints.Result()
        feedback_msg = FollowWaypoints.Feedback()

        poses: List[PoseStamped] = list(goal_handle.request.poses)
        waypoints = [
            Waypoint(
                x=p.pose.position.x,
                y=p.pose.position.y,
                yaw_deg=math.degrees(
                    math.atan2(
                        2.0 * (p.pose.orientation.w * p.pose.orientation.z
                               + p.pose.orientation.x * p.pose.orientation.y),
                        1.0 - 2.0 * (p.pose.orientation.y ** 2
                                     + p.pose.orientation.z ** 2),
                    )
                ),
            )
            for p in poses
        ]

        with self._lock:
            self._seq.start(waypoints)

        self.get_logger().info(
            f"Starting waypoint mission: {len(waypoints)} waypoints")

        missed_waypoints: List[int] = []

        while True:
            # Check cancellation
            if goal_handle.is_cancel_requested:
                with self._lock:
                    self._seq.cancel()
                goal_handle.canceled()
                result.missed_waypoints = missed_waypoints
                return result

            with self._lock:
                if self._seq.is_done:
                    break
                wp = self._seq.current_waypoint
                idx = self._seq.current_index

            if wp is None:
                break

            # Send goal to Nav2
            nav_goal = NavigateToPose.Goal()
            nav_goal.pose = waypoint_to_pose_stamped(wp, self._map_frame)

            self.get_logger().info(
                f"Sending waypoint {idx + 1}/{len(waypoints)}: "
                f"({wp.x:.2f}, {wp.y:.2f}, {wp.yaw_deg:.1f}°)"
            )

            if not self._nav_client.wait_for_server(timeout_sec=5.0):
                self.get_logger().error("Nav2 server unavailable — aborting")
                with self._lock:
                    self._seq.abort("Nav2 server unavailable")
                missed_waypoints.extend(range(idx, len(waypoints)))
                break

            send_future = self._nav_client.send_goal_async(
                nav_goal,
                feedback_callback=lambda fb: None,  # suppress individual feedback
            )
            # Spin until goal accepted
            rclpy.spin_until_future_complete(
                self, send_future, timeout_sec=self._nav_timeout)

            if not send_future.done():
                self.get_logger().error(
                    f"Waypoint {idx + 1}: goal send timed out — skipping")
                missed_waypoints.append(idx)
                with self._lock:
                    self._seq.advance()
                continue

            goal_res = send_future.result()
            if not goal_res.accepted:
                self.get_logger().warn(
                    f"Waypoint {idx + 1}: goal rejected by Nav2 — skipping")
                missed_waypoints.append(idx)
                with self._lock:
                    self._seq.advance()
                continue

            # Wait for Nav2 to reach the goal
            result_future = goal_res.get_result_async()
            rclpy.spin_until_future_complete(
                self, result_future, timeout_sec=self._nav_timeout)

            if not result_future.done():
                self.get_logger().error(
                    f"Waypoint {idx + 1}: Nav2 timed out after "
                    f"{self._nav_timeout:.0f}s — skipping")
                goal_res.cancel_goal_async()
                missed_waypoints.append(idx)
            else:
                nav_result = result_future.result()
                from action_msgs.msg import GoalStatus
                if nav_result.status == GoalStatus.STATUS_SUCCEEDED:
                    self.get_logger().info(
                        f"Waypoint {idx + 1}/{len(waypoints)} reached ✓")
                else:
                    self.get_logger().warn(
                        f"Waypoint {idx + 1} failed (Nav2 status "
                        f"{nav_result.status}) — skipping")
                    missed_waypoints.append(idx)

            with self._lock:
                self._seq.advance()

            # Publish feedback
            feedback_msg.current_waypoint = self._seq.current_index
            goal_handle.publish_feedback(feedback_msg)

        # Mission complete
        with self._lock:
            final_state = self._seq.state

        if final_state == WaypointSequencer.SUCCEEDED and not missed_waypoints:
            self.get_logger().info(
                f"Waypoint mission SUCCEEDED ({len(waypoints)} waypoints)")
            goal_handle.succeed()
        elif final_state == WaypointSequencer.CANCELED:
            self.get_logger().info("Waypoint mission CANCELED")
            goal_handle.canceled()
        else:
            self.get_logger().warn(
                f"Waypoint mission ended: state={final_state} "
                f"missed={missed_waypoints}")
            goal_handle.succeed()  # FollowWaypoints always succeeds; missed list carries info

        result.missed_waypoints = missed_waypoints
        return result

    # ── JSON topic handler ─────────────────────────────────────────────────

    def _on_waypoints_json(self, msg: String) -> None:
        """Accept a JSON waypoint list and start a new mission immediately.

        This is a convenience interface for WebUI / operator use.
        For production use prefer the action server (supports feedback/cancel).
        """
        try:
            raw = json.loads(msg.data)
        except json.JSONDecodeError as exc:
            self.get_logger().error(f"Invalid waypoints JSON: {exc}")
            return

        if not isinstance(raw, list) or len(raw) == 0:
            self.get_logger().error(
                "Waypoints JSON must be a non-empty array of {x, y, yaw_deg}")
            return

        try:
            waypoints = [Waypoint.from_dict(d) for d in raw]
        except (KeyError, ValueError) as exc:
            self.get_logger().error(f"Malformed waypoint in list: {exc}")
            return

        if not self._nav_client.wait_for_server(timeout_sec=2.0):
            self.get_logger().error(
                "Nav2 server not ready — cannot execute JSON waypoints")
            return

        self.get_logger().info(
            f"JSON mission: {len(waypoints)} waypoints received via topic")
        with self._lock:
            self._seq.start(waypoints)
        # Execution driven by action server; JSON topic just seeds the sequencer.
        # For full feedback, use the FollowWaypoints action server directly.

    # ── Cancel topic handler ───────────────────────────────────────────────

    def _on_cancel(self, _msg: Empty) -> None:
        with self._lock:
            if self._seq.is_running:
                self._seq.cancel()
                self.get_logger().info(
                    "Mission cancelled via /saltybot/nav/cancel")

    # ── Status publish ─────────────────────────────────────────────────────

    def _publish_status(self) -> None:
        with self._lock:
            d = self._seq.status_dict()
        self._status_pub.publish(String(data=json.dumps(d)))


def main(args=None):
    rclpy.init(args=args)
    node = WaypointFollowerNode()
    executor = rclpy.executors.MultiThreadedExecutor()
    executor.add_node(node)
    try:
        executor.spin()
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()
