rustytorch/crates/specialized/rtx-neuro-lsl/src/service.rs

//! LSL service for managing stream discovery and connections

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

use crate::buffer::StreamBuffer;
use crate::error::{LslError, Result};
use crate::stream::{LslSample, LslStreamInfo};

/// Handle to an active LSL stream
#[derive(Debug, Clone)]
pub struct LslHandle {
    /// Unique identifier for this connection
    pub id: String,
    /// Stream information
    pub info: LslStreamInfo,
    /// Whether currently connected
    pub connected: bool,
    /// Samples received since connection
    pub samples_received: usize,
    /// Samples dropped due to buffer overflow
    pub samples_dropped: usize,
}

/// Internal state for a connected stream
struct StreamConnection {
    handle: LslHandle,
    buffer: Arc<RwLock<StreamBuffer>>,
    // In a real implementation, this would hold the LSL inlet
    // and background task handles
}

/// LSL service for managing stream discovery and connections
pub struct LslService {
    /// Active stream connections
    connections: HashMap<String, StreamConnection>,
    /// Counter for generating unique IDs
    counter: usize,
}

impl LslService {
    /// Create a new LSL service
    pub fn new() -> Self {
        Self {
            connections: HashMap::new(),
            counter: 0,
        }
    }

    /// Generate a unique stream ID
    fn next_id(&mut self) -> String {
        self.counter += 1;
        format!("lsl_{}", self.counter)
    }

    /// Discover available LSL streams
    ///
    /// # Arguments
    /// * `timeout` - Discovery timeout in seconds
    ///
    /// # Returns
    /// List of discovered stream information
    ///
    /// # Note
    /// This is a mock implementation. Real discovery requires native LSL library.
    pub async fn discover_streams(&self, _timeout: f64) -> Result<Vec<LslStreamInfo>> {
        // In a real implementation, this would use lsl::resolve_streams()
        // For now, return an empty list as we don't have native LSL
        Ok(Vec::new())
    }

    /// Connect to a stream by name
    ///
    /// # Arguments
    /// * `stream_name` - Name of the stream to connect to
    /// * `buffer_duration` - Duration of data to buffer in seconds
    ///
    /// # Returns
    /// Handle to the connected stream
    pub async fn connect(
        &mut self,
        stream_info: LslStreamInfo,
        buffer_duration: f64,
    ) -> Result<LslHandle> {
        let id = self.next_id();

        let buffer = StreamBuffer::for_duration(buffer_duration, stream_info.nominal_srate);

        let handle = LslHandle {
            id: id.clone(),
            info: stream_info,
            connected: true,
            samples_received: 0,
            samples_dropped: 0,
        };

        let connection = StreamConnection {
            handle: handle.clone(),
            buffer: Arc::new(RwLock::new(buffer)),
        };

        self.connections.insert(id, connection);

        Ok(handle)
    }

    /// Disconnect from a stream
    pub async fn disconnect(&mut self, stream_id: &str) -> Result<()> {
        self.connections
            .remove(stream_id)
            .ok_or_else(|| LslError::StreamNotFound {
                name: stream_id.to_string(),
            })?;
        Ok(())
    }

    /// Get recent data from a stream
    ///
    /// # Arguments
    /// * `stream_id` - ID of the connected stream
    /// * `duration` - Duration of data to retrieve in seconds
    ///
    /// # Returns
    /// Tuple of (data, timestamps) where data is [samples][channels]
    pub async fn get_data(
        &self,
        stream_id: &str,
        duration: f64,
    ) -> Result<(Vec<Vec<f64>>, Vec<f64>)> {
        let conn = self
            .connections
            .get(stream_id)
            .ok_or_else(|| LslError::StreamNotFound {
                name: stream_id.to_string(),
            })?;

        let buffer = conn.buffer.read().await;
        Ok(buffer.get_recent(duration))
    }

    /// Push a sample to a stream buffer (for simulation/testing)
    pub async fn push_sample(&self, stream_id: &str, sample: LslSample) -> Result<()> {
        let conn = self
            .connections
            .get(stream_id)
            .ok_or_else(|| LslError::StreamNotFound {
                name: stream_id.to_string(),
            })?;

        let mut buffer = conn.buffer.write().await;
        buffer.push(sample);
        Ok(())
    }

    /// Get stream info for a connected stream
    pub fn get_stream_info(&self, stream_id: &str) -> Option<&LslStreamInfo> {
        self.connections.get(stream_id).map(|c| &c.handle.info)
    }

    /// List all active stream connections
    pub fn list_active(&self) -> Vec<LslHandle> {
        self.connections
            .values()
            .map(|c| c.handle.clone())
            .collect()
    }

    /// Check if a stream is connected
    pub fn is_connected(&self, stream_id: &str) -> bool {
        self.connections.contains_key(stream_id)
    }

    /// Get sample count for a stream
    pub async fn get_sample_count(&self, stream_id: &str) -> Result<usize> {
        let conn = self
            .connections
            .get(stream_id)
            .ok_or_else(|| LslError::StreamNotFound {
                name: stream_id.to_string(),
            })?;

        let buffer = conn.buffer.read().await;
        Ok(buffer.total_received())
    }

    /// Clear buffer for a stream
    pub async fn clear_buffer(&self, stream_id: &str) -> Result<()> {
        let conn = self
            .connections
            .get(stream_id)
            .ok_or_else(|| LslError::StreamNotFound {
                name: stream_id.to_string(),
            })?;

        let mut buffer = conn.buffer.write().await;
        buffer.clear();
        Ok(())
    }
}

impl Default for LslService {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[tokio::test]
    async fn test_service_creation() {
        let service = LslService::new();
        assert!(service.list_active().is_empty());
    }

    #[tokio::test]
    async fn test_connect_disconnect() {
        let mut service = LslService::new();

        let info = LslStreamInfo::new("TestEEG", "EEG", 64, 256.0);
        let handle = service.connect(info, 5.0).await.unwrap();

        assert!(service.is_connected(&handle.id));
        assert_eq!(service.list_active().len(), 1);

        service.disconnect(&handle.id).await.unwrap();
        assert!(!service.is_connected(&handle.id));
        assert!(service.list_active().is_empty());
    }

    #[tokio::test]
    async fn test_push_and_get_data() {
        let mut service = LslService::new();

        let info = LslStreamInfo::new("TestEEG", "EEG", 4, 100.0);
        let handle = service.connect(info, 2.0).await.unwrap();

        // Push 100 samples (1 second at 100 Hz)
        for i in 0..100 {
            let sample = LslSample::new(vec![i as f64; 4], i as f64 * 0.01);
            service.push_sample(&handle.id, sample).await.unwrap();
        }

        // Get last 0.5 seconds (50 samples)
        let (data, times) = service.get_data(&handle.id, 0.5).await.unwrap();
        assert_eq!(data.len(), 50);
        assert_eq!(times.len(), 50);
    }

    #[tokio::test]
    async fn test_discover_streams() {
        let service = LslService::new();
        let streams = service.discover_streams(1.0).await.unwrap();
        // Without native LSL, this returns empty
        assert!(streams.is_empty());
    }
}