"""test_waypoint_follower.py — Unit tests for Nav2-UWB waypoint logic (Issue #599).

Tests are ROS2-free.  Run with:
    python3 -m pytest \
        jetson/ros2_ws/src/saltybot_nav2_uwb/test/test_waypoint_follower.py -v

Coverage:
    • Waypoint dataclass construction, validation, serialisation
    • WaypointSequencer state machine (all transitions)
    • Progress tracking
    • UWB pose-bridge 2-D TF maths
    • waypoint_to_pose_stamped conversion (quaternion from yaw)
"""

from __future__ import annotations

import math
import sys
import os

# Allow import without ROS2 install
_pkg_dir = os.path.join(os.path.dirname(__file__), "..", "saltybot_nav2_uwb")
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, _pkg_dir)

import pytest

# ── Stub ROS2 modules before importing any node files ────────────────────────

def _install_stubs():
    """Inject minimal stubs so node modules are importable without ROS2."""
    from unittest.mock import MagicMock
    import types

    for mod_name in [
        "rclpy", "rclpy.node", "rclpy.action", "rclpy.time",
        "rclpy.callback_groups", "rclpy.executors",
        "geometry_msgs", "geometry_msgs.msg",
        "nav_msgs", "nav_msgs.msg",
        "nav2_msgs", "nav2_msgs.action",
        "std_msgs", "std_msgs.msg",
        "tf2_ros", "action_msgs", "action_msgs.msg",
    ]:
        sys.modules.setdefault(mod_name, MagicMock())

    # Make geometry_msgs.msg.PoseWithCovarianceStamped / TransformStamped real
    # enough to be used as type hints
    sys.modules["geometry_msgs.msg"].PoseWithCovarianceStamped = MagicMock
    sys.modules["geometry_msgs.msg"].TransformStamped          = MagicMock
    sys.modules["nav_msgs.msg"].Odometry                       = MagicMock
    sys.modules["std_msgs.msg"].Empty                          = MagicMock
    sys.modules["std_msgs.msg"].String                         = MagicMock
    sys.modules["tf2_ros"].TransformBroadcaster                = MagicMock

_install_stubs()

from waypoint_sequencer import Waypoint, WaypointSequencer

# ── helpers from uwb_pose_bridge_node (pure-python section) ──────────────────

from uwb_pose_bridge_node import (
    _yaw_from_q,
    _q_from_yaw,
    _invert_2d,
    _compose_2d,
    compute_map_odom_tf,
)


# ─── Waypoint dataclass ────────────────────────────────────────────────────────

class TestWaypointConstruction:
    def test_basic_construction(self):
        wp = Waypoint(x=1.0, y=2.0, yaw_deg=45.0)
        assert wp.x == 1.0
        assert wp.y == 2.0
        assert wp.yaw_deg == 45.0

    def test_default_yaw_is_zero(self):
        wp = Waypoint(x=0.0, y=0.0)
        assert wp.yaw_deg == 0.0

    def test_negative_coordinates(self):
        wp = Waypoint(x=-3.5, y=-1.2, yaw_deg=-90.0)
        assert wp.x == -3.5
        assert wp.yaw_deg == -90.0

    def test_yaw_boundary_360(self):
        wp = Waypoint(x=0.0, y=0.0, yaw_deg=360.0)
        assert wp.yaw_deg == 360.0

    def test_yaw_boundary_neg360(self):
        wp = Waypoint(x=0.0, y=0.0, yaw_deg=-360.0)
        assert wp.yaw_deg == -360.0

    def test_yaw_out_of_range_raises(self):
        with pytest.raises(ValueError, match="yaw_deg"):
            Waypoint(x=0.0, y=0.0, yaw_deg=400.0)

    def test_yaw_out_of_range_negative_raises(self):
        with pytest.raises(ValueError, match="yaw_deg"):
            Waypoint(x=0.0, y=0.0, yaw_deg=-400.0)


