fanuc_ucl 1.5.4

#![allow(clippy::unnecessary_map_on_constructor, clippy::useless_conversion)]

use std::{
    collections::VecDeque,
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::{Duration, Instant},
};

use cfg_mixin::cfg_mixin;
use event_listener::{Event, Listener};
use flume::{Receiver, Sender};

use crate::{
    joints::JointDataSizeError,
    stmo::{
        JointMovementLimit,
        proto::{
            CommandPositionRequestPacket, CommandPositionResponsePacket, MotionCommandPacket,
            RobotStatusPacket, RxPackets, StartPacket, StopPacket, ThresholdTableRequestPacket,
            TxPackets, VersionNumberRequestPacket,
        },
        stmo_handle::StmoHandle,
        types::{AxisMotionConstraint, JointMovementLimits, RxStorage, StreamMotionError},
    },
    thread_util::{GeneralThreadError, ThreadConfig, ThreadHandle},
};

use snare::mio::net::UdpSocket as MioUdpSocket;
use snare::mio::{Events, Interest, Poll, Token, Waker};

#[cfg(feature = "py")]
use pyo3::prelude::*;

const TOK_SOCKET: Token = Token(0);
const TOK_WAKER: Token = Token(1);

#[derive(Debug, Clone)]
enum MaybeMany<T: Clone> {
    One(T),
    Many(Vec<T>),
}

enum ToThreadMessage {
    Start(StartPacket),
    Stop(StopPacket),
    ThresholdTableRequest(ThresholdTableRequestPacket),
    MotionCommandDouble(MaybeMany<MotionCommandPacket>, Option<StmoHandle>),
}

#[derive(Debug)]
struct StreamMotionContext {
    socket: MioUdpSocket,
    from_driver: Receiver<ToThreadMessage>,
    to_driver: Sender<RxPackets>,
    protocol_version: u32,
    send_last_command: bool,
    last_command_position_request_time: Instant,
    motion_command_queue: VecDeque<(MaybeMany<MotionCommandPacket>, Option<StmoHandle>)>,
    itl: Arc<(Event, AtomicBool)>,
}

impl StreamMotionContext {
    const COMMAND_POSITION_RATE: Duration = Duration::from_millis(128);

    fn new(
        from_driver: Receiver<ToThreadMessage>,
        to_driver: Sender<RxPackets>,
        socket: MioUdpSocket,
        itl: Arc<(Event, AtomicBool)>,
        send_last_command: bool,
    ) -> Self {
        Self {
            from_driver,
            to_driver,
            socket,
            protocol_version: 0,
            last_command_position_request_time: Instant::now() - Self::COMMAND_POSITION_RATE,
            motion_command_queue: VecDeque::new(),
            itl,
            send_last_command,
        }
    }

    fn send(
        &mut self,
        tx: TxPackets,
        buf: &mut [u8],
        timeout: Option<Duration>,
        version_override: Option<u32>,
    ) -> Result<(), StreamMotionError> {
        let n = tx.encode_into(version_override.unwrap_or(self.protocol_version), buf)?;
        match self.socket.send(&buf[..n]) {
            Ok(_) => Ok(()),
            Err(ref e)
                if e.kind() == io::ErrorKind::WouldBlock
                    || e.kind() == io::ErrorKind::Interrupted =>
            {
                if let Some(to) = timeout {
                    match self.retry_sending(tx, to, buf) {
                        Ok(()) => Ok(()),
                        Err(e) => Err(e),
                    }
                } else {
                    log::warn!("STMO send would block and no timeout configured");
                    Err(StreamMotionError::Timeout)
                }
            }
            Err(e) => {
                log::error!("STMO UDP send error: {}", e);
                Err(StreamMotionError::from(e))
            }
        }
    }

    fn retry_sending(
        &mut self,
        tx: TxPackets,
        timeout: Duration,
        buf: &mut [u8],
    ) -> Result<(), StreamMotionError> {
        let start = Instant::now();
        let mut sent = false;
        let n = tx.encode_into(self.protocol_version, buf)?;
        let sleeper = spin_sleep::SpinSleeper::new(1_000_000);
        while !sent && start.elapsed() < timeout {
            match self.socket.send(&buf[..n]) {
                Ok(_) => sent = true,
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock
                        || e.kind() == io::ErrorKind::Interrupted =>
                {
                    sleeper.sleep(Duration::from_micros(500));
                }
                Err(e) => {
                    log::error!("Error sending packet: {:?}", e);
                    return Err(StreamMotionError::from(e));
                }
            }
        }
        if sent {
            Ok(())
        } else {
            Err(StreamMotionError::Timeout)
        }
    }

