#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist, TransformStamped, PoseWithCovarianceStamped
from sensor_msgs.msg import Imu, JointState, BatteryState
from nav_msgs.msg import Odometry
from std_msgs.msg import String
import serial
import struct
import time
import math

class ZLAC8015DModbus:
    """
    Lightweight, zero-dependency Python implementation of ZLAC8015D Modbus RTU master.
    """
    def __init__(self, port, baudrate=115200, timeout=0.02):
        self.ser = serial.Serial(port, baudrate, timeout=timeout)
        
    def calculate_crc(self, data):
        crc = 0xFFFF
        for pos in data:
            crc ^= pos
            for _ in range(8):
                if (crc & 1) != 0:
                    crc >>= 1
                    crc ^= 0xA001
                else:
                    crc >>= 1
        return crc

    def write_single(self, slave_id, reg, value):
        packet = struct.pack('>BBHH', slave_id, 0x06, reg, value)
        crc = self.calculate_crc(packet)
        packet += struct.pack('<H', crc)
        try:
            self.ser.write(packet)
            self.ser.flush()
            # Read 8 byte echo response
            self.ser.read(8)
            return True
        except Exception as e:
            return False

    def write_multiple(self, slave_id, start_reg, values):
        num_regs = len(values)
        header = struct.pack('>BBHHB', slave_id, 0x10, start_reg, num_regs, num_regs * 2)
        data = b''
        for val in values:
            data += struct.pack('>h', val)
        packet = header + data
        crc = self.calculate_crc(packet)
        packet += struct.pack('<H', crc)
        try:
            self.ser.write(packet)
            self.ser.flush()
            # Read 8 byte response
            self.ser.read(8)
            return True
        except Exception as e:
            return False

    def read_registers(self, slave_id, start_reg, num_regs):
        packet = struct.pack('>BBHH', slave_id, 0x03, start_reg, num_regs)
        crc = self.calculate_crc(packet)
        packet += struct.pack('<H', crc)
        try:
            # Clear input buffer
            self.ser.reset_input_buffer()
            self.ser.write(packet)
            self.ser.flush()
            
            expected_len = 5 + num_regs * 2
            response = self.ser.read(expected_len)
            if len(response) < expected_len:
                return None
            
            # Verify CRC
            resp_crc = self.calculate_crc(response[:-2])
            calc_crc = struct.unpack('<H', response[-2:])[0]
            if resp_crc != calc_crc:
                return None
                
            values = []
            for i in range(num_regs):
                val = struct.unpack('>h', response[3 + i*2 : 5 + i*2])[0]
                values.append(val)
            return values
        except Exception as e:
            return None