class TestWaypointFromDict:
    def test_full_dict(self):
        wp = Waypoint.from_dict({"x": 1.5, "y": -2.5, "yaw_deg": 30.0})
        assert wp.x == 1.5
        assert wp.y == -2.5
        assert wp.yaw_deg == 30.0

    def test_missing_yaw_defaults_zero(self):
        wp = Waypoint.from_dict({"x": 0.0, "y": 1.0})
        assert wp.yaw_deg == 0.0

    def test_string_values_coerced(self):
        wp = Waypoint.from_dict({"x": "1.0", "y": "2.0", "yaw_deg": "45"})
        assert wp.x == 1.0

    def test_missing_x_raises(self):
        with pytest.raises(KeyError):
            Waypoint.from_dict({"y": 1.0})

    def test_missing_y_raises(self):
        with pytest.raises(KeyError):
            Waypoint.from_dict({"x": 1.0})


class TestWaypointToDict:
    def test_roundtrip(self):
        wp = Waypoint(x=1.0, y=2.0, yaw_deg=90.0)
        d = wp.to_dict()
        wp2 = Waypoint.from_dict(d)
        assert wp2.x == wp.x
        assert wp2.y == wp.y
        assert wp2.yaw_deg == wp.yaw_deg

    def test_keys_present(self):
        d = Waypoint(x=0.0, y=0.0).to_dict()
        assert "x" in d and "y" in d and "yaw_deg" in d


# ─── WaypointSequencer state machine ─────────────────────────────────────────

def _three_wp():
    return [Waypoint(float(i), 0.0) for i in range(3)]


class TestSequencerInitialState:
    def setup_method(self):
        self.seq = WaypointSequencer()

    def test_initial_state_idle(self):
        assert self.seq.state == WaypointSequencer.IDLE

    def test_is_not_running(self):
        assert not self.seq.is_running

    def test_is_not_done(self):
        assert not self.seq.is_done

    def test_no_current_waypoint(self):
        assert self.seq.current_waypoint is None

    def test_total_zero(self):
        assert self.seq.total == 0

    def test_progress_zero(self):
        assert self.seq.progress == (0, 0)

    def test_no_abort_reason(self):
        assert self.seq.abort_reason is None


class TestSequencerStart:
    def setup_method(self):
        self.seq = WaypointSequencer()

    def test_transitions_to_running(self):
        self.seq.start(_three_wp())
        assert self.seq.state == WaypointSequencer.RUNNING

    def test_is_running_true(self):
        self.seq.start(_three_wp())
        assert self.seq.is_running

    def test_is_done_false(self):
        self.seq.start(_three_wp())
        assert not self.seq.is_done

    def test_current_index_zero(self):
        self.seq.start(_three_wp())
        assert self.seq.current_index == 0

    def test_total_set(self):
        self.seq.start(_three_wp())
        assert self.seq.total == 3

    def test_first_waypoint_returned(self):
        wps = _three_wp()
        self.seq.start(wps)
        assert self.seq.current_waypoint is wps[0]

    def test_empty_list_raises(self):
        with pytest.raises(ValueError, match="empty"):
            self.seq.start([])

    def test_restart_resets_state(self):
        self.seq.start(_three_wp())
        self.seq.advance()
        self.seq.start([Waypoint(9.0, 9.0)])
        assert self.seq.current_index == 0
        assert self.seq.total == 1


class TestSequencerAdvance:
    def setup_method(self):
        self.seq = WaypointSequencer()

    def test_advance_moves_to_next(self):
        self.seq.start(_three_wp())
        self.seq.advance()
        assert self.seq.current_index == 1

    def test_advance_twice(self):
        self.seq.start(_three_wp())
        self.seq.advance()
        self.seq.advance()
        assert self.seq.current_index == 2

    def test_advance_last_sets_succeeded(self):
        self.seq.start(_three_wp())
        for _ in range(3):
            self.seq.advance()
        assert self.seq.state == WaypointSequencer.SUCCEEDED

    def test_advance_last_sets_done(self):
        self.seq.start(_three_wp())
        for _ in range(3):
            self.seq.advance()
        assert self.seq.is_done

    def test_advance_on_idle_raises(self):
        with pytest.raises(RuntimeError, match="RUNNING"):
            self.seq.advance()

    def test_advance_on_aborted_raises(self):
        self.seq.start(_three_wp())
        self.seq.abort()
        with pytest.raises(RuntimeError, match="RUNNING"):
            self.seq.advance()

    def test_no_current_waypoint_after_succeed(self):
        self.seq.start(_three_wp())
        for _ in range(3):
            self.seq.advance()
        assert self.seq.current_waypoint is None