    fn next_motion_command(&mut self) -> Option<(MotionCommandPacket, Option<StmoHandle>)> {
        loop {
            let should_pop_entry = match self.motion_command_queue.front()? {
                (MaybeMany::One(_), _) => true,
                (MaybeMany::Many(vec), _) => vec.len() <= 1,
            };

            if should_pop_entry {
                let (cmds, handle) = self.motion_command_queue.pop_front()?;
                match cmds {
                    MaybeMany::One(cmd) => return Some((cmd, handle)),
                    MaybeMany::Many(mut vec) => {
                        if let Some(cmd) = vec.pop() {
                            return Some((cmd, handle));
                        } else {
                            // empty batch — fulfill handle and try next entry
                            if let Some(h) = handle {
                                h.set();
                            }
                            continue;
                        }
                    }
                }
            } else {
                // Many with >1 element — pop one without consuming the handle yet
                if let Some((MaybeMany::Many(vec), _)) = self.motion_command_queue.front_mut() {
                    return vec.pop().map(|c| (c, None));
                }
                return None;
            }
        }
    }

    pub fn context_loop(mut self, thread_handle: ThreadHandle, mut poll: Poll) {
        let mut events = Events::with_capacity(64);
        let mut rx_buf = [0u8; 2048];
        let mut tx_buf = [0u8; 1024];
        let mut prev_motion_packet: Option<MotionCommandPacket> = None;
        let mut prev_command_was_real = false;
        let mut consecutive_fillers: u32 = 0;
        // let mut status_cycle_count: u32 = 0;

        while thread_handle.should_live() {
            if let Err(e) = poll.poll(&mut events, None) {
                if e.kind() == io::ErrorKind::Interrupted {
                    continue;
                }
                log::error!("STMO poll error, breaking event loop: {}", e);
                break;
            }

            for ev in events.iter() {
                match ev.token() {
                    TOK_SOCKET => {
                        // drain UDP
                        loop {
                            match self.socket.recv(&mut rx_buf) {
                                Ok(n) if n > 0 => {
                                    if let Some(rx) = RxPackets::decode_from(&rx_buf[..n]) {
                                        let _ = self.to_driver.send(rx);
                                        log::trace!("Received packet: {:?}", rx);
                                        if let RxPackets::VersionNumberResponse(vn) = &rx {
                                            self.protocol_version = vn.version;
                                            log::info!(
                                                "Detected Stream Motion protocol version {}",
                                                self.protocol_version
                                            );
                                        }
                                        // respond with the best matching motion command if applicable
                                        if let RxPackets::RobotStatus(state) = &rx {
                                            if !state.status_bits().ready_for_commands() {
                                                continue;
                                            }
                                            // status_cycle_count = status_cycle_count.wrapping_add(1);
                                            if state.status_bits().packet_rate() as u32 != 0 {
                                                log::debug!(
                                                    "Robot status packet rate: {}",
                                                    state.status_bits().packet_rate()
                                                );
                                            }
                                            // status_cycle_count = 0;
                                            if let Some((mut cmd, handle)) =
                                                self.next_motion_command()
                                            {
                                                cmd.seq = state.seq;
                                                if consecutive_fillers > 0 {
                                                    log::debug!(
                                                        "STMO queue refilled (seq {}) after {} filler cycle(s) (~{}ms starved)",
                                                        state.seq,
                                                        consecutive_fillers,
                                                        consecutive_fillers * 8
                                                    );
                                                    consecutive_fillers = 0;
                                                }
                                                prev_motion_packet = Some(cmd);
                                                prev_command_was_real = true;
                                                if cmd.last_command {
                                                    log::trace!("Last motion command sent");
                                                }
                                                let _ = self.send(
                                                    TxPackets::MotionCommand(cmd),
                                                    &mut tx_buf,
                                                    Some(Duration::from_millis(6)),
                                                    None,
                                                );
                                                if let Some(h) = handle {
                                                    h.set();
                                                }
                                            } else if state.status_bits().command_received()
                                                && state.status_bits().ready_for_commands()
                                                && !self.itl.1.load(Ordering::SeqCst)
                                            {
                                                if let Some(prev_motion_packet) =
                                                    &prev_motion_packet
                                                {
                                                    let mut cmd = MotionCommandPacket::filler(
                                                        state,
                                                        prev_motion_packet,
                                                        self.send_last_command,
                                                    );
                                                    cmd.seq = state.seq;
                                                    consecutive_fillers += 1;
                                                    let held = prev_motion_packet.position();
                                                    let actual = state.joints_raw();
                                                    log::debug!(
                                                        "STMO queue starved: filler #{} (seq {}, prev_real={}) holding setpoint while robot moves — J1 held={:.4} actual={:.4} (Δ{:.4}), rail held={:.3} actual={:.3} (Δ{:.3})",
                                                        consecutive_fillers,
                                                        cmd.seq,
                                                        prev_command_was_real,
                                                        held[0],
                                                        actual[0],
                                                        actual[0] as f64 - held[0],
                                                        held[6],
                                                        actual[6],
                                                        actual[6] as f64 - held[6],
                                                    );
                                                    prev_command_was_real = false;
                                                    let _ = self.send(
                                                        TxPackets::MotionCommand(cmd),
                                                        &mut tx_buf,
                                                        Some(Duration::from_millis(6)),
                                                        None,
                                                    );
                                                }
                                            } else if self.itl.1.load(Ordering::SeqCst) {
                                                log::trace!(
                                                    "Notifying in the loop that we got a new status"
                                                );
                                                self.itl.0.notify(1);
                                            } else {
                                                log::trace!("No new status received in the loop");
                                            }
                                        }
                                    } else {
                                        log::warn!(
                                            "Received unknown packet: ({}) {:02X?}",
                                            n,
                                            &rx_buf[..n]
                                        );
                                    }
                                }
                                Ok(_) => {
                                    log::warn!("Received empty packet");
                                    break;
                                }
                                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => break,
                                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                                Err(e) => {
                                    log::error!("Error receiving packet: {:?}", e);
                                    break;
                                }
                            }
                        }

                        let req =
                            TxPackets::CommandPositionRequest(CommandPositionRequestPacket {});
                        if self.last_command_position_request_time.elapsed()
                            >= Self::COMMAND_POSITION_RATE
                            && self.protocol_version != 0
                        {
                            let _ =
                                self.send(req, &mut tx_buf, Some(Duration::from_millis(2)), None);
                            self.last_command_position_request_time = Instant::now();
                        }
                    }

                    TOK_WAKER => {
                        // drain commands from driver
                        while let Ok(tx) = self.from_driver.try_recv() {
                            match tx {
                                ToThreadMessage::Start(pkt) => {
                                    let _ = self.send(
                                        TxPackets::Start(pkt),
                                        &mut tx_buf,
                                        Some(Duration::from_millis(24)),
                                        Some(3),
                                    );
                                    let _ = self.send(
                                        TxPackets::VersionNumberRequest(
                                            VersionNumberRequestPacket {},
                                        ),
                                        &mut tx_buf,
                                        Some(Duration::from_millis(24)),
                                        Some(3),
                                    );
                                }
                                ToThreadMessage::MotionCommandDouble(pkt, handle) => {
                                    self.motion_command_queue.push_back((pkt, handle));
                                }
                                ToThreadMessage::Stop(pkt) => {
                                    let _ = self.send(
                                        TxPackets::Stop(pkt),
                                        &mut tx_buf,
                                        Some(Duration::from_millis(24)),
                                        None,
                                    );
                                }
                                ToThreadMessage::ThresholdTableRequest(pkt) => {
                                    let _ = self.send(
                                        TxPackets::ThresholdTableRequest(pkt),
                                        &mut tx_buf,
                                        None,
                                        None,
                                    );
                                    log::info!("Sent ThresholdTableRequest");
                                }
                            }
                        }
                    }

                    _ => {}
                }
            }
        }

        // graceful shutdown if asked to stop while loop breaks
        if thread_handle.should_live() {
            if self.protocol_version == 0 {
                // never started, nothing to do
                log::info!("StreamMotionContext exiting (never started)");
                thread_handle.has_died();
                return;
            }
            let stop_res = self.retry_sending(
                TxPackets::Stop(StopPacket {}),
                Duration::from_millis(24),
                &mut tx_buf,
            );
            if let Err(e) = stop_res {
                log::error!("Error sending stop packet during shutdown: {e:?}");
            }
        }
        log::info!("StreamMotionContext exited");
        thread_handle.has_died();
    }
}

