#!/usr/bin/env python3
"""cmd_vel multiplexer for SaltyBot.

Priority-based routing with 0.5s timeout per source.
Teleop (1) > Nav2 (2) > Docking (3)

Published topics:
  /cmd_vel (geometry_msgs/Twist) - Multiplexed velocity command
  /saltybot/cmd_vel_mux_status (std_msgs/String) - JSON mux status

Subscribed topics:
  /cmd_vel_teleop (geometry_msgs/Twist) - Teleop command (priority 1)
  /cmd_vel_nav2 (geometry_msgs/Twist) - Nav2 command (priority 2)
  /cmd_vel_docking (geometry_msgs/Twist) - Docking command (priority 3)
"""

import json
from enum import Enum
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from geometry_msgs.msg import Twist
from std_msgs.msg import String


class CmdSource(Enum):
    """Command velocity sources with priority."""
    TELEOP = 1
    NAV2 = 2
    DOCKING = 3
    NONE = 0


class CmdVelMuxNode(Node):
    """ROS2 node for priority-based cmd_vel multiplexing."""

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

        # Parameters
        self.declare_parameter("frequency", 50)  # Hz
        self.declare_parameter("source_timeout", 0.5)  # seconds

        frequency = self.get_parameter("frequency").value
        self.source_timeout = self.get_parameter("source_timeout").value

        # Last update times for each source
        self.last_update = {
            CmdSource.TELEOP: None,
            CmdSource.NAV2: None,
            CmdSource.DOCKING: None,
        }

        # Last received commands
        self.last_cmd = {
            CmdSource.TELEOP: None,
            CmdSource.NAV2: None,
            CmdSource.DOCKING: None,
        }

        # Current active source
        self.active_source = CmdSource.NONE

        # Subscriptions
        self.create_subscription(
            Twist, "/cmd_vel_teleop", self._on_cmd_vel_teleop, 10
        )
        self.create_subscription(
            Twist, "/cmd_vel_nav2", self._on_cmd_vel_nav2, 10
        )
        self.create_subscription(
            Twist, "/cmd_vel_docking", self._on_cmd_vel_docking, 10
        )

        # Publications
        self.pub_mux = self.create_publisher(Twist, "/cmd_vel", 10)
        self.pub_status = self.create_publisher(String, "/saltybot/cmd_vel_mux_status", 10)

        # Timer for multiplexing at 50Hz
        period = 1.0 / frequency
        self.timer: Timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"cmd_vel mux initialized at {frequency}Hz. "
            f"Source timeout: {self.source_timeout}s. "
            f"Priority: Teleop(1) > Nav2(2) > Docking(3)"
        )

    def _on_cmd_vel_teleop(self, msg: Twist) -> None:
        """Update teleop command."""
        self.last_update[CmdSource.TELEOP] = self.get_clock().now()
        self.last_cmd[CmdSource.TELEOP] = msg

    def _on_cmd_vel_nav2(self, msg: Twist) -> None:
        """Update Nav2 command."""
        self.last_update[CmdSource.NAV2] = self.get_clock().now()
        self.last_cmd[CmdSource.NAV2] = msg

    def _on_cmd_vel_docking(self, msg: Twist) -> None:
        """Update docking command."""
        self.last_update[CmdSource.DOCKING] = self.get_clock().now()
        self.last_cmd[CmdSource.DOCKING] = msg

    def _timer_callback(self) -> None:
        """Multiplex cmd_vel based on priority and timeout at 50Hz."""
        now = self.get_clock().now()

        # Check which sources are active (within timeout)
        active_sources = []

        for source in [CmdSource.TELEOP, CmdSource.NAV2, CmdSource.DOCKING]:
            last_time = self.last_update[source]

            if last_time is not None:
                elapsed = (now - last_time).nanoseconds / 1e9

                if elapsed <= self.source_timeout:
                    active_sources.append(source)

        # Select highest priority active source
        if active_sources:
            # Sources are already in priority order
            self.active_source = active_sources[0]
        else:
            self.active_source = CmdSource.NONE

        # Publish multiplexed command
        if self.active_source != CmdSource.NONE:
            cmd = self.last_cmd[self.active_source]
            self.pub_mux.publish(cmd)
        else:
            # No active source: publish zero command
            self.pub_mux.publish(Twist())

        # Publish mux status
        self._publish_status()

    def _publish_status(self) -> None:
        """Publish multiplexer status as JSON."""
        now = self.get_clock().now()

        # Build source status
        sources_status = {}
        for source in [CmdSource.TELEOP, CmdSource.NAV2, CmdSource.DOCKING]:
            last_time = self.last_update[source]

            if last_time is not None:
                elapsed = (now - last_time).nanoseconds / 1e9
                is_active = elapsed <= self.source_timeout
            else:
                elapsed = None
                is_active = False

            sources_status[source.name] = {
                "priority": source.value,
                "active": is_active,
                "elapsed_s": elapsed,
            }

        status = {
            "timestamp": now.nanoseconds / 1e9,
            "active_source": self.active_source.name,
            "active_priority": self.active_source.value,
            "sources": sources_status,
        }

        msg = String(data=json.dumps(status))
        self.pub_status.publish(msg)


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


if __name__ == "__main__":
    main()