class TestSequencerAbort:
    def setup_method(self):
        self.seq = WaypointSequencer()

    def test_abort_transitions_state(self):
        self.seq.start(_three_wp())
        self.seq.abort("Nav2 failed")
        assert self.seq.state == WaypointSequencer.ABORTED

    def test_abort_sets_reason(self):
        self.seq.start(_three_wp())
        self.seq.abort("timeout")
        assert self.seq.abort_reason == "timeout"

    def test_abort_sets_done(self):
        self.seq.start(_three_wp())
        self.seq.abort()
        assert self.seq.is_done

    def test_abort_no_reason(self):
        self.seq.start(_three_wp())
        self.seq.abort()
        assert self.seq.abort_reason == ""


class TestSequencerCancel:
    def setup_method(self):
        self.seq = WaypointSequencer()

    def test_cancel_transitions_state(self):
        self.seq.start(_three_wp())
        self.seq.cancel()
        assert self.seq.state == WaypointSequencer.CANCELED

    def test_cancel_sets_done(self):
        self.seq.start(_three_wp())
        self.seq.cancel()
        assert self.seq.is_done

    def test_cancel_clears_current_waypoint(self):
        self.seq.start(_three_wp())
        self.seq.cancel()
        assert self.seq.current_waypoint is None


class TestSequencerProgress:
    def setup_method(self):
        self.seq = WaypointSequencer()
        self.seq.start(_three_wp())

    def test_initial_progress(self):
        assert self.seq.progress == (0, 3)

    def test_progress_after_one(self):
        self.seq.advance()
        assert self.seq.progress == (1, 3)

    def test_remaining_initial(self):
        assert self.seq.remaining == 3

    def test_remaining_after_advance(self):
        self.seq.advance()
        assert self.seq.remaining == 2

    def test_remaining_after_all(self):
        for _ in range(3):
            self.seq.advance()
        assert self.seq.remaining == 0


class TestSequencerStatusDict:
    def setup_method(self):
        self.seq = WaypointSequencer()

    def test_idle_status_dict(self):
        d = self.seq.status_dict()
        assert d["state"] == "idle"
        assert d["total"] == 0
        assert d["current_wp"] is None

    def test_running_status_dict(self):
        self.seq.start(_three_wp())
        d = self.seq.status_dict()
        assert d["state"] == "running"
        assert d["total"] == 3
        assert d["current_wp"] is not None
        assert "x" in d["current_wp"]

    def test_succeeded_status_dict(self):
        self.seq.start(_three_wp())
        for _ in range(3):
            self.seq.advance()
        d = self.seq.status_dict()
        assert d["state"] == "succeeded"
        assert d["remaining"] == 0


# ─── 2-D TF maths (uwb_pose_bridge_node) ──────────────────────────────────────

EPS = 1e-9


class TestYawFromQ:
    def test_zero_yaw(self):
        qx, qy, qz, qw = _q_from_yaw(0.0)
        assert abs(_yaw_from_q(qx, qy, qz, qw)) < EPS

    def test_90_degrees(self):
        yaw = math.pi / 2
        qx, qy, qz, qw = _q_from_yaw(yaw)
        assert abs(_yaw_from_q(qx, qy, qz, qw) - yaw) < 1e-6

    def test_minus_90_degrees(self):
        yaw = -math.pi / 2
        qx, qy, qz, qw = _q_from_yaw(yaw)
        assert abs(_yaw_from_q(qx, qy, qz, qw) - yaw) < 1e-6

    def test_180_degrees(self):
        yaw = math.pi
        qx, qy, qz, qw = _q_from_yaw(yaw)
        # atan2 wraps to -π at exactly π; tolerate either
        result = _yaw_from_q(qx, qy, qz, qw)
        assert abs(abs(result) - math.pi) < 1e-6

    def test_roundtrip_arbitrary(self):
        for deg in (0, 30, 60, 90, 120, 150, 170, -45, -120):
            yaw = math.radians(deg)
            qx, qy, qz, qw = _q_from_yaw(yaw)
            back = _yaw_from_q(qx, qy, qz, qw)
            assert abs(back - yaw) < 1e-6, f"roundtrip failed for {deg}°"