#[allow(clippy::too_many_arguments)]
fn stream_motion_runtime(
    mut thread_handle: ThreadHandle,
    socket: snare::net::UdpSocket,
    thread_config: Option<ThreadConfig>,
    to_driver: Sender<RxPackets>,
    from_driver: Receiver<ToThreadMessage>,
    waker_tx: Sender<Arc<Waker>>,
    itl: Arc<(Event, AtomicBool)>,
    send_last_command: bool,
) -> Result<(), GeneralThreadError> {
    if let Some(cfg) = thread_config {
        cfg.configure_this_thread_print_failure();
    }

    let mut socket = MioUdpSocket::from_std(socket);

    let poll = Poll::new().map_err(|_| GeneralThreadError::FailedToCreatePoll)?;
    poll.registry()
        .register(&mut socket, TOK_SOCKET, Interest::READABLE)
        .map_err(|_| GeneralThreadError::FailedSocketRegistry)?;

    let waker = Arc::new(
        Waker::new(poll.registry(), TOK_WAKER)
            .map_err(|_| GeneralThreadError::FailedWakerCreation)?,
    );
    // send a clone to the driver so API calls can wake the poller
    waker_tx.send(waker.clone())?;
    thread_handle.set_waker_mio(waker);

    log::debug!("Stream motion thread started, entering context loop");

    let context = StreamMotionContext::new(from_driver, to_driver, socket, itl, send_last_command);
    context.context_loop(thread_handle, poll);

    Ok(())
}