class ForiSerialBridge(Node):
    def __init__(self):
        super().__init__('fori_serial_bridge')
        
        # --- [ROS Parameters] ---
        # "arduino" = PC -> Arduino Mega -> Drivers
        # "direct_pc" = PC -> USB-to-RS485 -> Drivers
        self.declare_parameter('control_method', 'direct_pc')
        self.declare_parameter('port', '/dev/ttyUSB0')  # Default port (ttyACM0 for Arduino, ttyUSB0 for direct)
        self.declare_parameter('baud', 115200)
        self.declare_parameter('target_linear_speed', 0.5)  
        self.declare_parameter('accel_limit', 0.5)         
        self.declare_parameter('wheel_base', 0.374)         
        self.declare_parameter('wheel_radius', 0.127)      

        self.control_method = self.get_parameter('control_method').value
        self.port = self.get_parameter('port').value
        self.baud = self.get_parameter('baud').value
        self.limit_v = self.get_parameter('target_linear_speed').value
        self.accel_limit = self.get_parameter('accel_limit').value
        self.wheel_base = self.get_parameter('wheel_base').value
        self.wheel_radius = self.get_parameter('wheel_radius').value

        # Motor parameters
        self.MAX_RPM = 250
        
        # --- [IMU Yaw Integration / Control] ---
        self.Kp_yaw = 0.50             
        self.Ki_yaw = 0.20             
        self.yaw_error_integral = 0.0  
        self.max_integral = 1.0        

        # --- [Robot states] ---
        self.target_v = 0.0
        self.target_w = 0.0
        self.current_v = 0.0
        self.current_w = 0.0
        self.current_imu_w = 0.0       
        self.last_time = self.get_clock().now()

        # Odometry states
        self.odom_x = 0.0
        self.odom_y = 0.0
        self.odom_th = 0.0

        # Wheel positions for JointState
        self.left_wheel_joint_pos = 0.0
        self.right_wheel_joint_pos = 0.0

        # Driver state tracker (direct mode only)
        self.drivers_initialized = False
        self.drivers_enabled = False
        self.is_parked = False  # Track if the robot is fully parked

        self.get_logger().info('==========================================')
        self.get_logger().info('     FORI BLDC ZLAC8015D CONTROL BRIDGE   ')
        self.get_logger().info('==========================================')
        self.get_logger().info(f'    Control Method: {self.control_method}')
        self.get_logger().info(f'    Serial Port: {self.port} at {self.baud} baud')
        self.get_logger().info('==========================================')

        self.is_mock_mode = False
        # --- [Serial/Modbus Connection] ---
        if self.control_method == 'direct_pc':
            try:
                self.modbus = ZLAC8015DModbus(self.port, self.baud)
                self.get_logger().info('BLDC 드라이버 초기화 대기 중...')
                self.init_direct_drivers()
                self.get_logger().info('BLDC 드라이버 초기화 완료!')
            except Exception as e:
                self.get_logger().warn(f'!!! 드라이버 연결 실패: {e} !!!')
                self.get_logger().warn('실제 모터 드라이버 연결을 찾을 수 없어 [가상 모의 주행(Mock Mode)]으로 동작합니다.')
                self.is_mock_mode = True
        elif self.control_method == 'arduino':
            try:
                self.ser = serial.Serial(self.port, self.baud, timeout=0.01)
                self.get_logger().info('아두이노 초기화 대기 중 (3초)...')
                time.sleep(3)
                self.get_logger().info('아두이노 메가 통신 준비 완료!')
            except Exception as e:
                self.get_logger().warn(f'!!! 아두이노 시리얼 연결 실패: {e} !!!')
                self.get_logger().warn('아두이노 연결을 찾을 수 없어 [가상 모의 주행(Mock Mode)]으로 동작합니다.')
                self.is_mock_mode = True

        # --- [ROS Publishers & Subscribers] ---
        qos_profile = rclpy.qos.QoSProfile(depth=1)
        self.create_subscription(Twist, '/cmd_vel', self.cmd_vel_callback, qos_profile)
        self.create_subscription(Imu, '/livox/imu', self.imu_callback, qos_profile)
        self.create_subscription(PoseWithCovarianceStamped, '/initialpose', self.initial_pose_callback, 10)
        self.create_subscription(String, '/robot_mode_status', self.mode_status_callback, 10)
        
        # Telemetry/Odometry/Battery publishers
        self.odom_pub = self.create_publisher(Odometry, '/odom_wheels', qos_profile)
        self.joint_pub = self.create_publisher(JointState, '/joint_states', qos_profile)
        self.battery_pub = self.create_publisher(BatteryState, '/battery_state', qos_profile)

        # Control loop timer (20Hz)
        self.create_timer(0.05, self.control_loop)
        # Battery status timer (1Hz)
        self.create_timer(1.0, self.publish_battery_status)

    def cmd_vel_callback(self, msg):
        self.target_v = max(min(msg.linear.x, self.limit_v), -self.limit_v)
        self.target_w = msg.angular.z

    def imu_callback(self, msg):
        self.current_imu_w = msg.angular_velocity.z

    def initial_pose_callback(self, msg):
        self.odom_x = msg.pose.pose.position.x
        self.odom_y = msg.pose.pose.position.y
        
        q = msg.pose.pose.orientation
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
        self.odom_th = math.atan2(siny_cosp, cosy_cosp)
        
        self.get_logger().info(f'[MOCK] Reset robot pose to: x={self.odom_x:.2f}, y={self.odom_y:.2f}, yaw={self.odom_th:.2f}')

    def mode_status_callback(self, msg):
        # Checks if the robot has successfully arrived in PARKED state
        self.is_parked = "State: PARKED" in msg.data

    def publish_battery_status(self):
        msg = BatteryState()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'base_link'
        msg.design_capacity = 60.0  # 60Ah LiFePO4
        msg.capacity = 60.0        # 60Ah LiFePO4
        msg.power_supply_technology = BatteryState.POWER_SUPPLY_TECHNOLOGY_LIFE # LiFePO4

        # === LiFePO4 8S Battery Voltage Bounds ===
        # Full charge: 8 × 3.65V = 29.2V
        # Cutoff:      8 × 2.80V = 22.4V
        v_min = 22.4
        v_max = 29.2

        # Initialize last-valid-voltage storage
        if not hasattr(self, 'last_valid_voltage'):
            self.last_valid_voltage = None  # None = no valid reading yet

        voltage = None  # Will be set from Modbus or fallback

        if not self.is_mock_mode and self.control_method == 'direct_pc':
            try:
                # Read registers 0x20A0 (External Voltage) and 0x20A1 (Bus Voltage)
                v_resp = self.modbus.read_registers(1, 0x20A0, 2)
                if v_resp and len(v_resp) > 0:
                    raw_a0 = abs(v_resp[0])
                    raw_a1 = abs(v_resp[1]) if len(v_resp) > 1 else raw_a0

                    # ZLAC8015D reports voltage in units of 0.1V (e.g. 283 = 28.3V)
                    # or sometimes 0.01V (e.g. 2830 = 28.3V). Detect scale by range.
                    raw_v = 0.0
                    if 200 <= raw_a0 <= 320:       # 20.0V ~ 32.0V in 0.1V steps
                        raw_v = raw_a0 * 0.1
                    elif 2000 <= raw_a0 <= 3200:   # 20.0V ~ 32.0V in 0.01V steps
                        raw_v = raw_a0 * 0.01
                    elif 200 <= raw_a1 <= 320:
                        raw_v = raw_a1 * 0.1
                    elif 2000 <= raw_a1 <= 3200:
                        raw_v = raw_a1 * 0.01
                    else:
                        # Fallback guess
                        raw_v = raw_a1 * 0.01 if raw_a1 > 1000 else (raw_a1 * 0.1 if raw_a1 > 100 else float(raw_a1))

                    # Sanity check: plausible battery voltage range (20V ~ 32V)
                    if 20.0 <= raw_v <= 32.0:
                        voltage = raw_v
                        self.last_valid_voltage = voltage
                        self.get_logger().info(
                            f'[BATTERY] Driver raw: {raw_v:.2f}V  |  '
                            f'{max(0.0, min(100.0, (raw_v - v_min) / (v_max - v_min) * 100)):.1f}%',
                            throttle_duration_sec=3.0
                        )
                    else:
                        self.get_logger().warn(
                            f'[BATTERY] Raw voltage out of plausible range: raw_a0={raw_a0}, raw_a1={raw_a1}, computed={raw_v:.2f}V',
                            throttle_duration_sec=5.0
                        )
                else:
                    self.get_logger().warn('[BATTERY] Modbus read returned empty response.', throttle_duration_sec=5.0)
            except Exception as e:
                self.get_logger().warn(f'[BATTERY] Modbus read failed: {e}', throttle_duration_sec=5.0)

        # Fallback priority: last valid reading > mock nominal
        if voltage is None:
            if self.last_valid_voltage is not None:
                voltage = self.last_valid_voltage
                self.get_logger().warn('[BATTERY] Using last valid voltage reading as fallback.', throttle_duration_sec=10.0)
            else:
                # Absolute fallback: use midpoint of nominal LiFePO4 8S range
                voltage = (v_min + v_max) / 2.0  # ~25.8V nominal
                self.get_logger().warn(f'[BATTERY] No valid reading yet - using nominal fallback: {voltage:.1f}V', throttle_duration_sec=10.0)

        # Sliding Window Median + Heavy EMA Filter (eliminates Modbus spikes)
        if not hasattr(self, 'voltage_history'):
            import collections
            self.voltage_history = collections.deque(maxlen=10)

        self.voltage_history.append(voltage)

        sorted_v = sorted(self.voltage_history)
        median_v = sorted_v[len(sorted_v) // 2]

        if not hasattr(self, 'filtered_battery_voltage') or self.filtered_battery_voltage is None:
            self.filtered_battery_voltage = median_v
        else:
            # EMA smoothing (Alpha=0.05 = ~10s transition)
            self.filtered_battery_voltage = self.filtered_battery_voltage * 0.95 + median_v * 0.05

        voltage = self.filtered_battery_voltage

        percentage = (voltage - v_min) / (v_max - v_min)
        percentage = max(0.0, min(1.0, percentage))

        msg.voltage = float(voltage)
        msg.percentage = float(percentage)
        msg.power_supply_status = BatteryState.POWER_SUPPLY_STATUS_DISCHARGING
        self.battery_pub.publish(msg)


    def init_direct_drivers(self):
        """
        Initializes ZLAC8015D parameters (Operating mode to Velocity Control).
        """
        # Reg 0x200D (Operating Mode) = 3 (Velocity Mode)
        self.modbus.write_single(1, 0x200D, 3)
        time.sleep(0.01)
        self.modbus.write_single(2, 0x200D, 3)
        time.sleep(0.01)
        self.disable_direct_drivers()
        self.drivers_initialized = True

    def enable_direct_drivers(self):
        if self.drivers_enabled:
            return
        # Release brakes (0x201A, 0x201B) = 0
        self.modbus.write_single(1, 0x201A, 0)
        self.modbus.write_single(1, 0x201B, 0)
        self.modbus.write_single(2, 0x201A, 0)
        self.modbus.write_single(2, 0x201B, 0)
        time.sleep(0.01)

        # Enable drivers (0x200E) = 8
        self.modbus.write_single(1, 0x200E, 8)
        self.modbus.write_single(2, 0x200E, 8)
        self.drivers_enabled = True
        self.get_logger().info('BLDC Drivers and Brakes Released (ENABLED)')

    def disable_direct_drivers(self):
        # Stop velocity
        self.modbus.write_multiple(1, 0x2088, [0, 0])
        self.modbus.write_multiple(2, 0x2088, [0, 0])

        # Lock brakes (0x201A, 0x201B) = 1
        self.modbus.write_single(1, 0x201A, 1)
        self.modbus.write_single(1, 0x201B, 1)
        self.modbus.write_single(2, 0x201A, 1)
        self.modbus.write_single(2, 0x201B, 1)
        time.sleep(0.01)

        # Disable drivers (0x200E) = 7
        self.modbus.write_single(1, 0x200E, 7)
        self.modbus.write_single(2, 0x200E, 7)
        self.drivers_enabled = False
        self.get_logger().info('BLDC Brakes Locked (DISABLED)')

    def control_loop(self):
        now = self.get_clock().now()
        dt = (now - self.last_time).nanoseconds / 1e9
        self.last_time = now
        if dt <= 0: return

        # 1. Acceleration Ramp (Smoothing)
        dv = self.target_v - self.current_v
        max_dv = self.accel_limit * dt
        self.current_v += max(min(dv, max_dv), -max_dv)

        dw = self.target_w - self.current_w
        max_dw = 2.0 * dt 
        self.current_w += max(min(dw, max_dw), -max_dw)

        # 2. Yaw rate PI Control
        yaw_error = self.current_w - self.current_imu_w
        
        # Stop condition
        is_idle = (abs(self.target_v) < 0.01 and abs(self.target_w) < 0.01)
        if is_idle:
            self.yaw_error_integral = 0.0
            correction = 0.0
            self.current_v = 0.0  
            self.current_w = 0.0
        else:
            self.yaw_error_integral += yaw_error * dt
            self.yaw_error_integral = max(min(self.yaw_error_integral, self.max_integral), -self.max_integral)
            correction = (self.Kp_yaw * yaw_error) + (self.Ki_yaw * self.yaw_error_integral)

        # 3. Inverse Kinematics (Wheel Linear Speeds)
        v_left = self.current_v - ((self.current_w + correction) * self.wheel_base / 2.0)
        v_right = self.current_v + ((self.current_w + correction) * self.wheel_base / 2.0)

        # Convert to RPM (RPM = v / (2 * pi * r) * 60)
        rpm_left = int(round(v_left / (2.0 * math.pi * self.wheel_radius) * 60.0))
        rpm_right = int(round(v_right / (2.0 * math.pi * self.wheel_radius) * 60.0))

        # Clip speed limits
        rpm_left = max(min(rpm_left, self.MAX_RPM), -self.MAX_RPM)
        rpm_right = max(min(rpm_right, self.MAX_RPM), -self.MAX_RPM)

        # 4. Command the Motors
        feedback_speeds = [0, 0, 0, 0] # FL, FR, RL, RR

        if self.is_mock_mode:
            feedback_speeds[0] = rpm_left
            feedback_speeds[1] = rpm_right
            feedback_speeds[2] = rpm_left
            feedback_speeds[3] = rpm_right
        else:
            if self.control_method == 'direct_pc':
                # Only disable drivers and lock brakes when idle AND fully parked
                if is_idle and self.is_parked:
                    if self.drivers_enabled:
                        self.disable_direct_drivers()
                else:
                    if not self.drivers_enabled:
                        self.enable_direct_drivers()
                    # Write velocities (invert right motor target RPM due to mirrored physical mounting)
                    self.modbus.write_multiple(1, 0x2088, [rpm_left, -rpm_right])
                    self.modbus.write_multiple(2, 0x2088, [rpm_left, -rpm_right])

                # Read actual speeds (invert right motor feedback RPM due to mirrored physical mounting)
                front_resp = self.modbus.read_registers(1, 0x20AD, 2)
                rear_resp = self.modbus.read_registers(2, 0x20AD, 2)
                if front_resp:
                    feedback_speeds[0] = front_resp[0]
                    feedback_speeds[1] = -front_resp[1]
                if rear_resp:
                    feedback_speeds[2] = rear_resp[0]
                    feedback_speeds[3] = -rear_resp[1]

            elif self.control_method == 'arduino':
                # Send command packet
                # Format: [0xFE, left_rpm_H, left_rpm_L, right_rpm_H, right_rpm_L, control_state, checksum]
                packet = bytearray(7)
                packet[0] = 0xFE
                struct.pack_into('>h', packet, 1, rpm_left)
                struct.pack_into('>h', packet, 3, rpm_right)
                # control_state: 0 (disable) only when idle AND fully parked, otherwise 1 (enable)
                packet[5] = 0 if (is_idle and self.is_parked) else 1
                packet[6] = sum(packet[1:6]) & 0xFF
                
                try:
                    self.ser.write(bytes(packet))
                    self.ser.flush()
                except Exception as e:
                    self.get_logger().error(f'아두이노 송신 오류: {e}')

                # Read feedback packet
                # Format: [0xFD, f_lh, f_ll, f_rh, f_rl, r_lh, r_ll, r_rh, r_rl, checksum]
                try:
                    if self.ser.in_waiting >= 10:
                        head = self.ser.read(1)
                        if head == b'\xfd':
                            data = self.ser.read(9)
                            if len(data) == 9:
                                checksum = sum(data[:-1]) & 0xFF
                                if checksum == data[-1]:
                                    f_lh, f_rh, r_lh, r_rh = struct.unpack('>hhhh', data[:-1])
                                    feedback_speeds[0] = f_lh
                                    feedback_speeds[1] = f_rh
                                    feedback_speeds[2] = r_lh
                                    feedback_speeds[3] = r_rh
                except Exception as e:
                    self.get_logger().error(f'아두이노 수신 오류: {e}')

        # 5. Odometry & Telemetry Processing
        self.process_feedback(feedback_speeds, dt)

        # Print debug log (2s throttle)
        self.get_logger().info(
            f'[BLDC FEEDBACK] FL:{feedback_speeds[0]} FR:{feedback_speeds[1]} | '
            f'RL:{feedback_speeds[2]} RR:{feedback_speeds[3]} RPM | '
            f'Target_W:{self.current_w:.2f} IMU_W:{self.current_imu_w:.2f}',
            throttle_duration_sec=2.0
        )

    def process_feedback(self, speeds, dt):
        # Average front/rear left and right wheel speeds
        actual_rpm_l = (speeds[0] + speeds[2]) / 2.0
        actual_rpm_r = (speeds[1] + speeds[3]) / 2.0

        # Convert RPM to linear speeds (v = rpm * 2 * pi * r / 60)
        v_l = actual_rpm_l * 2.0 * math.pi * self.wheel_radius / 60.0
        v_r = actual_rpm_r * 2.0 * math.pi * self.wheel_radius / 60.0

        # Calculate robot linear and angular velocities
        linear_vel = (v_r + v_l) / 2.0
        angular_vel = (v_r - v_l) / self.wheel_base

        # Integrate pose (Odometry)
        delta_th = angular_vel * dt
        self.odom_th += delta_th
        self.odom_x += linear_vel * math.cos(self.odom_th) * dt
        self.odom_y += linear_vel * math.sin(self.odom_th) * dt

        # Update joint states
        self.left_wheel_joint_pos += (v_l / self.wheel_radius) * dt
        self.right_wheel_joint_pos += (v_r / self.wheel_radius) * dt

        # Publish JointState
        joint_state = JointState()
        joint_state.header.stamp = self.get_clock().now().to_msg()
        joint_state.name = ['front_left_wheel_joint', 'front_right_wheel_joint', 
                            'rear_left_wheel_joint', 'rear_right_wheel_joint']
        # Calculate velocity in rad/s
        rads_l = v_l / self.wheel_radius
        rads_r = v_r / self.wheel_radius
        joint_state.position = [self.left_wheel_joint_pos, self.right_wheel_joint_pos,
                                self.left_wheel_joint_pos, self.right_wheel_joint_pos]
        joint_state.velocity = [rads_l, rads_r, rads_l, rads_r]
        self.joint_pub.publish(joint_state)

        # Publish Odometry msg
        odom = Odometry()
        odom.header.stamp = joint_state.header.stamp
        odom.header.frame_id = 'odom'
        odom.child_frame_id = 'base_link'

        # Set positions
        odom.pose.pose.position.x = self.odom_x
        odom.pose.pose.position.y = self.odom_y
        odom.pose.pose.position.z = 0.0
        
        # Quaternion from yaw
        cy = math.cos(self.odom_th * 0.5)
        sy = math.sin(self.odom_th * 0.5)
        odom.pose.pose.orientation.w = cy
        odom.pose.pose.orientation.z = sy

        # Set velocities
        odom.twist.twist.linear.x = linear_vel
        odom.twist.twist.angular.z = angular_vel

        self.odom_pub.publish(odom)

def main(args=None):
    rclpy.init(args=args)
    node = ForiSerialBridge()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        if node.control_method == 'direct_pc':
            node.disable_direct_drivers()
        elif node.control_method == 'arduino':
            try:
                # Send explicit stop packet
                packet = bytearray(7)
                packet[0] = 0xFE
                packet[5] = 0
                packet[6] = sum(packet[1:6]) & 0xFF
                node.ser.write(bytes(packet))
                node.ser.flush()
            except:
                pass
        node.get_logger().info('ROS2 BRIDGE STOPPED')
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