class TestInvert2D:
    def test_identity(self):
        xi, yi, thi = _invert_2d(0.0, 0.0, 0.0)
        assert abs(xi) < EPS and abs(yi) < EPS and abs(thi) < EPS

    def test_translation_only(self):
        # invert pure translation: inv(T(3,4,0)) = T(-3,-4,0)
        xi, yi, thi = _invert_2d(3.0, 4.0, 0.0)
        assert abs(xi + 3.0) < EPS
        assert abs(yi + 4.0) < EPS
        assert abs(thi) < EPS

    def test_rotation_only(self):
        # invert pure rotation by 90°
        xi, yi, thi = _invert_2d(0.0, 0.0, math.pi / 2)
        assert abs(thi + math.pi / 2) < EPS

    def test_compose_with_inverse_is_identity(self):
        x, y, th = 2.5, -1.3, math.radians(47)
        ix, iy, ith = _invert_2d(x, y, th)
        rx, ry, rth = _compose_2d(x, y, th, ix, iy, ith)
        assert abs(rx) < 1e-9
        assert abs(ry) < 1e-9
        assert abs(rth % (2 * math.pi)) < 1e-9 or abs(rth) < 1e-9


class TestComputeMapOdomTf:
    def test_identity_pose(self):
        """If both poses are identity, map→odom is identity."""
        mx, my, mth = compute_map_odom_tf(0, 0, 0, 0, 0, 0)
        assert abs(mx) < EPS and abs(my) < EPS and abs(mth) < EPS

    def test_odom_is_map(self):
        """If robot is at same position in both frames, map→odom is identity."""
        pose = (1.0, 2.0, math.radians(45))
        mx, my, mth = compute_map_odom_tf(*pose, *pose)
        assert abs(mx) < 1e-9 and abs(my) < 1e-9 and abs(mth) < 1e-9

    def test_pure_translation_offset(self):
        """map_bl at (5,0,0), odom_bl at (2,0,0) → map_odom at (3,0,0)."""
        mx, my, mth = compute_map_odom_tf(5.0, 0.0, 0.0, 2.0, 0.0, 0.0)
        assert abs(mx - 3.0) < 1e-9
        assert abs(my) < 1e-9
        assert abs(mth) < 1e-9

    def test_tf_reconstruction(self):
        """T_map_odom ⊗ T_odom_bl == T_map_bl (round-trip check)."""
        map_bl  = (3.0, 1.5, math.radians(30))
        odom_bl = (1.0, 0.5, math.radians(10))
        mo = compute_map_odom_tf(*map_bl, *odom_bl)
        # Recompose: T_map_odom ⊗ T_odom_bl should equal T_map_bl
        rx, ry, rth = _compose_2d(*mo, *odom_bl)
        assert abs(rx - map_bl[0]) < 1e-9
        assert abs(ry - map_bl[1]) < 1e-9
        assert abs(rth - map_bl[2]) < 1e-9

    def test_rotated_frame(self):
        """Verify correctness with a 90° heading offset."""
        # map: robot at (1,0) facing 90°; odom: robot at (0,1) facing 0°
        mo_x, mo_y, mo_th = compute_map_odom_tf(
            1.0, 0.0, math.pi / 2,
            0.0, 1.0, 0.0,
        )
        # Recompose and check
        rx, ry, rth = _compose_2d(mo_x, mo_y, mo_th, 0.0, 1.0, 0.0)
        assert abs(rx - 1.0) < 1e-9
        assert abs(ry - 0.0) < 1e-9
        assert abs(rth - math.pi / 2) < 1e-9
