mecha10-core 0.6.2

//! Priority queue for message processing
//!
//! Enables quality-of-service (QoS) by prioritizing critical messages over
//! low-priority ones. Useful for safety-critical systems where some messages
//! (e.g., emergency stop) must be processed before others (e.g., telemetry).
//!
//! # Features
//!
//! - **Priority Levels**: Define message importance (Critical, High, Normal, Low)
//! - **Fair Processing**: Round-robin within same priority level
//! - **Starvation Prevention**: Ensures low-priority messages eventually process
//! - **Backpressure Support**: Integrates with bounded channels
//! - **Zero-Copy**: Messages moved, not copied
//!
//! # Example
//!
//! ```rust
//! use mecha10::prelude::*;
//! use mecha10::priority_queue::{PriorityReceiver, Priority};
//!
//! # async fn example(ctx: &Context) -> Result<()> {
//! // Create priority receiver
//! let mut prio_rx = PriorityReceiver::new();
//!
//! // Subscribe to multiple topics with different priorities
//! let emergency = ctx.subscribe("/safety/emergency").await?;
//! let commands = ctx.subscribe("/motor/command").await?;
//! let telemetry = ctx.subscribe("/sensor/telemetry").await?;
//!
//! // Add with priorities
//! prio_rx.add_source(emergency, Priority::Critical);
//! prio_rx.add_source(commands, Priority::High);
//! prio_rx.add_source(telemetry, Priority::Low);
//!
//! // Process messages in priority order
//! while let Some((msg, priority)) = prio_rx.recv().await {
//!     match priority {
//!         Priority::Critical => handle_emergency(msg).await?,
//!         Priority::High => handle_command(msg).await?,
//!         _ => handle_telemetry(msg).await?,
//!     }
//! }
//! # Ok(())
//! # }
//! ```

use crate::context::Receiver;
use std::collections::VecDeque;

/// Message priority levels
///
/// Higher priority messages are processed before lower priority ones.
/// Within the same priority level, messages are processed FIFO.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
pub enum Priority {
    /// Lowest priority - background tasks, non-critical telemetry
    Low = 0,
    /// Normal priority - standard operations
    #[default]
    Normal = 1,
    /// High priority - time-sensitive commands
    High = 2,
    /// Highest priority - safety-critical, emergency stop
    Critical = 3,
}

/// Configuration for priority queue behavior
#[derive(Debug, Clone)]
pub struct PriorityQueueConfig {
    /// Maximum messages to buffer per priority level
    /// None = unbounded (risk of memory growth)
    pub max_per_priority: Option<usize>,

    /// Enable starvation prevention
    /// If true, occasionally process lower priority messages even when
    /// higher priority messages are available
    pub prevent_starvation: bool,

    /// Starvation threshold: after processing N high-priority messages,
    /// process one lower-priority message (if available)
    pub starvation_threshold: usize,
}

impl Default for PriorityQueueConfig {
    fn default() -> Self {
        Self {
            max_per_priority: Some(1000),
            prevent_starvation: true,
            starvation_threshold: 100,
        }
    }
}

/// Priority receiver that processes messages by priority
///
/// Combines multiple receivers and delivers messages in priority order.
/// Messages with higher priority are processed before lower priority ones.
pub struct PriorityReceiver<T> {
    sources: Vec<PrioritySource<T>>,
    config: PriorityQueueConfig,
    high_priority_count: usize,
}

struct PrioritySource<T> {
    receiver: Receiver<T>,
    priority: Priority,
    buffer: VecDeque<T>,
}

