#!/usr/bin/env python3
"""底盘轮式里程计节点 (chassis_odometry)

直连底盘 MCU (lzucar 板) 串口 `/dev/ttylzucar`，**只读**解析 32 字节里程计数据包，
发布 nav_msgs/Odometry (/odom) 与 TF (odom -> base_footprint)。

数据包协议（移植自 lzu_robot/src/cgbot/cgbot/seriallzucar.py，已上真机验证的同款底盘）：
  - 帧尾标志 3 字节: 'L', 'Z', 'U'
  - 整包 32 字节，以 'LZU' 结尾对齐
  - packet[0]      = 校验和 = sum(packet[1:30]) % 256
  - packet[1:29]   = <fffffff> = x, y, yaw, posFR, posFL, posBL, posBR
                     (x/y 单位 mm，需 /1000 转米; yaw 单位弧度)
  - packet[29]     = 模式字节, 'D' 表示航位推算有效
  - packet[29:32]  = 帧尾 'L','Z','U'（对齐用）

本节点不发送任何命令；既有底盘速度控制 (UDP->ESP32->UART) 不受影响，互不抢占串口。
若后续需把底盘命令也迁到直连串口，应在此节点内合并发送逻辑（参考 seriallzucar 的 pack_floats_send）。
"""

import math
import struct
import time
from threading import Lock, Thread

import rclpy
import serial
from geometry_msgs.msg import Quaternion, TransformStamped
from nav_msgs.msg import Odometry
from rclpy.node import Node
from tf2_ros import TransformBroadcaster

PACKET_SIZE = 32
MODE_DEAD_RECKON = ord('D')


def yaw_to_quaternion(yaw: float) -> Quaternion:
    """绕 Z 轴偏航角 -> 四元数。"""
    q = Quaternion()
    q.x = 0.0
    q.y = 0.0
    q.z = math.sin(yaw / 2.0)
    q.w = math.cos(yaw / 2.0)
    return q


