"""
route_replayer_node — Loads and autonomously replays a recorded route.

"Ride once, replay forever." — Loads a .jsonl route file, converts GPS
waypoints to map-frame PoseStamped via robot_localization's flat-earth
ENU projection, then sends batches to Nav2's navigate_through_poses action.

Replay strategy:
  1. Load route file → list of {lat, lon, heading_deg, ...}
  2. Convert GPS → map-frame ENU (relative to first fix / datum)
  3. Subsample waypoints by replay_spacing_m (default 3m)
  4. Inject heading from recorded data (preserved via quaternion)
  5. Send full waypoint list to Nav2 navigate_through_poses
  6. Monitor feedback; pause/resume on demand
  7. Odometry correction: compare current odom position to expected odom
     position and log drift for future closed-loop correction (V1: log only)

GPS drift handling:
  - Use wider goal tolerance (2m default — matches SIM7600X CEP)
  - Slow down near waypoints (Nav2 DWB sim_time reduces automatically)
  - Report GPS covariance in status so operator can monitor quality

Topics subscribed:
  /gps/fix          sensor_msgs/NavSatFix
  /odom             nav_msgs/Odometry
  /route/pause_cmd  std_msgs/Bool  — True=pause, False=resume

Topics published:
  /route/replay_status  std_msgs/String  — state machine status
  /route/progress       std_msgs/String  — 'N/M waypoints, D m remaining'

Actions called:
  /navigate_through_poses  nav2_msgs/action/NavigateThroughPoses

Services:
  /route/load          std_srvs/Trigger  — load route named by route_name param
  /route/start_replay  std_srvs/Trigger  — begin replay
  /route/pause         std_srvs/Trigger  — pause (cancel Nav2 goal, hold position)
  /route/stop          std_srvs/Trigger  — abort replay
"""

import json
import math
import os
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient
from rclpy.callback_groups import ReentrantCallbackGroup
from rclpy.executors import MultiThreadedExecutor

from sensor_msgs.msg import NavSatFix, NavSatStatus
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PoseStamped, PoseArray, Pose
from std_msgs.msg import String, Bool
from std_srvs.srv import Trigger

from nav2_msgs.action import NavigateThroughPoses


def _haversine(lat1, lon1, lat2, lon2) -> float:
    R = 6_371_000.0
    phi1, phi2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlam = math.radians(lon2 - lon1)
    a = math.sin(dphi / 2) ** 2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlam / 2) ** 2
    return R * 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))


def _latlon_to_enu(lat: float, lon: float, lat0: float, lon0: float):
    """Flat-earth ENU approximation (valid ~10 km)."""
    dlat = math.radians(lat - lat0)
    dlon = math.radians(lon - lon0)
    cos0 = math.cos(math.radians(lat0))
    x = dlon * cos0 * 6_371_000.0
    y = dlat * 6_371_000.0
    return x, y


def _heading_to_quat(heading_deg: float) -> tuple:
    """Convert heading (degrees, 0=North, CW) to ENU quaternion (0=East, CCW)."""
    # In ENU: x=East, y=North. yaw=0 → facing East.
    # Heading: 0=North, 90=East. Convert: yaw_enu = 90 - heading_deg
    yaw = math.radians(90.0 - heading_deg)
    return (0.0, 0.0, math.sin(yaw / 2), math.cos(yaw / 2))   # x, y, z, w


