"""
route_recorder_node — Records GPS + odometry path during follow-me rides.

"Ride once, replay forever." — During a UWB follow-me session, this node
samples the robot's GPS position, heading, and speed into a timestamped
JSON-Lines file. The result is a portable route file that route_replayer_node
can load to autonomously repeat the ride.

Recording strategy:
  - Samples at up to 1 Hz (configurable)
  - Skips waypoints within waypoint_spacing_m of the previous one
  - Each waypoint: {t, lat, lon, heading_deg, speed_mps, x, y, source}
  - source: 'gps' | 'odom' (used when GPS quality drops)

JSON-Lines format (one JSON object per line, .jsonl):
  {"version": 1, "name": "...", "recorded_at": "...", "waypoints": [...]}
  Written as a single-line header on line 1 (metadata), then appended to.

File location: /data/routes/<name>.jsonl

Topics subscribed:
  /gps/fix          sensor_msgs/NavSatFix   — GPS position (SIM7600X)
  /odom             nav_msgs/Odometry       — wheel encoder odometry
  /uwb/target       geometry_msgs/PoseStamped — UWB target position (for context)
  /imu/data         sensor_msgs/Imu         — heading from IMU

Topics published:
  /route/status     std_msgs/String  — 'idle'|'recording:<name>'|'saved:<name>'

Services:
  /route/start_recording  std_srvs/Trigger  — begin recording (name via param)
  /route/stop_recording   std_srvs/Trigger  — stop and buffer (don't save yet)
  /route/save             std_srvs/Trigger  — flush buffer to disk
  /route/discard          std_srvs/Trigger  — discard unsaved recording

Parameters:
  save_dir            (default /data/routes)
  waypoint_spacing_m  (default 2.0) — skip closer waypoints
  sample_rate_hz      (default 1.0) — max recording rate
  route_name          (default 'route') — set before start_recording
  min_fix_status      (default 0)  — reject STATUS_NO_FIX
"""

import json
import math
import os
import time
from datetime import datetime, timezone
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.callback_groups import ReentrantCallbackGroup

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


