bcms 0.1.12

use crate::data::IoTData;
use crate::Result;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use chrono::Utc;

/// In-memory data store for IoT data with time-based filtering
/// Equivalent to Python data_store.py
#[derive(Debug)]
pub struct DataStore {
    data: Arc<RwLock<HashMap<String, Vec<IoTData>>>>,
}

impl DataStore {
    pub fn new() -> Self {
        Self {
            data: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Add data to the store
    pub async fn add_data(&self, data: IoTData) -> Result<()> {
        let mut store = self.data.write().await;
        let address = data.address().to_string();
        
        store.entry(address)
            .or_insert_with(Vec::new)
            .push(data);
        
        Ok(())
    }

    /// Add multiple data points to the store
    pub async fn add_data_batch(&self, data_batch: Vec<IoTData>) -> Result<()> {
        let mut store = self.data.write().await;
        
        for data in data_batch {
            let address = data.address().to_string();
            store.entry(address)
                .or_insert_with(Vec::new)
                .push(data);
        }
        
        Ok(())
    }

    /// Get all data for a specific device
    pub async fn get_device_data(&self, address: &str) -> Vec<IoTData> {
        let store = self.data.read().await;
        store.get(address)
            .cloned()
            .unwrap_or_default()
    }

    /// Get data for a device since a specific timestamp
    pub async fn get_device_data_since(&self, address: &str, since_timestamp: i64) -> Vec<IoTData> {
        let store = self.data.read().await;
        if let Some(device_data) = store.get(address) {
            device_data.iter()
                .filter(|data| data.timestamp() > since_timestamp)
                .cloned()
                .collect()
        } else {
            Vec::new()
        }
    }

    /// Get all data from all devices
    pub async fn get_all_data(&self) -> Vec<IoTData> {
        let store = self.data.read().await;
        store.values()
            .flat_map(|device_data| device_data.iter())
            .cloned()
            .collect()
    }

    /// Get all data since a specific timestamp
    pub async fn get_all_data_since(&self, since_timestamp: i64) -> Vec<IoTData> {
        let store = self.data.read().await;
        store.values()
            .flat_map(|device_data| device_data.iter())
            .filter(|data| data.timestamp() > since_timestamp)
            .cloned()
            .collect()
    }

    /// Get data for submission (all unsent data)
    /// Returns data grouped by device address
    pub async fn get_data_for_submission(&self) -> HashMap<String, Vec<IoTData>> {
        let store = self.data.read().await;
        store.clone()
    }

    /// Clear old data older than the specified timestamp
    pub async fn clear_old_data(&self, older_than_timestamp: i64) -> Result<usize> {
        let mut store = self.data.write().await;
        let mut removed_count = 0;
        
        for device_data in store.values_mut() {
            let initial_len = device_data.len();
            device_data.retain(|data| data.timestamp() > older_than_timestamp);
            removed_count += initial_len - device_data.len();
        }
        
        // Remove empty device entries
        store.retain(|_, device_data| !device_data.is_empty());
        
        Ok(removed_count)
    }

    /// Clear all data for a specific device
    pub async fn clear_device_data(&self, address: &str) -> Result<usize> {
        let mut store = self.data.write().await;
        if let Some(device_data) = store.remove(address) {
            Ok(device_data.len())
        } else {
            Ok(0)
        }
    }

    /// Clear all data from the store
    pub async fn clear_all_data(&self) -> Result<usize> {
        let mut store = self.data.write().await;
        let total_count = store.values()
            .map(|device_data| device_data.len())
            .sum();
        
        store.clear();
        Ok(total_count)
    }

    /// Get data count for a specific device
    pub async fn get_device_data_count(&self, address: &str) -> usize {
        let store = self.data.read().await;
        store.get(address)
            .map(|device_data| device_data.len())
            .unwrap_or(0)
    }

    /// Get total data count across all devices
    pub async fn get_total_data_count(&self) -> usize {
        let store = self.data.read().await;
        store.values()
            .map(|device_data| device_data.len())
            .sum()
    }

    /// Get list of devices that have data
    pub async fn get_devices_with_data(&self) -> Vec<String> {
        let store = self.data.read().await;
        store.keys().cloned().collect()
    }

    /// Get the latest data point for a device
    pub async fn get_latest_device_data(&self, address: &str) -> Option<IoTData> {
        let store = self.data.read().await;
        if let Some(device_data) = store.get(address) {
            device_data.iter()
                .max_by_key(|data| data.timestamp())
                .cloned()
        } else {
            None
        }
    }

    /// Mark data as submitted by removing it from the store
    /// This is equivalent to the Python version that removes submitted data
    pub async fn mark_data_submitted(&self, submitted_data: &[IoTData]) -> Result<()> {
        let mut store = self.data.write().await;
        
        for submitted in submitted_data {
            if let Some(device_data) = store.get_mut(submitted.address()) {
                // Remove the specific data point
                device_data.retain(|data| {
                    !(data.address() == submitted.address() && 
                      data.timestamp() == submitted.timestamp() &&
                      data.name() == submitted.name())
                });
            }
        }
        
        // Remove empty device entries
        store.retain(|_, device_data| !device_data.is_empty());
        
        Ok(())
    }

    /// Perform cleanup of old data
    /// This should be called periodically to prevent memory growth
    pub async fn cleanup_old_data(&self, max_age_seconds: i64) -> Result<usize> {
        let now = Utc::now().timestamp();
        let cutoff_timestamp = now - max_age_seconds;
        self.clear_old_data(cutoff_timestamp).await
    }
    
    /// Get data since a specific timestamp (for incremental submission)
    pub async fn get_data_since(&self, since_timestamp: i64) -> Vec<IoTData> {
        let store = self.data.read().await;
        store.values()
            .flat_map(|device_data| device_data.iter())
            .filter(|data| data.timestamp() > since_timestamp)
            .cloned()
            .collect()
    }

    /// Get data with per-device timestamp filtering.
    /// Each device is filtered against its own last-submitted timestamp.
    pub async fn get_data_since_per_device(&self, since_timestamps: &HashMap<String, i64>) -> Vec<IoTData> {
        let store = self.data.read().await;
        store.iter()
            .flat_map(|(address, device_data)| {
                let since = since_timestamps.get(address).copied().unwrap_or(0);
                device_data.iter()
                    .filter(move |data| data.timestamp() > since)
                    .cloned()
            })
            .collect()
    }
    
    /// Limit IoT data sample rate (equivalent to Python's limit_iot_data_sample_rate)
    /// Keeps only 1 sample per (device, data_type) combination every N seconds
    pub fn limit_sample_rate(data: Vec<IoTData>, samples_every_seconds: i64) -> Vec<IoTData> {
        if data.is_empty() {
            return data;
        }
        
        use std::collections::HashMap;
        
        // Group data by device address
        let mut data_by_device: HashMap<String, Vec<IoTData>> = HashMap::new();
        for d in data {
            data_by_device
                .entry(d.address().to_string())
                .or_insert_with(Vec::new)
                .push(d);
        }
        
        // Group by device and data type
        let mut data_by_device_and_type: HashMap<String, HashMap<String, Vec<IoTData>>> = HashMap::new();
        for (address, device_data) in data_by_device {
            let mut data_by_type: HashMap<String, Vec<IoTData>> = HashMap::new();
            for d in device_data {
                data_by_type
                    .entry(d.name().to_string())
                    .or_insert_with(Vec::new)
                    .push(d);
            }
            data_by_device_and_type.insert(address, data_by_type);
        }
        
        // Filter data - keep only 1 sample per (device, type) every N seconds
        let mut filtered_data = Vec::new();
        for (_address, data_by_type) in data_by_device_and_type {
            for (_data_type, mut type_data) in data_by_type {
                // Sort by timestamp to process in chronological order
                type_data.sort_by_key(|d| d.timestamp());
                
                let mut last_timestamp: Option<i64> = None;
                for d in type_data {
                    if last_timestamp.is_none() || d.timestamp() - last_timestamp.unwrap() > samples_every_seconds {
                        filtered_data.push(d.clone());
                        last_timestamp = Some(d.timestamp());
                    }
                }
            }
        }
        
        filtered_data
    }
    
    /// Get data count per device for monitoring
    pub async fn get_data_count_per_device(&self) -> HashMap<String, usize> {
        let store = self.data.read().await;
        store.iter()
            .map(|(address, device_data)| (address.clone(), device_data.len()))
            .collect()
    }
    
    /// Enforce memory limits by removing oldest data when limits are exceeded
    pub async fn enforce_limits(&self, max_data_points_per_device: usize, max_total_data_points: usize) -> Result<(usize, usize)> {
        let mut store = self.data.write().await;
        let mut total_removed = 0;
        let mut total_remaining = 0;
        
        // First pass: enforce per-device limits
        for device_data in store.values_mut() {
            if device_data.len() > max_data_points_per_device {
                // Sort by timestamp and keep only the newest data
                device_data.sort_by_key(|d| d.timestamp());
                let to_remove = device_data.len() - max_data_points_per_device;
                device_data.drain(0..to_remove);
                total_removed += to_remove;
            }
            total_remaining += device_data.len();
        }
        
        // Second pass: enforce total limit if still exceeded
        if total_remaining > max_total_data_points {
            // Collect all data with timestamps, sort globally, and keep newest
            let mut all_data_with_device: Vec<(String, IoTData)> = Vec::new();
            for (address, device_data) in store.iter() {
                for data in device_data {
                    all_data_with_device.push((address.clone(), data.clone()));
                }
            }
            
            // Sort by timestamp (oldest first)
            all_data_with_device.sort_by_key(|(_, data)| data.timestamp());
            
            // Remove oldest data to reach limit
            let to_remove_global = total_remaining - max_total_data_points;
            let data_to_remove = &all_data_with_device[0..to_remove_global];
            
            // Remove the selected data from the store
            for (address, data_to_remove) in data_to_remove {
                if let Some(device_data) = store.get_mut(address) {
                    device_data.retain(|data| {
                        !(data.address() == data_to_remove.address() &&
                          data.timestamp() == data_to_remove.timestamp() &&
                          data.name() == data_to_remove.name())
                    });
                }
            }
            
            total_removed += to_remove_global;
            total_remaining = max_total_data_points;
        }
        
        // Remove empty device entries
        store.retain(|_, device_data| !device_data.is_empty());
        
        Ok((total_removed, total_remaining))
    }
    
    /// Get unique data types for a specific device (for supportedDataTypes inference)
    pub async fn get_device_data_types(&self, address: &str) -> Vec<String> {
        let store = self.data.read().await;
        let mut data_types = std::collections::HashSet::new();

        if let Some(device_data) = store.get(address) {
            for data in device_data {
                data_types.insert(data.name().to_string());
            }
        }

        data_types.into_iter().collect()
    }

    /// Get data summary for a specific device (data type counts and latest timestamps)
    pub async fn get_device_data_summary(&self, address: &str) -> HashMap<String, (u32, i64)> {
        let store = self.data.read().await;
        let mut summary = HashMap::new();
        
        if let Some(device_data) = store.get(address) {
            for data in device_data {
                let data_type = data.name().to_string();
                let timestamp = data.timestamp();
                
                match summary.get_mut(&data_type) {
                    Some((count, last_timestamp)) => {
                        *count += 1;
                        if timestamp > *last_timestamp {
                            *last_timestamp = timestamp;
                        }
                    }
                    None => {
                        summary.insert(data_type, (1, timestamp));
                    }
                }
            }
        }
        
        summary
    }
}

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

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

    #[test]
    fn test_rate_limit_keeps_first_discards_rapid_followups() {
        // Simulates the 0/0 then real-data race condition (now fixed by removing 0/0)
        let data = vec![
            helpers::create_blood_pressure_data("device1".into(), 0, 0, 100),
            helpers::create_blood_pressure_data("device1".into(), 120, 80, 101), // 1s later
        ];
        let result = DataStore::limit_sample_rate(data, 2);

        // With 2-second window, only first sample is kept
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].timestamp(), 100);
    }

    #[test]
    fn test_rate_limit_keeps_samples_outside_window() {
        let data = vec![
            helpers::create_blood_pressure_data("device1".into(), 120, 80, 100),
            helpers::create_blood_pressure_data("device1".into(), 125, 85, 103), // 3s later, > 2s window
        ];
        let result = DataStore::limit_sample_rate(data, 2);

        // Both samples should be kept (3s apart > 2s window)
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_rate_limit_independent_per_data_type() {
        // blood_pressure and heart_rate should be rate-limited separately
        let data = vec![
            helpers::create_blood_pressure_data("device1".into(), 120, 80, 100),
            helpers::create_heart_rate_data("device1".into(), 72, 100), // same timestamp, different type
        ];
        let result = DataStore::limit_sample_rate(data, 2);

        // Both kept - different data types
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_rate_limit_independent_per_device() {
        let data = vec![
            helpers::create_blood_pressure_data("device1".into(), 120, 80, 100),
            helpers::create_blood_pressure_data("device2".into(), 130, 85, 100), // same timestamp, different device
        ];
        let result = DataStore::limit_sample_rate(data, 2);

        // Both kept - different devices
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_rate_limit_empty_input() {
        let data: Vec<IoTData> = vec![];
        let result = DataStore::limit_sample_rate(data, 2);
        assert!(result.is_empty());
    }

    #[test]
    fn test_rate_limit_single_sample() {
        let data = vec![
            helpers::create_blood_pressure_data("device1".into(), 120, 80, 100),
        ];
        let result = DataStore::limit_sample_rate(data, 2);
        assert_eq!(result.len(), 1);
    }

    #[test]
    fn test_rate_limit_boundary_exactly_at_window() {
        // Test exact boundary: 2s window, samples exactly 2s apart
        let data = vec![
            helpers::create_blood_pressure_data("device1".into(), 120, 80, 100),
            helpers::create_blood_pressure_data("device1".into(), 125, 85, 102), // exactly 2s later
        ];
        let result = DataStore::limit_sample_rate(data, 2);

        // Second sample filtered (uses strict > not >=)
        assert_eq!(result.len(), 1);
    }

    #[tokio::test]
    async fn test_get_device_data_types_empty_store() {
        let store = DataStore::new();
        let types = store.get_device_data_types("nonexistent").await;
        assert!(types.is_empty());
    }

    #[tokio::test]
    async fn test_get_device_data_types_single_type() {
        let store = DataStore::new();
        store.add_data(helpers::create_blood_pressure_data("device1".into(), 120, 80, 100)).await.unwrap();

        let types = store.get_device_data_types("device1").await;
        assert_eq!(types.len(), 1);
        assert!(types.contains(&"blood_pressure".to_string()));
    }

    #[tokio::test]
    async fn test_get_device_data_types_multiple_types() {
        let store = DataStore::new();
        store.add_data(helpers::create_blood_pressure_data("device1".into(), 120, 80, 100)).await.unwrap();
        store.add_data(helpers::create_heart_rate_data("device1".into(), 72, 101)).await.unwrap();
        store.add_data(helpers::create_blood_pressure_data("device1".into(), 125, 85, 102)).await.unwrap();

        let types = store.get_device_data_types("device1").await;
        assert_eq!(types.len(), 2);
        assert!(types.contains(&"blood_pressure".to_string()));
        assert!(types.contains(&"heart_rate".to_string()));
    }

    #[tokio::test]
    async fn test_get_device_data_types_different_devices() {
        let store = DataStore::new();
        store.add_data(helpers::create_blood_pressure_data("device1".into(), 120, 80, 100)).await.unwrap();
        store.add_data(helpers::create_heart_rate_data("device2".into(), 72, 100)).await.unwrap();

        let types1 = store.get_device_data_types("device1").await;
        let types2 = store.get_device_data_types("device2").await;

        assert_eq!(types1.len(), 1);
        assert!(types1.contains(&"blood_pressure".to_string()));

        assert_eq!(types2.len(), 1);
        assert!(types2.contains(&"heart_rate".to_string()));
    }

    #[tokio::test]
    async fn test_get_data_since_strict_inequality() {
        let store = DataStore::new();
        store.add_data(helpers::create_blood_pressure_data("device1".into(), 120, 80, 100)).await.unwrap();

        // Exact timestamp should NOT be included (strict >)
        let result = store.get_data_since(100).await;
        assert!(result.is_empty());

        // One less should include the data
        let result = store.get_data_since(99).await;
        assert_eq!(result.len(), 1);
    }
}