#!/usr/bin/env python3
"""
ntrip_client_node.py — NTRIP client for Here4 RTK corrections
Issue: https://gitea.vayrette.com/seb/saltylab-firmware/issues/725

Connects to an NTRIP caster (default: rtk2go.com:2101), streams RTCM3
correction data, and forwards it to the Here4 GPS via DroneCAN
(uavcan.equipment.gnss.RTCMStream) on the CAN bus.

Publishes:
  /ntrip/status  std_msgs/String  — CONNECTED / DISCONNECTED / ERROR:<reason>

Parameters:
  ntrip_caster    (str)   rtk2go.com
  ntrip_port      (int)   2101
  ntrip_mount     (str)   ''       — required, e.g. 'RTCM3_GENERIC'
  ntrip_user      (str)   ''       — leave empty for anonymous casters
  ntrip_password  (str)   ''
  can_interface   (str)   can0
  can_bitrate     (int)   1000000
  can_node_id     (int)   126      — DroneCAN local node ID for NTRIP node
  reconnect_delay (float) 5.0     — seconds between reconnect attempts
"""

import base64
import socket
import threading
import time

import rclpy
from rclpy.node import Node
from std_msgs.msg import String

try:
    import dronecan
except ImportError:
    dronecan = None

_NTRIP_AGENT  = 'saltybot-ntrip/1.0'
_RTCM_CHUNK   = 512   # bytes per DroneCAN RTCMStream message (max 128 per frame)
_RECV_BUF     = 4096


class NtripClientNode(Node):
    def __init__(self) -> None:
        super().__init__('ntrip_client')

        self.declare_parameter('ntrip_caster',    'rtk2go.com')
        self.declare_parameter('ntrip_port',      2101)
        self.declare_parameter('ntrip_mount',     '')
        self.declare_parameter('ntrip_user',      '')
        self.declare_parameter('ntrip_password',  '')
        self.declare_parameter('can_interface',   'can0')
        self.declare_parameter('can_bitrate',     1000000)
        self.declare_parameter('can_node_id',     126)
        self.declare_parameter('reconnect_delay', 5.0)

        self._caster   = self.get_parameter('ntrip_caster').value
        self._port     = self.get_parameter('ntrip_port').value
        self._mount    = self.get_parameter('ntrip_mount').value
        self._user     = self.get_parameter('ntrip_user').value
        self._password = self.get_parameter('ntrip_password').value
        self._iface    = self.get_parameter('can_interface').value
        self._bitrate  = self.get_parameter('can_bitrate').value
        self._node_id  = self.get_parameter('can_node_id').value
        self._reconnect_delay = self.get_parameter('reconnect_delay').value

        self._status_pub = self.create_publisher(String, '/ntrip/status', 10)

        if dronecan is None:
            self.get_logger().error(
                'python-dronecan not installed. Run: pip install python-dronecan'
            )
            self._publish_status('ERROR:dronecan_not_installed')
            return

        if not self._mount:
            self.get_logger().error(
                'ntrip_mount parameter is required (e.g. RTCM3_GENERIC)'
            )
            self._publish_status('ERROR:no_mount_point')
            return

        self._dc_node = dronecan.make_node(
            self._iface,
            node_id=self._node_id,
            bitrate=self._bitrate,
        )

        self._stop_event = threading.Event()
        self._thread = threading.Thread(
            target=self._ntrip_loop, daemon=True
        )
        self._thread.start()

        self.get_logger().info(
            f'NTRIP client started — '
            f'{self._caster}:{self._port}/{self._mount}'
        )

    # ── NTRIP loop (reconnects on any error) ──────────────────────────────────

    def _ntrip_loop(self) -> None:
        while not self._stop_event.is_set() and rclpy.ok():
            try:
                self._connect_and_stream()
            except Exception as exc:
                self.get_logger().warn(f'NTRIP disconnected: {exc}')
                self._publish_status(f'DISCONNECTED')
            if not self._stop_event.is_set() and rclpy.ok():
                self.get_logger().info(
                    f'Reconnecting in {self._reconnect_delay}s…'
                )
                time.sleep(self._reconnect_delay)

    def _connect_and_stream(self) -> None:
        sock = socket.create_connection(
            (self._caster, self._port), timeout=10.0
        )
        sock.settimeout(30.0)

        # Send NTRIP HTTP/1.1 GET request
        request = self._build_request()
        sock.sendall(request.encode('ascii'))

        # Read response header
        header = b''
        while b'\r\n\r\n' not in header:
            chunk = sock.recv(256)
            if not chunk:
                raise ConnectionError('Connection closed during header read')
            header += chunk

        header_str = header.split(b'\r\n\r\n')[0].decode('ascii', errors='replace')
        if 'ICY 200 OK' not in header_str and '200 OK' not in header_str:
            raise ConnectionError(f'NTRIP rejected: {header_str[:120]}')

        self.get_logger().info(
            f'NTRIP connected to {self._caster}:{self._port}/{self._mount}'
        )
        self._publish_status('CONNECTED')

        # Any trailing bytes after header are RTCM data
        leftover = header.split(b'\r\n\r\n', 1)[1] if b'\r\n\r\n' in header else b''
        if leftover:
            self._forward_rtcm(leftover)

        while not self._stop_event.is_set() and rclpy.ok():
            data = sock.recv(_RECV_BUF)
            if not data:
                raise ConnectionError('NTRIP stream closed by server')
            self._forward_rtcm(data)

    def _build_request(self) -> str:
        lines = [
            f'GET /{self._mount} HTTP/1.1',
            f'Host: {self._caster}:{self._port}',
            f'User-Agent: {_NTRIP_AGENT}',
            'Ntrip-Version: Ntrip/2.0',
            'Connection: close',
        ]
        if self._user:
            creds = base64.b64encode(
                f'{self._user}:{self._password}'.encode()
            ).decode()
            lines.append(f'Authorization: Basic {creds}')
        lines += ['', '']
        return '\r\n'.join(lines)

    # ── DroneCAN RTCMStream forwarding ────────────────────────────────────────

    def _forward_rtcm(self, data: bytes) -> None:
        """Chunk RTCM data into DroneCAN RTCMStream messages."""
        for i in range(0, len(data), _RTCM_CHUNK):
            chunk = data[i:i + _RTCM_CHUNK]
            try:
                msg = dronecan.uavcan.equipment.gnss.RTCMStream()
                msg.data = list(chunk)
                self._dc_node.broadcast(msg)
                self._dc_node.spin(timeout=0.0)
            except Exception as exc:
                self.get_logger().warn(f'DroneCAN send error: {exc}')

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

    def destroy_node(self) -> None:
        self._stop_event.set()
        super().destroy_node()


def main(args=None) -> None:
    rclpy.init(args=args)
    node = NtripClientNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
