#!/usr/bin/env python3
"""Battery-aware speed limiter for SaltyBot.

Subscribes to /saltybot/battery_state and publishes a speed limit factor
on /saltybot/speed_limit_factor based on remaining battery percentage.

Thresholds:
  100-50%:  1.0 (full speed)
  50-25%:   0.7 (70% speed)
  25-15%:   0.4 (40% speed)
  <15%:     0.0 (stop, preserve battery)
"""

import rclpy
from rclpy.node import Node
from std_msgs.msg import Float32
from sensor_msgs.msg import BatteryState


class BatterySpeedLimiterNode(Node):
    """ROS2 node for battery-aware speed limiting."""

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

        # Parameters
        self.declare_parameter('threshold_full_1', 100)       # 100%
        self.declare_parameter('threshold_full_2', 50)        # 50%
        self.declare_parameter('threshold_reduced_1', 50)     # 50%
        self.declare_parameter('threshold_reduced_2', 25)     # 25%
        self.declare_parameter('threshold_critical_1', 25)    # 25%
        self.declare_parameter('threshold_critical_2', 15)    # 15%
        self.declare_parameter('threshold_stop', 15)          # <15%

        self.declare_parameter('speed_factor_full', 1.0)
        self.declare_parameter('speed_factor_reduced', 0.7)
        self.declare_parameter('speed_factor_critical', 0.4)
        self.declare_parameter('speed_factor_stop', 0.0)

        self.declare_parameter('publish_frequency', 10)
        self.declare_parameter('enable_limiting', True)

        # Read parameters
        self.threshold_full_1 = self.get_parameter('threshold_full_1').value
        self.threshold_full_2 = self.get_parameter('threshold_full_2').value
        self.threshold_reduced_1 = self.get_parameter('threshold_reduced_1').value
        self.threshold_reduced_2 = self.get_parameter('threshold_reduced_2').value
        self.threshold_critical_1 = self.get_parameter('threshold_critical_1').value
        self.threshold_critical_2 = self.get_parameter('threshold_critical_2').value
        self.threshold_stop = self.get_parameter('threshold_stop').value

        self.speed_factor_full = self.get_parameter('speed_factor_full').value
        self.speed_factor_reduced = self.get_parameter('speed_factor_reduced').value
        self.speed_factor_critical = self.get_parameter('speed_factor_critical').value
        self.speed_factor_stop = self.get_parameter('speed_factor_stop').value

        publish_frequency = self.get_parameter('publish_frequency').value
        self.enable_limiting = self.get_parameter('enable_limiting').value

        # Current battery percentage
        self.current_battery_percent = 100.0

        # Subscription to battery state
        self.sub_battery = self.create_subscription(
            BatteryState, '/saltybot/battery_state', self._on_battery_state, 10
        )

        # Publisher for speed limit factor
        self.pub_speed_limit = self.create_publisher(
            Float32, '/saltybot/speed_limit_factor', 10
        )

        # Timer for publishing at fixed frequency
        period = 1.0 / publish_frequency
        self.timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"Battery speed limiter initialized. "
            f"Thresholds: Full={self.threshold_full_2}%, "
            f"Reduced={self.threshold_reduced_2}%, "
            f"Critical={self.threshold_critical_2}%, "
            f"Stop={self.threshold_stop}%"
        )

    def _on_battery_state(self, msg: BatteryState) -> None:
        """Callback for battery state messages.

        Extract percentage from BatteryState message.
        Handles both direct percentage field and calculated percentage.
        """
        # Try to get percentage from message
        if msg.percentage >= 0:
            self.current_battery_percent = msg.percentage * 100.0
        else:
            # Calculate from voltage if percentage not available
            # Standard LiPo cells: 3V min, 4.2V max per cell
            # Assume 2S (8.4V max) or 3S (12.6V max) pack
            # This is a fallback - ideally percentage is provided
            if msg.voltage > 0:
                # Normalize to 0-1 based on typical LiPo range
                min_voltage = 6.0  # 2S minimum
                max_voltage = 8.4  # 2S maximum
                normalized = (msg.voltage - min_voltage) / (max_voltage - min_voltage)
                self.current_battery_percent = max(0.0, min(100.0, normalized * 100.0))

    def _calculate_speed_limit(self, battery_percent: float) -> float:
        """Calculate speed limit factor based on battery percentage.

        Args:
            battery_percent: Battery state as percentage (0-100)

        Returns:
            Speed limit factor (0.0-1.0)
        """
        if not self.enable_limiting:
            return 1.0

        # Apply thresholds
        if battery_percent >= self.threshold_full_2:
            # 100-50%: Full speed
            return self.speed_factor_full
        elif battery_percent >= self.threshold_reduced_2:
            # 50-25%: Reduced speed (70%)
            return self.speed_factor_reduced
        elif battery_percent >= self.threshold_critical_2:
            # 25-15%: Critical speed (40%)
            return self.speed_factor_critical
        else:
            # <15%: Stop to preserve battery
            return self.speed_factor_stop

    def _timer_callback(self) -> None:
        """Periodically publish speed limit factor based on battery state."""
        speed_limit = self._calculate_speed_limit(self.current_battery_percent)
        msg = Float32(data=speed_limit)
        self.pub_speed_limit.publish(msg)


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


if __name__ == '__main__':
    main()