#[derive(Debug)]
struct StreamMotionConnection {
    thread_handle: ThreadHandle,
    to_thread: Sender<ToThreadMessage>,
    from_thread: Receiver<RxPackets>,
    is_started: bool,
    err_flag: Arc<AtomicBool>,
    itl: Arc<(Event, AtomicBool)>,
}

/// Driver for FANUC Stream Motion (STMO), a UDP protocol in which the controller
/// requests a position command every interpolation cycle (typically 8ms).
/// A dedicated I/O thread answers each cycle from a queue of motion commands;
/// dropping the driver disconnects, joining that thread.
#[cfg_attr(feature = "py", pyo3::pyclass(str))]
#[derive(Debug)]
pub struct StreamMotionDriver {
    remote_addr: IpAddr,
    send_last_command: bool,
    connection: Option<StreamMotionConnection>,
    cached_movement_limits: Option<JointMovementLimits>,
    rx_storage: RxStorage,
}

impl StreamMotionDriver {
    #[inline]
    fn send_packet(&self, tx: ToThreadMessage) {
        if let Some(conn) = &self.connection {
            if let Err(e) = conn.to_thread.send(tx) {
                log::error!("Error sending packet to thread: {:?}", e);
            }
            let _ = conn.thread_handle.wake();
        }
    }
}

#[cfg(feature = "py")]
type DriverResult<T> = pyo3::PyResult<T>;
#[cfg(not(feature = "py"))]
type DriverResult<T> = Result<T, StreamMotionError>;

#[cfg_mixin(feature = "py")]
#[cfg_attr(feature = "py", pyo3::pymethods)]
impl StreamMotionDriver {
    /// Creates a driver targeting the controller at `addr`. `send_last_command`
    /// sets the last-command flag on filler packets, ending the stream once the
    /// command queue runs dry instead of holding position indefinitely.
    ///
    /// # Errors
    /// `addr` is not a valid IP address.
    #[cfg(on)]
    #[on(pyo3(signature = (addr, send_last_command = false)))]
    #[on(new)]
    pub fn new(addr: Bound<PyAny>, send_last_command: bool) -> DriverResult<Self> {
        let addr = addr.extract::<IpAddr>()?;

        Ok(Self {
            remote_addr: addr,
            send_last_command,
            connection: None,
            cached_movement_limits: None,
            rx_storage: RxStorage::new(),
        })
    }

    /// Creates a driver targeting the controller at `remote_addr`; `send_last_command`
    /// sets the last-command flag on filler packets, ending the stream once the
    /// command queue runs dry instead of holding position indefinitely.
    #[cfg(off)]
    pub fn new<T: Into<IpAddr>>(remote_addr: T, send_last_command: bool) -> Self {
        let remote_addr = remote_addr.into();
        Self {
            remote_addr,
            send_last_command,
            connection: None,
            cached_movement_limits: None,
            rx_storage: RxStorage::new(),
        }
    }

    /// Returns the controller's IP address as a string.
    #[on(pyo3(signature = ()))]
    pub fn get_remote_addr(&self) -> String {
        self.remote_addr.to_string()
    }

    /// Drains packets received by the I/O thread into the driver's internal buffers.
    pub fn refresh(&mut self) {
        let connection = match &self.connection {
            Some(c) => c,
            None => return,
        };
        while let Ok(pkt) = connection.from_thread.try_recv() {
            match pkt {
                RxPackets::RobotStatus(state) => self.rx_storage.status.push_back(state),
                RxPackets::ThresholdTableResponse(threshold) => {
                    self.rx_storage.threshold_table.push_back(threshold)
                }
                RxPackets::CommandPositionResponse(cmd_pos) => {
                    self.rx_storage.command_position.push_back(cmd_pos)
                }
                _ => {}
            }
        }
        self.rx_storage.prune();
    }

    /// Queues motion commands; the I/O thread sends one per controller cycle.
    /// The returned handle is set once the whole batch has been sent.
    ///
    /// # Errors
    /// [`StreamMotionError::NotConnected`] or [`StreamMotionError::NotStarted`] if
    /// [`connect`](Self::connect) and [`start`](Self::start) have not succeeded.
    pub fn command_motion(
        &mut self,
        mut motions: Vec<MotionCommandPacket>,
    ) -> DriverResult<StmoHandle> {
        if self.connection.is_none() {
            return Err(StreamMotionError::NotConnected).map_err(Into::into);
        }
        if !self.is_started() {
            return Err(StreamMotionError::NotStarted).map_err(Into::into);
        }
        let handle = StmoHandle::new();
        if motions.is_empty() {
            handle.set();
            return Ok(handle);
        }
        motions.reverse();
        self.send_packet(ToThreadMessage::MotionCommandDouble(
            MaybeMany::Many(motions),
            Some(handle.clone()),
        ));
        self.refresh();
        Ok(handle)
    }

