#!/usr/bin/env python3
"""
vesc_odometry_bridge.py — Differential drive odometry from dual VESC CAN motors (Issue #646).

Subscribes to per-motor VESC telemetry topics (configurable, defaulting to
/vesc/left/state and /vesc/right/state as published by the telemetry node)
for CAN IDs 56 (left) and 68 (right),
applies differential drive kinematics via DiffDriveOdometry, and publishes:

  /odom                  nav_msgs/Odometry   — consumed by Nav2 / slam_toolbox
  /saltybot/wheel_odom   nav_msgs/Odometry   — consumed by saltybot_pose_fusion EKF
  TF odom → base_link

Input topics (std_msgs/String, JSON):
  /vesc/left/state   {"rpm": <eRPM>, "voltage_v": ..., "current_a": ..., ...}
  /vesc/right/state  same schema

Topic names are configurable via left_state_topic / right_state_topic params.
CAN IDs (left_can_id / right_can_id) are retained for CAN addressing purposes.

Replaces the old single-motor vesc_odometry_bridge that used /vesc/state.
"""

import json
import math

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

from .diff_drive_odom import DiffDriveOdometry


def _yaw_to_quaternion(yaw: float) -> Quaternion:
    """Convert yaw angle (rad) to geometry_msgs/Quaternion (rotation about Z)."""
    half = yaw * 0.5
    q = Quaternion()
    q.w = math.cos(half)
    q.x = 0.0
    q.y = 0.0
    q.z = math.sin(half)
    return q


