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

Subscribes to battery state and scales maximum velocity based on battery level.
Prevents over-discharge and extends operational range.

Subscribed topics:
  /saltybot/battery_state (sensor_msgs/BatteryState) - Battery status

Published topics:
  /saltybot/speed_scale (std_msgs/Float32) - Speed scaling factor (0.0-1.0)

Battery level thresholds:
  100-50%:  1.0 scale (full speed)
  50-20%:   0.7 scale (70% speed)
  <20%:     0.4 scale (40% speed - critical)
"""

from typing import Optional

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


class BatterySpeedScalerNode(Node):
    """ROS2 node for battery-aware speed scaling."""

    def __init__(self):
        super().__init__("battery_speed_scaler")

        # Parameters
        self.declare_parameter("frequency", 1)  # Hz
        frequency = self.get_parameter("frequency").value

        # Battery thresholds (percentage)
        self.declare_parameter("critical_threshold", 20.0)
        self.declare_parameter("warning_threshold", 50.0)

        # Speed scaling factors
        self.declare_parameter("full_scale", 1.0)
        self.declare_parameter("warning_scale", 0.7)
        self.declare_parameter("critical_scale", 0.4)

        self.critical_threshold = self.get_parameter("critical_threshold").value
        self.warning_threshold = self.get_parameter("warning_threshold").value
        self.full_scale = self.get_parameter("full_scale").value
        self.warning_scale = self.get_parameter("warning_scale").value
        self.critical_scale = self.get_parameter("critical_scale").value

        # Latest battery state
        self.battery_state: Optional[BatteryState] = None

        # Subscription
        self.create_subscription(
            BatteryState, "/saltybot/battery_state", self._on_battery_state, 10
        )

        # Publisher for speed scale
        self.pub_scale = self.create_publisher(Float32, "/saltybot/speed_scale", 10)

        # Timer for speed scaling at configured frequency
        period = 1.0 / frequency
        self.timer: Timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"Battery speed scaler initialized at {frequency}Hz. "
            f"Thresholds: warning={self.warning_threshold}%, "
            f"critical={self.critical_threshold}%. "
            f"Scale factors: full={self.full_scale}, "
            f"warning={self.warning_scale}, critical={self.critical_scale}"
        )

    def _on_battery_state(self, msg: BatteryState) -> None:
        """Update battery state from subscription."""
        self.battery_state = msg

    def _timer_callback(self) -> None:
        """Compute and publish speed scale based on battery level."""
        if self.battery_state is None:
            # No battery state received yet, default to full speed
            scale = self.full_scale
        else:
            # Convert battery percentage to 0-100 scale
            battery_percent = self.battery_state.percentage * 100.0

            # Determine speed scale based on battery level
            if battery_percent >= self.warning_threshold:
                # Good battery level: full speed
                scale = self.full_scale
            elif battery_percent >= self.critical_threshold:
                # Moderate discharge: warning speed
                scale = self.warning_scale
            else:
                # Critical battery: reduced speed
                scale = self.critical_scale

        # Publish speed scale
        scale_msg = Float32()
        scale_msg.data = scale
        self.pub_scale.publish(scale_msg)


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


if __name__ == "__main__":
    main()