def _haversine(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """Great-circle distance in metres."""
    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 _quat_to_yaw_deg(q) -> float:
    """Convert quaternion to yaw in degrees (-180..180)."""
    siny = 2.0 * (q.w * q.z + q.x * q.y)
    cosy = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
    return math.degrees(math.atan2(siny, cosy))


class RouteRecorderNode(Node):

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

        self.declare_parameter('save_dir',           '/data/routes')
        self.declare_parameter('waypoint_spacing_m', 2.0)
        self.declare_parameter('sample_rate_hz',     1.0)
        self.declare_parameter('route_name',         'route')
        self.declare_parameter('min_fix_status',     NavSatStatus.STATUS_FIX)

        self._recording = False
        self._waypoints: list[dict] = []
        self._last_wp:   Optional[dict] = None
        self._last_sample_t = 0.0

        # Latest sensor state
        self._fix:    Optional[NavSatFix]     = None
        self._odom:   Optional[Odometry]      = None
        self._imu:    Optional[Imu]           = None
        self._uwb:    Optional[PoseStamped]   = 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(Imu,          '/imu/data',    self._imu_cb,  10)
        self.create_subscription(PoseStamped,  '/uwb/target',  self._uwb_cb,  10)

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

        # ── Services ─────────────────────────────────────────────────────────
        self.create_service(Trigger, '/route/start_recording', self._start_cb,
                            callback_group=self._cb_group)
        self.create_service(Trigger, '/route/stop_recording',  self._stop_cb,
                            callback_group=self._cb_group)
        self.create_service(Trigger, '/route/save',            self._save_cb,
                            callback_group=self._cb_group)
        self.create_service(Trigger, '/route/discard',         self._discard_cb,
                            callback_group=self._cb_group)

        # ── Sample timer ─────────────────────────────────────────────────────
        rate = self.get_parameter('sample_rate_hz').value
        self.create_timer(1.0 / rate, self._sample)

        self._pub_status('idle')
        self.get_logger().info('route_recorder: ready. Call /route/start_recording.')

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

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

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

    def _imu_cb(self, msg: Imu) -> None:
        self._imu = msg

    def _uwb_cb(self, msg: PoseStamped) -> None:
        self._uwb = msg

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

    def _start_cb(self, _req, response):
        if self._recording:
            response.success = False
            response.message = 'Already recording'
            return response

        name = self.get_parameter('route_name').value
        self._waypoints = []
        self._last_wp = None
        self._recording = True
        self.get_logger().info(f'route_recorder: started recording "{name}"')
        self._pub_status(f'recording:{name}')
        response.success = True
        response.message = f'Recording started: {name}'
        return response

    def _stop_cb(self, _req, response):
        if not self._recording:
            response.success = False
            response.message = 'Not recording'
            return response
        self._recording = False
        n = len(self._waypoints)
        self.get_logger().info(f'route_recorder: stopped. {n} waypoints buffered.')
        self._pub_status(f'stopped:{n}_waypoints')
        response.success = True
        response.message = f'Recording stopped. {n} waypoints buffered. Call /route/save.'
        return response

    def _save_cb(self, _req, response):
        if not self._waypoints:
            response.success = False
            response.message = 'No waypoints to save'
            return response

        name = self.get_parameter('route_name').value
        save_dir = self.get_parameter('save_dir').value
        os.makedirs(save_dir, exist_ok=True)
        path = os.path.join(save_dir, f'{name}.jsonl')

        total_dist = self._total_distance_m()
        meta = {
            'version':      1,
            'name':         name,
            'recorded_at':  datetime.now(timezone.utc).isoformat(),
            'waypoint_count': len(self._waypoints),
            'total_distance_m': round(total_dist, 1),
            'duration_s':   self._duration_s(),
            'waypoints':    self._waypoints,
        }
        with open(path, 'w') as f:
            json.dump(meta, f)
            f.write('\n')

        self.get_logger().info(
            f'route_recorder: saved {len(self._waypoints)} waypoints → {path} '
            f'({total_dist:.0f} m)'
        )
        self._pub_status(f'saved:{name}')
        self._waypoints = []
        response.success = True
        response.message = f'Saved: {path} ({len(meta["waypoints"])} waypoints, {total_dist:.0f} m)'
        return response

    def _discard_cb(self, _req, response):
        n = len(self._waypoints)
        self._waypoints = []
        self._recording = False
        self._last_wp = None
        self._pub_status('idle')
        response.success = True
        response.message = f'Discarded {n} waypoints'
        return response

    # ── Sample timer ──────────────────────────────────────────────────────────

    def _sample(self) -> None:
        if not self._recording:
            return
        if self._fix is None:
            return

        # GPS quality check
        min_status = self.get_parameter('min_fix_status').value
        if self._fix.status.status < min_status:
            return

        lat = self._fix.latitude
        lon = self._fix.longitude

        # Skip if too close to last waypoint
        spacing = self.get_parameter('waypoint_spacing_m').value
        if self._last_wp is not None:
            d = _haversine(lat, lon, self._last_wp['lat'], self._last_wp['lon'])
            if d < spacing:
                return

        # Heading: prefer IMU yaw, fall back to GPS bearing
        heading_deg = 0.0
        if self._imu is not None:
            heading_deg = _quat_to_yaw_deg(self._imu.orientation)
        elif self._last_wp is not None:
            dlat = lat - self._last_wp['lat']
            dlon = lon - self._last_wp['lon']
            heading_deg = math.degrees(math.atan2(dlon, dlat))

        # Speed from odometry
        speed_mps = 0.0
        odom_x, odom_y = 0.0, 0.0
        if self._odom is not None:
            vx = self._odom.twist.twist.linear.x
            vy = self._odom.twist.twist.linear.y
            speed_mps = math.hypot(vx, vy)
            odom_x = self._odom.pose.pose.position.x
            odom_y = self._odom.pose.pose.position.y

        wp = {
            't':           self.get_clock().now().nanoseconds / 1e9,
            'lat':         lat,
            'lon':         lon,
            'heading_deg': round(heading_deg, 1),
            'speed_mps':   round(speed_mps, 2),
            'odom_x':      round(odom_x, 3),
            'odom_y':      round(odom_y, 3),
            'gps_status':  self._fix.status.status,
        }

        self._waypoints.append(wp)
        self._last_wp = wp

        name = self.get_parameter('route_name').value
        self.get_logger().debug(
            f'route_recorder: wp #{len(self._waypoints)} '
            f'lat={lat:.6f} lon={lon:.6f} hdg={heading_deg:.0f}° spd={speed_mps:.1f}m/s'
        )
        self._pub_status(f'recording:{name}:{len(self._waypoints)}_wps')

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

    def _total_distance_m(self) -> float:
        if len(self._waypoints) < 2:
            return 0.0
        total = 0.0
        for i in range(1, len(self._waypoints)):
            a, b = self._waypoints[i - 1], self._waypoints[i]
            total += _haversine(a['lat'], a['lon'], b['lat'], b['lon'])
        return total

    def _duration_s(self) -> float:
        if len(self._waypoints) < 2:
            return 0.0
        return self._waypoints[-1]['t'] - self._waypoints[0]['t']

    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 = RouteRecorderNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == '__main__':
    main()