class VESCCANOdometryNode(Node):
    """
    Differential drive odometry from dual VESC CAN motor telemetry.

    Integrates DiffDriveOdometry at a fixed rate using the latest RPM values
    received from each motor's CAN telemetry topic.
    """

    def __init__(self) -> None:
        super().__init__("vesc_can_odometry")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("left_can_id",       56)      # CAN ID for left motor
        self.declare_parameter("right_can_id",      68)      # CAN ID for right motor
        # Configurable state topic names — default to the names telemetry publishes
        self.declare_parameter("left_state_topic",  "/vesc/left/state")
        self.declare_parameter("right_state_topic", "/vesc/right/state")
        self.declare_parameter("wheel_radius",       0.10)   # metres
        self.declare_parameter("wheel_separation",   0.35)   # metres (track width)
        self.declare_parameter("motor_poles",        7)      # BLDC pole pairs
        self.declare_parameter("invert_right",       True)   # flip right eRPM sign
        self.declare_parameter("odom_frame_id",      "odom")
        self.declare_parameter("base_frame_id",      "base_link")
        self.declare_parameter("publish_tf",         True)
        self.declare_parameter("publish_wheel_odom", True)
        self.declare_parameter("update_frequency",   50.0)   # Hz
        # Covariance tuning (diagonal variances)
        self.declare_parameter("pose_cov_xx",  1e-3)   # x        σ² (m²)
        self.declare_parameter("pose_cov_yy",  1e-3)   # y        σ² (m²)
        self.declare_parameter("pose_cov_aa",  1e-2)   # heading  σ² (rad²)
        self.declare_parameter("twist_cov_vx", 1e-2)   # vx       σ² (m/s)²
        self.declare_parameter("twist_cov_wz", 2e-2)   # ωz       σ² (rad/s)²

        left_id         = self.get_parameter("left_can_id").value
        right_id        = self.get_parameter("right_can_id").value
        left_topic      = self.get_parameter("left_state_topic").value
        right_topic     = self.get_parameter("right_state_topic").value
        freq            = self.get_parameter("update_frequency").value

        self._odom_frame      = self.get_parameter("odom_frame_id").value
        self._base_frame      = self.get_parameter("base_frame_id").value
        self._publish_tf      = self.get_parameter("publish_tf").value
        self._pub_wheel_en    = self.get_parameter("publish_wheel_odom").value

        self._cov_xx = self.get_parameter("pose_cov_xx").value
        self._cov_yy = self.get_parameter("pose_cov_yy").value
        self._cov_aa = self.get_parameter("pose_cov_aa").value
        self._cov_vx = self.get_parameter("twist_cov_vx").value
        self._cov_wz = self.get_parameter("twist_cov_wz").value

        # ── Kinematics engine ─────────────────────────────────────────────────
        self._odom = DiffDriveOdometry(
            wheel_radius     = self.get_parameter("wheel_radius").value,
            wheel_separation = self.get_parameter("wheel_separation").value,
            motor_poles      = self.get_parameter("motor_poles").value,
            invert_right     = self.get_parameter("invert_right").value,
        )

        # Latest eRPM from each motor (updated by subscription callbacks)
        self._erpm_left:  float = 0.0
        self._erpm_right: float = 0.0

        self._last_time: float | None = None

        # ── Publishers ─────────────────────────────────────────────────────────
        self._pub_odom = self.create_publisher(Odometry, "/odom", 10)

        if self._pub_wheel_en:
            self._pub_wheel = self.create_publisher(Odometry, "/saltybot/wheel_odom", 10)
        else:
            self._pub_wheel = None

        if self._publish_tf:
            self._tf_br = TransformBroadcaster(self)
        else:
            self._tf_br = None

        # ── Subscriptions ──────────────────────────────────────────────────────
        self.create_subscription(
            String,
            left_topic,
            self._on_left,
            10,
        )
        self.create_subscription(
            String,
            right_topic,
            self._on_right,
            10,
        )

        # ── Integration timer ──────────────────────────────────────────────────
        self.create_timer(1.0 / freq, self._on_timer)

        self.get_logger().info(
            f"vesc_can_odometry: left=CAN{left_id}({left_topic}) "
            f"right=CAN{right_id}({right_topic}) "
            f"r={self._odom.wheel_radius}m sep={self._odom.wheel_separation}m "
            f"poles={self._odom.motor_poles} invert_right={self._odom.invert_right} "
            f"{self._odom_frame}→{self._base_frame} @ {freq:.0f}Hz"
        )

    # ── RPM callbacks ──────────────────────────────────────────────────────────

    def _parse_erpm(self, msg: String) -> float | None:
        """Extract eRPM from JSON state message; return None on parse error."""
        try:
            return float(json.loads(msg.data).get("rpm", 0))
        except (json.JSONDecodeError, ValueError, TypeError):
            return None

    def _on_left(self, msg: String) -> None:
        erpm = self._parse_erpm(msg)
        if erpm is not None:
            self._erpm_left = erpm

    def _on_right(self, msg: String) -> None:
        erpm = self._parse_erpm(msg)
        if erpm is not None:
            self._erpm_right = erpm

    # ── Integration timer ──────────────────────────────────────────────────────

    def _on_timer(self) -> None:
        """Integrate odometry and publish at fixed rate."""
        now = self.get_clock().now()
        now_s = now.nanoseconds * 1e-9

        if self._last_time is None:
            self._last_time = now_s
            return

        dt = now_s - self._last_time
        self._last_time = now_s

        if dt <= 0.0 or dt > 1.0:   # guard against clock jumps / stalls
            return

        v_lin, v_ang = self._odom.update(self._erpm_left, self._erpm_right, dt)

        stamp = now.to_msg()
        odom_msg = self._build_odom_msg(stamp, v_lin, v_ang)
        self._pub_odom.publish(odom_msg)

        if self._pub_wheel is not None:
            wheel_msg = self._build_odom_msg(stamp, v_lin, v_ang)
            self._pub_wheel.publish(wheel_msg)

        if self._tf_br is not None:
            self._publish_tf(stamp)

    # ── Message builders ───────────────────────────────────────────────────────

    def _build_odom_msg(self, stamp, v_lin: float, v_ang: float) -> Odometry:
        msg = Odometry()
        msg.header.stamp    = stamp
        msg.header.frame_id = self._odom_frame
        msg.child_frame_id  = self._base_frame

        msg.pose.pose.position.x  = self._odom.x
        msg.pose.pose.position.y  = self._odom.y
        msg.pose.pose.position.z  = 0.0
        msg.pose.pose.orientation = _yaw_to_quaternion(self._odom.theta)

        # 6×6 row-major pose covariance [x, y, z, roll, pitch, yaw]
        cov_p = [0.0] * 36
        cov_p[0]  = self._cov_xx   # x,x
        cov_p[7]  = self._cov_yy   # y,y
        cov_p[14] = 1e-4            # z,z  (not estimated)
        cov_p[21] = 1e-4            # roll  (not estimated)
        cov_p[28] = 1e-4            # pitch (not estimated)
        cov_p[35] = self._cov_aa   # yaw,yaw
        msg.pose.covariance = cov_p

        msg.twist.twist.linear.x  = v_lin
        msg.twist.twist.linear.y  = 0.0
        msg.twist.twist.linear.z  = 0.0
        msg.twist.twist.angular.z = v_ang

        # 6×6 row-major twist covariance [vx, vy, vz, wx, wy, wz]
        cov_t = [0.0] * 36
        cov_t[0]  = self._cov_vx   # vx,vx
        cov_t[7]  = 1e-4            # vy (constrained to ~0)
        cov_t[14] = 1e-4            # vz
        cov_t[21] = 1e-4            # wx
        cov_t[28] = 1e-4            # wy
        cov_t[35] = self._cov_wz   # wz,wz
        msg.twist.covariance = cov_t

        return msg

    def _publish_tf(self, stamp) -> None:
        tf = TransformStamped()
        tf.header.stamp             = stamp
        tf.header.frame_id          = self._odom_frame
        tf.child_frame_id           = self._base_frame
        tf.transform.translation.x  = self._odom.x
        tf.transform.translation.y  = self._odom.y
        tf.transform.translation.z  = 0.0
        tf.transform.rotation       = _yaw_to_quaternion(self._odom.theta)
        self._tf_br.sendTransform(tf)


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


if __name__ == "__main__":
    main()