    pub(crate) fn command_motion_single(
        &mut self,
        motion: MotionCommandPacket,
    ) -> DriverResult<()> {
        if self.connection.is_none() {
            return Err(StreamMotionError::NotConnected).map_err(Into::into);
        }
        if !self.is_started() {
            return Err(StreamMotionError::NotStarted).map_err(Into::into);
        }
        self.send_packet(ToThreadMessage::MotionCommandDouble(
            MaybeMany::One(motion),
            None,
        ));
        self.refresh();
        Ok(())
    }

    /// Sends a stop packet, halting the stream on the controller side.
    pub fn stop(&mut self) {
        self.send_packet(ToThreadMessage::Stop(StopPacket {}));
        self.refresh();
    }

    /// Binds a local UDP socket to the controller's Stream Motion port (60015)
    /// and spawns the I/O thread. No-op if already connected.
    ///
    /// # Errors
    /// I/O failure binding or connecting the socket, or failure to spawn the I/O thread.
    #[on(pyo3(signature = (thread_config=None)))]
    pub fn connect(&mut self, thread_config: Option<ThreadConfig>) -> DriverResult<()> {
        log::info!(
            "Attempting to connect StreamMotionDriver to {}",
            self.remote_addr
        );
        if let Some(conn) = &self.connection
            && conn.thread_handle.is_alive()
        {
            return Ok(());
        }
        let port = openport::pick_unused_port(57000..60000).unwrap_or(60000);
        let local_addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), port);
        let socket = snare::net::UdpSocket::bind(local_addr).map_err(StreamMotionError::from)?;
        socket
            .connect(SocketAddr::new(self.remote_addr, 60015))
            .map_err(StreamMotionError::from)?;
        socket
            .set_nonblocking(true)
            .map_err(StreamMotionError::from)?;

        let (to_thread, from_driver) = flume::unbounded();
        let (to_driver, from_thread) = flume::unbounded();

        let mut thread_handle = ThreadHandle::new();
        let thread_handle_mv = thread_handle.to_pass_in();

        let local_err_flag = Arc::new(AtomicBool::new(false));
        let thread_err_flag = local_err_flag.clone();

        let itl = Arc::new((Event::new(), AtomicBool::new(false)));
        let thread_itl = itl.clone();

        let (waker_tx, waker_rx) = flume::bounded(1);

        let send_last_command = self.send_last_command;

        let thread = snare::thread::Builder::new()
            .name("fanuc-stmo-runner".to_string())
            .spawn(move || {
                if let Err(e) = stream_motion_runtime(
                    thread_handle_mv,
                    socket,
                    thread_config,
                    to_driver,
                    from_driver,
                    waker_tx,
                    thread_itl,
                    send_last_command,
                ) {
                    log::error!("Stream motion thread error: {:?}", e);
                    thread_err_flag.store(true, std::sync::atomic::Ordering::SeqCst);
                }
            })?;

        let thread_waker = waker_rx
            .recv()
            .map_err(|_| StreamMotionError::NotConnected)?;
        thread_handle.set_waker_mio(thread_waker);
        thread_handle.set_handle(thread);

        self.connection = Some(StreamMotionConnection {
            thread_handle,
            to_thread,
            from_thread,
            is_started: false,
            err_flag: local_err_flag,
            itl,
        });

        log::info!("StreamMotionDriver connected to {}", self.remote_addr);

        Ok(())
    }

    /// Returns `true` if the I/O thread exited with an error.
    pub fn has_connection_errored(&self) -> bool {
        if let Some(conn) = &self.connection {
            conn.err_flag.load(std::sync::atomic::Ordering::SeqCst)
        } else {
            false
        }
    }

    /// Sends the start packet and waits for the controller's version response.
    ///
    /// # Errors
    /// [`StreamMotionError::NotConnected`] if [`connect`](Self::connect) has not succeeded;
    /// [`StreamMotionError::Timeout`] if no version response arrives within `timeout_secs`.
    #[on(pyo3(signature = (timeout_secs=2.0)))]
    pub fn start(&mut self, timeout_secs: f32) -> DriverResult<()> {
        let timeout = Duration::from_secs_f32(timeout_secs);
        let start_time = Instant::now();
        let end_time = start_time + timeout;
        if let Some(conn) = &self.connection {
            self.send_packet(ToThreadMessage::Start(StartPacket {}));
            let mut started = false;
            while start_time.elapsed() < timeout {
                let remaining = end_time.saturating_duration_since(Instant::now());
                if let Ok(RxPackets::VersionNumberResponse(_)) =
                    conn.from_thread.recv_timeout(remaining)
                {
                    started = true;
                }
            }
            if !started {
                log::error!(
                    "STMO start timed out after {:.1}s waiting for version response",
                    timeout.as_secs_f32()
                );
                Err(StreamMotionError::Timeout)?;
            }
        } else {
            Err(StreamMotionError::NotConnected)?;
        };
        if let Some(conn) = &mut self.connection {
            conn.is_started = true;
        }
        log::info!("StreamMotionDriver started on {}", self.remote_addr);
        Ok(())
    }

    /// Stops the stream and joins the I/O thread, blocking until it exits.
    /// Called automatically on drop.
    pub fn disconnect(&mut self) {
        if let Some(conn) = self.connection.take() {
            log::info!("StreamMotionDriver disconnecting from {}", self.remote_addr);
            let _ = conn.to_thread.send(ToThreadMessage::Stop(StopPacket {}));
            conn.thread_handle.join();
            log::info!("StreamMotionDriver disconnected from {}", self.remote_addr);
        }
        self.rx_storage.clear();
    }

    /// Returns `true` if the I/O thread is alive.
    pub fn is_connected(&self) -> bool {
        if let Some(conn) = &self.connection {
            conn.thread_handle.is_alive()
        } else {
            false
        }
    }

    /// Returns `true` once [`start`](Self::start) has completed on the current connection.
    pub fn is_started(&self) -> bool {
        if let Some(conn) = &self.connection {
            conn.is_started
        } else {
            false
        }
    }

    /// Requests the per-axis velocity, acceleration, and jerk threshold tables, blocking
    /// until all are received. `extra_axis` is the number of axes beyond the standard six.
    /// Results are cached after the first successful fetch.
    ///
    /// # Errors
    /// [`StreamMotionError::NotConnected`] or [`StreamMotionError::NotStarted`] before
    /// [`connect`](Self::connect) and [`start`](Self::start), or if the connection drops
    /// mid-fetch; [`StreamMotionError::JointDataSizeError`] if `extra_axis > 3`.
    #[on(pyo3(signature = (extra_axis=0)))]
    #[allow(clippy::needless_range_loop)]
    pub fn fetch_movement_limits(&mut self, extra_axis: u8) -> DriverResult<JointMovementLimits> {
        if !self.is_connected() {
            return Err(StreamMotionError::NotConnected).map_err(Into::into);
        }
        if !self.is_started() {
            return Err(StreamMotionError::NotStarted).map_err(Into::into);
        }
        if let Some(cached) = self.cached_movement_limits {
            return Ok(cached);
        }
        if extra_axis > 3 {
            return Err(StreamMotionError::JointDataSizeError(JointDataSizeError(9)))
                .map_err(Into::into);
        }

        let axis_cnt = 6 + extra_axis as usize;

        let mut seen = vec![[false; 3]; axis_cnt];
        let mut limits = JointMovementLimits::default();

        let mut last_send = Instant::now()
            .checked_sub(Duration::from_millis(50))
            .unwrap_or_else(Instant::now);

        let all_filled = |seen: &Vec<[bool; 3]>| seen.iter().flatten().all(|&b| b);

        while !all_filled(&seen) && self.is_connected() {
            if last_send.elapsed() >= Duration::from_millis(48) {
                for joint_idx in 0..axis_cnt {
                    for deriv_idx in 0..3 {
                        if !seen[joint_idx][deriv_idx] {
                            let req = ThresholdTableRequestPacket::try_from((
                                joint_idx as u32 + 1,
                                deriv_idx as u32,
                            ));
                            match req {
                                Ok(r) => {
                                    self.send_packet(ToThreadMessage::ThresholdTableRequest(r))
                                }
                                Err(e) => log::error!(
                                    "Invalid ThresholdTableRequestPacket parameters: {:?}",
                                    e
                                ),
                            }
                            std::thread::sleep(Duration::from_millis(24));
                        }
                    }
                }
                last_send = Instant::now();
            }

            self.refresh();

            while let Some(pkt) = self.rx_storage.threshold_table.pop_front() {
                log::debug!(
                    "Received movement limit: axis {}, type {}, vmax {}",
                    pkt.axis_number,
                    pkt.limit_type,
                    pkt.vmax,
                );
                let axis = pkt.axis_number as usize - 1;
                let deriv = pkt.limit_type as usize;

                if axis < axis_cnt && deriv < 3 && !seen[axis][deriv] {
                    let entry = &mut limits.joints[axis];

                    // set vmax once (first response wins)
                    if limits.vmax == 0 {
                        limits.vmax = pkt.vmax;
                    }

                    let cons = AxisMotionConstraint {
                        no_payload: pkt.no_payload,
                        max_payload: pkt.max_payload,
                    };

                    if entry.is_none() {
                        *entry = Some(JointMovementLimit::default());
                    }

                    if let Some(entry) = entry {
                        match deriv {
                            0 => entry.velocity = cons,
                            1 => entry.acceleration = cons,
                            2 => entry.jerk = cons,
                            _ => {}
                        }
                    }
                    seen[axis][deriv] = true;
                }
            }

            std::thread::sleep(Duration::from_millis(25));
        }

        if self.is_connected() && all_filled(&seen) {
            self.cached_movement_limits = Some(limits);
            Ok(limits)
        } else {
            Err(StreamMotionError::NotConnected).map_err(Into::into)
        }
    }

    /// Drains and returns all buffered robot status packets.
    pub fn pull_states(&mut self) -> Vec<RobotStatusPacket> {
        self.refresh();
        self.rx_storage.status.drain(..).collect()
    }

    /// Drains and returns all buffered command position packets.
    pub fn pull_command_positions(&mut self) -> Vec<CommandPositionResponsePacket> {
        self.refresh();
        self.rx_storage.command_position.drain(..).collect()
    }

    /// Blocks until a command position packet arrives, or returns `None` after `timeout_secs`.
    #[on(pyo3(signature = (timeout_secs = 0.2)))]
    pub fn wait_for_command_position(
        &mut self,
        timeout_secs: f64,
    ) -> Option<CommandPositionResponsePacket> {
        let start = Instant::now();
        while start.elapsed() < Duration::from_secs_f64(timeout_secs) {
            self.refresh();
            if let Some(pkt) = self.rx_storage.command_position.pop_front() {
                return Some(pkt);
            }
            std::thread::sleep(Duration::from_millis(1));
        }
        None
    }
}