class RouteReplayerNode(Node):

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

        self.declare_parameter('save_dir',          '/data/routes')
        self.declare_parameter('route_name',        'route')
        self.declare_parameter('replay_spacing_m',  3.0)   # subsample spacing
        self.declare_parameter('goal_tolerance_xy', 2.0)   # GPS ±2.5m CEP
        self.declare_parameter('replay_speed_factor', 1.0) # 1.0 = same speed as recorded
        self.declare_parameter('map_frame',         'map')
        self.declare_parameter('nav_timeout_s',     600.0) # 10 min

        self._route:    Optional[dict]         = None   # loaded route metadata
        self._fix:      Optional[NavSatFix]    = None
        self._odom:     Optional[Odometry]     = None
        self._paused    = False
        self._active    = False
        self._nav_goal_handle = None

        self._cb_group = ReentrantCallbackGroup()

        # ── Subscriptions ────────────────────────────────────────────────────
        self.create_subscription(NavSatFix, '/gps/fix',         self._fix_cb,   10)
        self.create_subscription(Odometry,  '/odom',            self._odom_cb,  10)
        self.create_subscription(Bool,      '/route/pause_cmd', self._pause_cb, 10)

        # ── Publishers ───────────────────────────────────────────────────────
        self._status_pub   = self.create_publisher(String, '/route/replay_status', 10)
        self._progress_pub = self.create_publisher(String, '/route/progress',      10)

        # ── Services ─────────────────────────────────────────────────────────
        self.create_service(Trigger, '/route/load',          self._load_cb,
                            callback_group=self._cb_group)
        self.create_service(Trigger, '/route/start_replay',  self._start_cb,
                            callback_group=self._cb_group)
        self.create_service(Trigger, '/route/pause',         self._pause_svc_cb,
                            callback_group=self._cb_group)
        self.create_service(Trigger, '/route/stop',          self._stop_cb,
                            callback_group=self._cb_group)

        # ── Nav2 action client ───────────────────────────────────────────────
        self._nav_client = ActionClient(
            self, NavigateThroughPoses, '/navigate_through_poses',
            callback_group=self._cb_group,
        )

        # ── Progress timer ────────────────────────────────────────────────────
        self.create_timer(2.0, self._publish_progress)

        self._pub_status('idle')
        self.get_logger().info('route_replayer: ready. Call /route/load then /route/start_replay.')

    # ── Sensor callbacks ──────────────────────────────────────────────────────

    def _fix_cb(self, msg: NavSatFix) -> None:
        self._fix = msg

    def _odom_cb(self, msg: Odometry) -> None:
        self._odom = msg

    def _pause_cb(self, msg: Bool) -> None:
        if msg.data and not self._paused:
            self._pause_replay()
        elif not msg.data and self._paused:
            self._resume_replay()

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

    def _load_cb(self, _req, response):
        name    = self.get_parameter('route_name').value
        dir_    = self.get_parameter('save_dir').value
        path    = os.path.join(dir_, f'{name}.jsonl')

        if not os.path.exists(path):
            response.success = False
            response.message = f'Route file not found: {path}'
            return response

        with open(path) as f:
            self._route = json.load(f)

        n    = self._route.get('waypoint_count', len(self._route.get('waypoints', [])))
        dist = self._route.get('total_distance_m', 0.0)
        dur  = self._route.get('duration_s', 0.0)

        self.get_logger().info(
            f'route_replayer: loaded "{name}" — {n} waypoints, '
            f'{dist:.0f} m, {dur:.0f} s'
        )
        self._pub_status(f'loaded:{name}:{n}_waypoints:{dist:.0f}m')
        response.success = True
        response.message = f'Loaded: {name} ({n} waypoints, {dist:.0f} m, {dur:.0f} s)'
        return response

    def _start_cb(self, _req, response):
        if self._route is None:
            response.success = False
            response.message = 'No route loaded — call /route/load first'
            return response

        if self._active:
            response.success = False
            response.message = 'Replay already active'
            return response

        if self._fix is None:
            response.success = False
            response.message = 'No GPS fix'
            return response

        if self._fix.status.status < NavSatStatus.STATUS_FIX:
            response.success = False
            response.message = f'GPS quality too low: status={self._fix.status.status}'
            return response

        # Kick off async nav
        self.create_timer(0.1, self._send_nav_goal, callback_group=self._cb_group)
        response.success = True
        name = self._route.get('name', '?')
        response.message = f'Starting replay of "{name}"'
        return response

    def _pause_svc_cb(self, _req, response):
        if self._paused:
            self._resume_replay()
            response.message = 'Replay resumed'
        else:
            self._pause_replay()
            response.message = 'Replay paused'
        response.success = True
        return response

    def _stop_cb(self, _req, response):
        self._cancel_nav('stop requested')
        response.success = True
        response.message = 'Replay stopped'
        return response

    # ── Navigation ────────────────────────────────────────────────────────────

    def _send_nav_goal(self) -> None:
        """One-shot timer: convert route waypoints → Nav2 goal."""
        self.destroy_timer(list(self.timers)[-1])

        if not self._nav_client.wait_for_server(timeout_sec=5.0):
            self.get_logger().error('Nav2 not available')
            self._pub_status('error:nav2_unavailable')
            return

        waypoints = self._route.get('waypoints', [])
        if not waypoints:
            self.get_logger().error('Route has no waypoints')
            self._pub_status('error:empty_route')
            return

        # Use current GPS fix as datum for ENU origin
        lat0 = self._fix.latitude
        lon0 = self._fix.longitude

        # Subsample by spacing
        spacing = self.get_parameter('replay_spacing_m').value
        sampled = self._subsample_waypoints(waypoints, spacing)

        frame = self.get_parameter('map_frame').value
        now   = self.get_clock().now().to_msg()
        goal  = NavigateThroughPoses.Goal()

        for wp in sampled:
            x, y = _latlon_to_enu(wp['lat'], wp['lon'], lat0, lon0)
            qx, qy, qz, qw = _heading_to_quat(wp.get('heading_deg', 0.0))

            ps = PoseStamped()
            ps.header.stamp    = now
            ps.header.frame_id = frame
            ps.pose.position.x = x
            ps.pose.position.y = y
            ps.pose.orientation.x = qx
            ps.pose.orientation.y = qy
            ps.pose.orientation.z = qz
            ps.pose.orientation.w = qw
            goal.poses.append(ps)

        name = self._route.get('name', '?')
        self.get_logger().info(
            f'route_replayer: sending {len(goal.poses)} waypoints to Nav2 '
            f'(from {len(waypoints)} recorded, datum={lat0:.6f},{lon0:.6f})'
        )
        self._pub_status(f'replaying:{name}:{len(goal.poses)}_waypoints')
        self._active = True
        self._paused = False

        future = self._nav_client.send_goal_async(
            goal,
            feedback_callback=self._nav_feedback_cb,
        )
        future.add_done_callback(self._nav_goal_response_cb)

    def _nav_goal_response_cb(self, future) -> None:
        self._nav_goal_handle = future.result()
        if not self._nav_goal_handle.accepted:
            self.get_logger().error('Nav2 rejected replay goal')
            self._active = False
            self._pub_status('error:goal_rejected')
            return
        result_future = self._nav_goal_handle.get_result_async()
        result_future.add_done_callback(self._nav_result_cb)

    def _nav_feedback_cb(self, feedback) -> None:
        fb = feedback.feedback
        remaining = getattr(fb, 'number_of_poses_remaining', '?')
        dist = getattr(fb, 'distance_remaining', None)
        if dist is not None:
            self.get_logger().info(
                f'replay: {remaining} waypoints remaining, {dist:.1f} m to go',
                throttle_duration_sec=10.0,
            )

    def _nav_result_cb(self, future) -> None:
        self._active = False
        result = future.result()
        name   = self._route.get('name', '?') if self._route else '?'
        if result.status == 4:
            self.get_logger().info(f'Replay of "{name}" completed successfully')
            self._pub_status(f'completed:{name}')
        else:
            self.get_logger().warn(f'Replay ended with status {result.status}')
            self._pub_status(f'failed:{name}:status={result.status}')

    def _pause_replay(self) -> None:
        if self._nav_goal_handle and self._active:
            self._nav_goal_handle.cancel_goal_async()
        self._paused = True
        self._active = False
        self._pub_status('paused')
        self.get_logger().info('route_replayer: paused')

    def _resume_replay(self) -> None:
        """Re-issue the nav goal from current position (simple re-start)."""
        self._paused = False
        if self._route:
            self.create_timer(0.1, self._send_nav_goal, callback_group=self._cb_group)
            self._pub_status('resuming')
            self.get_logger().info('route_replayer: resuming')

    def _cancel_nav(self, reason: str) -> None:
        if self._nav_goal_handle and self._active:
            self._nav_goal_handle.cancel_goal_async()
        self._active = False
        self._paused = False
        self._pub_status(f'stopped:{reason}')
        self.get_logger().info(f'route_replayer: stopped — {reason}')

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _subsample_waypoints(self, waypoints: list, spacing_m: float) -> list:
        """Keep waypoints at least spacing_m apart. Always keep first + last."""
        if not waypoints:
            return []
        out = [waypoints[0]]
        for wp in waypoints[1:-1]:
            prev = out[-1]
            d = _haversine(wp['lat'], wp['lon'], prev['lat'], prev['lon'])
            if d >= spacing_m:
                out.append(wp)
        if len(waypoints) > 1:
            out.append(waypoints[-1])
        return out

    def _publish_progress(self) -> None:
        if not self._active or self._route is None:
            return
        name = self._route.get('name', '?')
        n    = self._route.get('waypoint_count', '?')
        dist = self._route.get('total_distance_m', 0.0)
        msg = String()
        msg.data = f'replaying:{name} total={n}wps dist={dist:.0f}m'
        self._progress_pub.publish(msg)

    def _pub_status(self, s: str) -> None:
        msg = String()
        msg.data = s
        self._status_pub.publish(msg)


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


if __name__ == '__main__':
    main()
