#!/usr/bin/env python3
"""Speed limiter node for SaltyBot.

Enforces safety constraints from terrain type, battery level, and emergency system.
Computes minimum scale factor and applies to cmd_vel at 50Hz (20ms).

Published topics:
  /saltybot/cmd_vel_limited (geometry_msgs/Twist)
  /saltybot/speed_limit_info (std_msgs/String) - JSON telemetry

Subscribed topics:
  /cmd_vel (geometry_msgs/Twist)
  /saltybot/terrain_speed_scale (std_msgs/Float32) - 0.0 to 1.0
  /saltybot/speed_limit (std_msgs/Float32) - 0.0 to 1.0, from battery level
  /saltybot/emergency (saltybot_emergency_msgs/EmergencyEvent)
"""

import json
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from geometry_msgs.msg import Twist
from std_msgs.msg import Float32, String
from saltybot_emergency_msgs.msg import EmergencyEvent


class SpeedLimiterNode(Node):
    """ROS2 node that enforces speed limits from multiple constraints."""

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

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

        # Last received values (with defaults)
        self.terrain_scale = 1.0
        self.battery_limit = 1.0
        self.emergency_scale = 1.0
        self.last_cmd_vel = None

        # Subscriptions
        self.sub_cmd_vel = self.create_subscription(
            Twist, "/cmd_vel", self._on_cmd_vel, 10
        )
        self.sub_terrain = self.create_subscription(
            Float32, "/saltybot/terrain_speed_scale", self._on_terrain_scale, 10
        )
        self.sub_battery = self.create_subscription(
            Float32, "/saltybot/speed_limit", self._on_battery_limit, 10
        )
        self.sub_emergency = self.create_subscription(
            EmergencyEvent, "/saltybot/emergency", self._on_emergency, 10
        )

        # Publications
        self.pub_cmd_vel_limited = self.create_publisher(
            Twist, "/saltybot/cmd_vel_limited", 10
        )
        self.pub_info = self.create_publisher(
            String, "/saltybot/speed_limit_info", 10
        )

        # Timer for periodic publishing at specified frequency
        period = 1.0 / frequency
        self.timer: Timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"Speed limiter initialized at {frequency}Hz. "
            "Awaiting terrain_speed_scale, speed_limit, and emergency signals."
        )

    def _on_cmd_vel(self, msg: Twist) -> None:
        """Cache the last received cmd_vel."""
        self.last_cmd_vel = msg

    def _on_terrain_scale(self, msg: Float32) -> None:
        """Update terrain speed scale (0.0 to 1.0)."""
        self.terrain_scale = max(0.0, min(1.0, msg.data))

    def _on_battery_limit(self, msg: Float32) -> None:
        """Update battery-driven speed limit (0.0 to 1.0)."""
        self.battery_limit = max(0.0, min(1.0, msg.data))

    def _on_emergency(self, msg: EmergencyEvent) -> None:
        """Update emergency-driven speed scale based on system state."""
        # Map emergency state to speed reduction factor
        if msg.state == "ESCALATED":
            self.emergency_scale = 0.0  # Complete stop
        elif msg.state == "STOPPING":
            self.emergency_scale = 0.1  # Minimal movement for safe deceleration
        elif msg.state == "RECOVERING":
            self.emergency_scale = 0.3  # Cautious recovery
        else:  # NOMINAL
            self.emergency_scale = 1.0  # Full speed allowed

    def _timer_callback(self) -> None:
        """Process and publish speed-limited cmd_vel at 50Hz."""
        if self.last_cmd_vel is None:
            return

        # Compute minimum scale factor from all constraints
        scale_factor = min(self.terrain_scale, self.battery_limit, self.emergency_scale)

        # Apply scale factor to cmd_vel
        limited_cmd_vel = self._scale_twist(self.last_cmd_vel, scale_factor)

        # Publish limited command
        self.pub_cmd_vel_limited.publish(limited_cmd_vel)

        # Publish telemetry
        info = {
            "timestamp": self.get_clock().now().nanoseconds / 1e9,
            "scale_factor": float(scale_factor),
            "terrain_scale": float(self.terrain_scale),
            "battery_limit": float(self.battery_limit),
            "emergency_scale": float(self.emergency_scale),
            "linear_x": float(limited_cmd_vel.linear.x),
            "linear_y": float(limited_cmd_vel.linear.y),
            "linear_z": float(limited_cmd_vel.linear.z),
            "angular_x": float(limited_cmd_vel.angular.x),
            "angular_y": float(limited_cmd_vel.angular.y),
            "angular_z": float(limited_cmd_vel.angular.z),
        }
        info_msg = String(data=json.dumps(info))
        self.pub_info.publish(info_msg)

    @staticmethod
    def _scale_twist(msg: Twist, scale: float) -> Twist:
        """Scale all velocities in a Twist message by a factor."""
        scaled = Twist()
        scaled.linear.x = msg.linear.x * scale
        scaled.linear.y = msg.linear.y * scale
        scaled.linear.z = msg.linear.z * scale
        scaled.angular.x = msg.angular.x * scale
        scaled.angular.y = msg.angular.y * scale
        scaled.angular.z = msg.angular.z * scale
        return scaled


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


if __name__ == "__main__":
    main()