impl Drop for StreamMotionDriver {
    fn drop(&mut self) {
        self.disconnect();
    }
}

impl std::fmt::Display for StreamMotionDriver {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let opening = if cfg!(feature = "py") { "(" } else { "{" };
        let closing = if cfg!(feature = "py") { ")" } else { "}" };
        write!(
            f,
            "StreamMotionDriver{}remote_addr: {}, connected: {}{}",
            opening,
            self.remote_addr,
            self.is_connected(),
            closing
        )
    }
}

/// A cycle-by-cycle control session: suspends the I/O thread's automatic filler
/// replies so the caller can answer each robot status with [`send_command`](Self::send_command).
/// Filler replies resume on drop.
#[derive(Debug)]
pub struct StmoControlLoop<'a> {
    driver: &'a mut StreamMotionDriver,
}

impl<'a> StmoControlLoop<'a> {
    /// Begins a control session on the given driver.
    ///
    /// # Errors
    /// [`StreamMotionError::NotConnected`] if the driver has no live connection.
    pub fn try_new(driver: &'a mut StreamMotionDriver) -> Result<Self, StreamMotionError> {
        if let Some(cnx) = &mut driver.connection {
            cnx.itl.1.store(true, Ordering::SeqCst);
            Ok(Self { driver })
        } else {
            Err(StreamMotionError::NotConnected)
        }
    }