impl<T> PriorityReceiver<T>
where
    T: Send + 'static,
{
    /// Create a new priority receiver with default configuration
    pub fn new() -> Self {
        Self::with_config(PriorityQueueConfig::default())
    }

    /// Create a new priority receiver with custom configuration
    pub fn with_config(config: PriorityQueueConfig) -> Self {
        Self {
            sources: Vec::new(),
            config,
            high_priority_count: 0,
        }
    }

    /// Add a message source with the specified priority
    ///
    /// # Example
    ///
    /// ```rust
    /// # use mecha10::prelude::*;
    /// # use mecha10::priority_queue::{PriorityReceiver, Priority};
    /// # async fn example(ctx: &Context) -> Result<()> {
    /// let mut prio_rx = PriorityReceiver::new();
    ///
    /// let emergency = ctx.subscribe("/safety/emergency").await?;
    /// prio_rx.add_source(emergency, Priority::Critical);
    ///
    /// let telemetry = ctx.subscribe("/sensor/telemetry").await?;
    /// prio_rx.add_source(telemetry, Priority::Low);
    /// # Ok(())
    /// # }
    /// ```
    pub fn add_source(&mut self, receiver: Receiver<T>, priority: Priority) {
        self.sources.push(PrioritySource {
            receiver,
            priority,
            buffer: VecDeque::new(),
        });
    }

    /// Receive the next message in priority order
    ///
    /// Returns `None` when all sources are closed.
    /// Returns `Some((message, priority))` with the highest priority message available.
    ///
    /// # Starvation Prevention
    ///
    /// If configured, this method will occasionally process lower-priority messages
    /// even when higher-priority messages are available, to prevent starvation.
    pub async fn recv(&mut self) -> Option<(T, Priority)> {
        if self.sources.is_empty() {
            return None;
        }

        loop {
            // Try to fill buffers from all sources (non-blocking)
            self.try_fill_buffers();

            // Check if we should process a lower-priority message for starvation prevention
            if self.config.prevent_starvation && self.high_priority_count >= self.config.starvation_threshold {
                // Try to get a lower-priority message
                if let Some((msg, prio)) = self.take_lower_priority_message() {
                    self.high_priority_count = 0;
                    return Some((msg, prio));
                }
            }

            // Get highest priority message from buffers
            if let Some((msg, prio)) = self.take_highest_priority_message() {
                if prio >= Priority::High {
                    self.high_priority_count += 1;
                }
                return Some((msg, prio));
            }

            // No messages in buffers, wait for any source to have data
            if !self.wait_for_message().await {
                // All sources closed
                return None;
            }
        }
    }

    /// Try to receive messages from all sources without blocking
    fn try_fill_buffers(&mut self) {
        for source in &mut self.sources {
            // Respect buffer size limits
            let max_size = self.config.max_per_priority.unwrap_or(usize::MAX);

            while source.buffer.len() < max_size {
                match source.receiver.try_recv() {
                    Ok(msg) => source.buffer.push_back(msg),
                    Err(_) => break, // Empty or closed
                }
            }
        }
    }

    /// Wait for at least one message from any source
    async fn wait_for_message(&mut self) -> bool {
        if self.sources.is_empty() {
            return false;
        }

        // Use tokio::select! to wait on all sources simultaneously
        // This is efficient and doesn't poll
        loop {
            for source in &mut self.sources {
                // Try each source
                if let Some(msg) = source.receiver.recv().await {
                    source.buffer.push_back(msg);
                    return true;
                }
            }

            // Remove closed sources
            self.sources.retain(|s| !s.buffer.is_empty() || !s.receiver.is_empty());

            if self.sources.is_empty() {
                return false;
            }
        }
    }

    /// Take the highest priority message from buffers
    fn take_highest_priority_message(&mut self) -> Option<(T, Priority)> {
        // Find source with highest priority that has messages.
        // Seed with `None` rather than a sentinel `Priority` value: a sentinel
        // collides with real priority levels (e.g. `Priority::Low`), which would
        // make a strict `>` comparison never select sources sitting exactly at
        // the sentinel even though they're the only messages available.
        let mut best: Option<(usize, Priority)> = None;

        for (idx, source) in self.sources.iter().enumerate() {
            if source.buffer.is_empty() {
                continue;
            }
            match best {
                None => best = Some((idx, source.priority)),
                Some((_, best_priority)) if source.priority > best_priority => {
                    best = Some((idx, source.priority));
                }
                _ => {}
            }
        }

        best.and_then(|(idx, _)| {
            self.sources[idx]
                .buffer
                .pop_front()
                .map(|msg| (msg, self.sources[idx].priority))
        })
    }

    /// Take a lower-priority message for starvation prevention
    fn take_lower_priority_message(&mut self) -> Option<(T, Priority)> {
        // Find source with lowest priority that has messages. See
        // `take_highest_priority_message` for why we seed with `None` instead
        // of a sentinel `Priority` value.
        let mut best: Option<(usize, Priority)> = None;

        for (idx, source) in self.sources.iter().enumerate() {
            if source.buffer.is_empty() {
                continue;
            }
            match best {
                None => best = Some((idx, source.priority)),
                Some((_, best_priority)) if source.priority < best_priority => {
                    best = Some((idx, source.priority));
                }
                _ => {}
            }
        }

        best.and_then(|(idx, _)| {
            self.sources[idx]
                .buffer
                .pop_front()
                .map(|msg| (msg, self.sources[idx].priority))
        })
    }

    /// Get statistics about the priority queue
    pub fn stats(&self) -> PriorityQueueStats {
        let mut counts_by_priority = [0; 4];
        let mut total_buffered = 0;

        for source in &self.sources {
            let idx = source.priority as usize;
            counts_by_priority[idx] += source.buffer.len();
            total_buffered += source.buffer.len();
        }

        PriorityQueueStats {
            total_buffered,
            critical_buffered: counts_by_priority[Priority::Critical as usize],
            high_buffered: counts_by_priority[Priority::High as usize],
            normal_buffered: counts_by_priority[Priority::Normal as usize],
            low_buffered: counts_by_priority[Priority::Low as usize],
            source_count: self.sources.len(),
        }
    }
}