class ChassisOdometry(Node):
    def __init__(self) -> None:
        super().__init__('chassis_odometry')

        # ---- 参数 ----
        self.port = self.declare_parameter('port', '/dev/ttylzucar').value
        self.baudrate = int(self.declare_parameter('baudrate', 115200).value)
        self.odom_frame_id = self.declare_parameter('odom_frame_id', 'odom').value
        self.base_frame_id = self.declare_parameter('base_frame_id', 'base_footprint').value
        self.odom_topic = self.declare_parameter('odom_topic', 'odom').value
        self.publish_tf = bool(self.declare_parameter('publish_tf', True).value)
        # 仅当模式字节为 'D'(航位推算有效) 时才发布；置 False 可在调试期放行所有包
        self.require_mode_d = bool(self.declare_parameter('require_mode_d', True).value)
        # 串口打开后下发的使能命令，让底盘板持续上报里程计；空字符串=不发。
        # lzu_robot/seriallzucar.open() 与 move_try/serial_bridge 都在打开后发 "UP0LZU"，
        # 缺它则底盘上电后只发几帧就停，odom TF 变陈旧、gmapping 丢弃所有 scan。
        self.enable_cmd = str(self.declare_parameter('enable_cmd', 'UP0LZU').value)
        # >0 时按该周期(秒)重发使能命令做保活；0=仅打开时发一次(与参考实现一致)
        self.enable_cmd_period = float(self.declare_parameter('enable_cmd_period', 0.0).value)

        # ---- 发布器 ----
        self.odom_pub = self.create_publisher(Odometry, self.odom_topic, 10)
        self.tf_broadcaster = TransformBroadcaster(self) if self.publish_tf else None

        # ---- 串口接收环形缓冲 ----
        self.BUFFER_SIZE = 1024
        self.buffer = bytearray(self.BUFFER_SIZE)
        self.write_index = 0
        self.buffer_lock = Lock()

        # ---- 上一帧状态（用于差分求速度）----
        self.prev_x = None
        self.prev_y = None
        self.prev_yaw = None
        self.prev_t = None

        # ---- 串口 + 接收线程 ----
        self.ser = None
        self._running = True
        self._reader = Thread(target=self._receiver_thread, daemon=True)
        self._reader.start()

        # 可选：周期性重发使能命令做保活
        if self.enable_cmd and self.enable_cmd_period > 0.0:
            self.create_timer(self.enable_cmd_period, self._send_enable_cmd)

        self.get_logger().info(
            f'chassis_odometry 启动: port={self.port}@{self.baudrate} '
            f'odom_frame={self.odom_frame_id} base_frame={self.base_frame_id} '
            f'publish_tf={self.publish_tf}'
        )

    # ------------------------------------------------------------------ #
    # 串口接收 / 重连
    # ------------------------------------------------------------------ #
    def _open_serial(self) -> bool:
        try:
            self.ser = serial.Serial(self.port, self.baudrate, timeout=1)
            self.get_logger().info(f'串口已连接 {self.port}')
            self._send_enable_cmd()
            return True
        except Exception as exc:  # noqa: BLE001
            self.get_logger().warning(f'串口打开失败 {self.port}: {exc}')
            self.ser = None
            return False

    def _send_enable_cmd(self) -> None:
        """下发里程计使能命令(默认 UP0LZU)；空命令或串口未开则跳过。"""
        if not self.enable_cmd or self.ser is None or not self.ser.is_open:
            return
        try:
            self.ser.write(self.enable_cmd.encode('utf-8'))
            self.ser.flush()
            self.get_logger().info(f'已下发里程计使能命令: {self.enable_cmd}')
        except Exception as exc:  # noqa: BLE001
            self.get_logger().warning(f'使能命令发送失败: {exc}')

    def _receiver_thread(self) -> None:
        while self._running:
            try:
                if self.ser is None or not self.ser.is_open:
                    if not self._open_serial():
                        time.sleep(1.0)
                    continue
                waiting = self.ser.in_waiting
                if waiting > 0:
                    self._write_to_buffer(self.ser.read(waiting))
                else:
                    time.sleep(0.005)
            except serial.SerialException as exc:
                self.get_logger().warning(f'串口中断，准备重连: {exc}')
                self._close_serial()
                time.sleep(1.0)
            except Exception as exc:  # noqa: BLE001
                self.get_logger().error(f'接收线程异常: {exc}')
                self._close_serial()
                time.sleep(1.0)

    def _close_serial(self) -> None:
        if self.ser is not None:
            try:
                self.ser.close()
            except Exception:  # noqa: BLE001
                pass
        self.ser = None

    # ------------------------------------------------------------------ #
    # 帧解析（逐字节移植自 seriallzucar，保持已验证行为）
    # ------------------------------------------------------------------ #
    def _read_byte(self, index: int) -> int:
        return self.buffer[(index + self.BUFFER_SIZE) % self.BUFFER_SIZE]

    def _read_from_buffer(self, start_index: int, size: int) -> bytearray:
        data = bytearray()
        for i in range(size):
            data.append(self.buffer[(start_index + i) % self.BUFFER_SIZE])
        return data

    def _write_to_buffer(self, data: bytes) -> None:
        with self.buffer_lock:
            for byte in data:
                self.buffer[self.write_index] = byte
                # 检测帧尾 'L','Z','U'
                if (byte == ord('U')
                        and self._read_byte(self.write_index - 1) == ord('Z')
                        and self._read_byte(self.write_index - 2) == ord('L')):
                    start_index = (self.write_index - PACKET_SIZE + self.BUFFER_SIZE) % self.BUFFER_SIZE
                    packet = self._read_from_buffer(start_index, PACKET_SIZE)
                    self._process_packet(packet)
                self.write_index = (self.write_index + 1) % self.BUFFER_SIZE

    def _process_packet(self, packet: bytearray) -> None:
        # 校验和
        if (sum(packet[1:30]) % 256) != packet[0]:
            return
        try:
            x, y, yaw, _fr, _fl, _bl, _br = struct.unpack('<fffffff', packet[1:29])
            mode = packet[29]
        except struct.error:
            return

        if self.require_mode_d and mode != MODE_DEAD_RECKON:
            return

        x /= 1000.0  # mm -> m
        y /= 1000.0
        self._publish(x, y, yaw)

    # ------------------------------------------------------------------ #
    # 发布 Odometry + TF
    # ------------------------------------------------------------------ #
    def _publish(self, x: float, y: float, yaw: float) -> None:
        now = self.get_clock().now()
        stamp = now.to_msg()
        quat = yaw_to_quaternion(yaw)

        # 差分求速度（body 系）；首帧或 dt 过小时置零
        vx = vy = wz = 0.0
        t = now.nanoseconds * 1e-9
        if self.prev_t is not None:
            dt = t - self.prev_t
            if dt > 1e-3:
                dx = x - self.prev_x
                dy = y - self.prev_y
                dyaw = math.atan2(math.sin(yaw - self.prev_yaw),
                                  math.cos(yaw - self.prev_yaw))
                # 世界系位移旋到 body 系
                vx = (dx * math.cos(yaw) + dy * math.sin(yaw)) / dt
                vy = (-dx * math.sin(yaw) + dy * math.cos(yaw)) / dt
                wz = dyaw / dt
        self.prev_x, self.prev_y, self.prev_yaw, self.prev_t = x, y, yaw, t

        odom = Odometry()
        odom.header.stamp = stamp
        odom.header.frame_id = self.odom_frame_id
        odom.child_frame_id = self.base_frame_id
        odom.pose.pose.position.x = x
        odom.pose.pose.position.y = y
        odom.pose.pose.position.z = 0.0
        odom.pose.pose.orientation = quat
        odom.twist.twist.linear.x = vx
        odom.twist.twist.linear.y = vy
        odom.twist.twist.angular.z = wz
        # 平面机器人协方差: x,y,yaw 给小值, z/roll/pitch 不可观给大值
        odom.pose.covariance[0] = 0.01    # x
        odom.pose.covariance[7] = 0.01    # y
        odom.pose.covariance[14] = 1e6    # z
        odom.pose.covariance[21] = 1e6    # roll
        odom.pose.covariance[28] = 1e6    # pitch
        odom.pose.covariance[35] = 0.02   # yaw
        odom.twist.covariance[0] = 0.01
        odom.twist.covariance[7] = 0.01
        odom.twist.covariance[35] = 0.02
        self.odom_pub.publish(odom)

        if self.tf_broadcaster is not None:
            tf = TransformStamped()
            tf.header.stamp = stamp
            tf.header.frame_id = self.odom_frame_id
            tf.child_frame_id = self.base_frame_id
            tf.transform.translation.x = x
            tf.transform.translation.y = y
            tf.transform.translation.z = 0.0
            tf.transform.rotation = quat
            self.tf_broadcaster.sendTransform(tf)

    def destroy_node(self) -> bool:
        self._running = False
        self._close_serial()
        return super().destroy_node()


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


if __name__ == '__main__':
    main()