    /// Blocks until the next robot status packet arrives.
    ///
    /// # Errors
    /// [`StreamMotionError::Timeout`] if no status arrives within `timeout`;
    /// [`StreamMotionError::NotConnected`] or [`StreamMotionError::NotStarted`]
    /// if the connection or session is gone.
    pub fn wait_for_status(
        &mut self,
        timeout: Duration,
    ) -> Result<RobotStatusPacket, StreamMotionError> {
        // Register the listener and verify ITL state before draining pending
        // packets — otherwise a status that arrives between refresh() and
        // listen() fires notify() with no listeners and is silently lost.
        let listener = match &self.driver.connection {
            Some(cnx) => {
                if !cnx.itl.1.load(Ordering::SeqCst) {
                    return Err(StreamMotionError::NotStarted);
                }
                cnx.itl.0.listen()
            }
            None => return Err(StreamMotionError::NotConnected),
        };
        self.driver.refresh();
        if let Some(pkt) = self.driver.rx_storage.status.pop_back() {
            return Ok(pkt);
        }
        if listener.wait_timeout(timeout).is_some() {
            self.driver.refresh();
            if let Some(pkt) = self.driver.rx_storage.status.pop_back() {
                return Ok(pkt);
            }
        }
        Err(StreamMotionError::Timeout)
    }

    /// Sends a single motion command in reply to the most recent status.
    ///
    /// # Errors
    /// [`StreamMotionError::NotConnected`] or [`StreamMotionError::NotStarted`] if the
    /// driver is no longer connected and started.
    #[inline]
    pub fn send_command(&mut self, motion: MotionCommandPacket) -> DriverResult<()> {
        self.driver
            .command_motion_single(motion)
            .map_err(Into::into)
    }
}

impl Drop for StmoControlLoop<'_> {
    fn drop(&mut self) {
        if let Some(cnx) = &mut self.driver.connection {
            cnx.itl.1.store(false, Ordering::SeqCst);
        }
    }
}