impl<T> Default for PriorityReceiver<T>
where
    T: Send + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_source(priority: Priority, messages: Vec<u32>) -> PrioritySource<u32> {
        let (_tx, rx) = tokio::sync::mpsc::unbounded_channel();
        PrioritySource {
            receiver: Receiver::from(rx),
            priority,
            buffer: messages.into_iter().collect(),
        }
    }

    /// Regression test: when every buffered message is `Priority::Low` (the
    /// sentinel `take_highest_priority_message` used to seed `best_priority`
    /// with), the strict `>` comparison used to never select any of them,
    /// even though they were the only messages present. This starved those
    /// sources forever in `recv()`, since the underlying channel had already
    /// been drained into the buffer via `try_recv` and `wait_for_message`
    /// only awaits *new* channel messages, not buffered ones.
    #[test]
    fn take_highest_priority_message_selects_all_low_messages() {
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.sources.push(make_source(Priority::Low, vec![1, 2, 3]));

        let (msg, priority) = prio_rx
            .take_highest_priority_message()
            .expect("a Low-priority message should still be returned when it's the only one buffered");
        assert_eq!(msg, 1);
        assert_eq!(priority, Priority::Low);

        // Subsequent calls should keep draining the remaining Low messages.
        assert_eq!(prio_rx.take_highest_priority_message(), Some((2, Priority::Low)));
        assert_eq!(prio_rx.take_highest_priority_message(), Some((3, Priority::Low)));
        assert_eq!(prio_rx.take_highest_priority_message(), None);
    }

    /// Regression test: mirror of the above for `take_lower_priority_message`,
    /// whose sentinel was `Priority::Critical` compared with strict `<`.
    #[test]
    fn take_lower_priority_message_selects_all_critical_messages() {
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.sources.push(make_source(Priority::Critical, vec![10, 20]));

        let (msg, priority) = prio_rx
            .take_lower_priority_message()
            .expect("a Critical-priority message should still be returned when it's the only one buffered");
        assert_eq!(msg, 10);
        assert_eq!(priority, Priority::Critical);

        assert_eq!(prio_rx.take_lower_priority_message(), Some((20, Priority::Critical)));
        assert_eq!(prio_rx.take_lower_priority_message(), None);
    }

    /// End-to-end regression test through `recv()`: a single source producing
    /// only `Priority::Low` messages must have those messages delivered
    /// rather than hanging forever waiting on `wait_for_message`.
    #[tokio::test]
    async fn recv_delivers_messages_from_a_low_priority_only_source() {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<u32>();
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.add_source(Receiver::from(rx), Priority::Low);

        tx.send(42).unwrap();
        drop(tx);

        let result = tokio::time::timeout(std::time::Duration::from_secs(5), prio_rx.recv())
            .await
            .expect("recv() should not hang starving a Low-only source");
        assert_eq!(result, Some((42, Priority::Low)));
    }

    /// take_highest_priority_message should still prefer higher priorities
    /// over lower ones when both are present.
    #[test]
    fn take_highest_priority_message_prefers_higher_priority() {
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.sources.push(make_source(Priority::Low, vec![1]));
        prio_rx.sources.push(make_source(Priority::Critical, vec![2]));

        assert_eq!(prio_rx.take_highest_priority_message(), Some((2, Priority::Critical)));
    }
}

/// Statistics about priority queue state
#[derive(Debug, Clone)]
pub struct PriorityQueueStats {
    /// Total messages buffered across all priorities
    pub total_buffered: usize,
    /// Messages buffered at Critical priority
    pub critical_buffered: usize,
    /// Messages buffered at High priority
    pub high_buffered: usize,
    /// Messages buffered at Normal priority
    pub normal_buffered: usize,
    /// Messages buffered at Low priority
    pub low_buffered: usize,
    /// Number of active sources
    pub source_count: usize,
}