impl StreamMotionDriver {
    /// Begins an [`StmoControlLoop`] session on this driver.
    ///
    /// # Errors
    /// [`StreamMotionError::NotConnected`] if the driver has no live connection.
    pub fn control_loop(&mut self) -> Result<StmoControlLoop<'_>, StreamMotionError> {
        StmoControlLoop::try_new(self)
    }
}

/// Python bindings for the STMO driver.
#[cfg(feature = "py")]
pub mod py {
    use crate::stmo::types::JointMovementLimit;

    use super::*;

    /// Python context-manager counterpart of [`StmoControlLoop`].
    #[derive(Debug)]
    #[pyclass(name = "StmoControlLoop")]
    pub struct PyStmoControlLoop {
        inner: Py<StreamMotionDriver>,
    }

    #[pymethods]
    impl PyStmoControlLoop {
        fn __enter__<'p>(slf: PyRef<'p, Self>, py: Python<'p>) -> PyResult<PyRef<'p, Self>> {
            if let Some(cnx) = &mut slf.inner.borrow_mut(py).connection {
                cnx.itl.1.store(true, Ordering::SeqCst);
            } else {
                return Err(StreamMotionError::NotConnected.into());
            }
            Ok(slf)
        }

        fn __exit__<'a>(
            &mut self,
            py: Python<'a>,
            _exc_type: Bound<'a, PyAny>,
            _exc_value: Bound<'a, PyAny>,
            _traceback: Bound<'a, PyAny>,
        ) -> PyResult<()> {
            if let Some(cnx) = &mut self.inner.borrow_mut(py).connection {
                cnx.itl.1.store(false, Ordering::SeqCst);
            }
            Ok(())
        }

        /// Blocks (GIL released) until the next robot status packet arrives.
        ///
        /// # Errors
        /// Timeout if no status arrives within `timeout_secs`; not-connected or
        /// not-started if the connection or session is gone.
        pub fn wait_for_status(
            &mut self,
            py: Python<'_>,
            timeout_secs: f32,
        ) -> PyResult<RobotStatusPacket> {
            let timeout = Duration::from_secs_f32(timeout_secs);

            // Clone the shared itl Arc and register the listener BEFORE
            // refresh()/drain — otherwise a status that arrives between
            // refresh() and listen() fires notify() with no listeners and is
            // silently lost (event_listener doesn't buffer for unsubscribed
            // listeners). The borrow is also dropped before blocking so other
            // stmo_driver methods can run during the timeout window.
            let listener = {
                let driver = self.inner.borrow(py);
                match &driver.connection {
                    Some(cnx) => {
                        if !cnx.itl.1.load(Ordering::SeqCst) {
                            return Err(StreamMotionError::NotStarted.into());
                        }
                        cnx.itl.0.listen()
                    }
                    None => return Err(StreamMotionError::NotConnected.into()),
                }
            };

            // Drain any status that arrived before we registered the listener.
            {
                let mut driver = self.inner.borrow_mut(py);
                driver.refresh();
                if let Some(pkt) = driver.rx_storage.status.pop_back() {
                    return Ok(pkt);
                }
            }

            // Wait for the next status notification with the GIL released so
            // other Python threads (and other stmo_driver methods) can run.
            let woke = py.detach(|| listener.wait_timeout(timeout).is_some());

            if woke {
                let mut driver = self.inner.borrow_mut(py);
                driver.refresh();
                if let Some(pkt) = driver.rx_storage.status.pop_back() {
                    return Ok(pkt);
                }
            }
            Err(StreamMotionError::Timeout.into())
        }

        /// Sends a single motion command in reply to the most recent status.
        ///
        /// # Errors
        /// Not-connected or not-started if used outside an entered context manager.
        pub fn send_command(
            &mut self,
            py: Python<'_>,
            motion: MotionCommandPacket,
        ) -> PyResult<()> {
            let mut driver = self.inner.borrow_mut(py);
            if let Some(cnx) = &mut driver.connection {
                if !cnx.itl.1.load(Ordering::SeqCst) {
                    return Err(StreamMotionError::NotStarted.into());
                }
                driver.command_motion_single(motion)
            } else {
                Err(StreamMotionError::NotConnected.into())
            }
        }
    }

    #[pymethods]
    impl StreamMotionDriver {
        /// Returns a control-loop context manager for this driver.
        #[pyo3(name = "control_loop")]
        pub fn py_control_loop(slf: Bound<'_, StreamMotionDriver>) -> PyResult<PyStmoControlLoop> {
            Ok(PyStmoControlLoop {
                inner: slf.unbind(),
            })
        }
    }

    /// Registers the STMO driver classes on the given Python module.
    pub fn register(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
        parent_module.add_class::<AxisMotionConstraint>()?;
        parent_module.add_class::<JointMovementLimit>()?;
        parent_module.add_class::<JointMovementLimits>()?;
        parent_module.add_class::<StreamMotionDriver>()?;
        parent_module.add_class::<PyStmoControlLoop>()?;

        Ok(())
    }
}