tool.meshservice/src/store.rs

//! In-memory message store with eviction policies.

use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

use crate::message::ServiceMessage;

/// Configuration for the message store.
#[derive(Debug, Clone)]
pub struct StoreConfig {
    /// Maximum messages per service.
    pub max_per_service: usize,
    /// Maximum messages per sender (per service).
    pub max_per_sender: usize,
    /// Maximum total messages.
    pub max_total: usize,
    /// Prune interval in seconds.
    pub prune_interval_secs: u64,
}

impl Default for StoreConfig {
    fn default() -> Self {
        Self {
            max_per_service: 10_000,
            max_per_sender: 100,
            max_total: 50_000,
            prune_interval_secs: 300,
        }
    }
}

/// A stored message with metadata.
#[derive(Debug, Clone)]
pub struct StoredMessage {
    pub message: ServiceMessage,
    /// Sender's public key (needed for verification).
    pub sender_public_key: [u8; 32],
    /// When we stored this message.
    pub stored_at: u64,
}

/// Generic service message store.
///
/// Organized by service_id, then by sender_address, then by message_id.
pub struct ServiceStore {
    config: StoreConfig,
    /// service_id -> sender_address -> message_id -> StoredMessage
    messages: HashMap<u32, HashMap<[u8; 16], HashMap<[u8; 16], StoredMessage>>>,
    /// Total message count.
    total_count: usize,
    /// Last prune timestamp.
    last_prune: u64,
}

impl ServiceStore {
    /// Create a new store with default config.
    pub fn new() -> Self {
        Self::with_config(StoreConfig::default())
    }

    /// Create with custom config.
    pub fn with_config(config: StoreConfig) -> Self {
        Self {
            config,
            messages: HashMap::new(),
            total_count: 0,
            last_prune: 0,
        }
    }

    /// Store a message, returning true if it was new.
    pub fn store(&mut self, message: ServiceMessage, sender_public_key: [u8; 32]) -> bool {
        // Prune if interval passed
        self.maybe_prune();

        let service_id = message.service_id;
        let sender_address = message.sender_address;
        let message_id = message.id;

        // Check per-service limit and evict if needed
        {
            let service_count: usize = self.messages
                .get(&service_id)
                .map(|s| s.values().map(|m| m.len()).sum())
                .unwrap_or(0);
            if service_count >= self.config.max_per_service {
                self.evict_oldest_in_service(service_id);
            }
        }

        // Check per-sender limit and evict if needed
        {
            let sender_count = self.messages
                .get(&service_id)
                .and_then(|s| s.get(&sender_address))
                .map(|m| m.len())
                .unwrap_or(0);
            if sender_count >= self.config.max_per_sender {
                self.evict_oldest_from_sender(service_id, sender_address);
            }
        }

        // Get or create maps
        let service_map = self.messages.entry(service_id).or_default();
        let sender_map = service_map.entry(sender_address).or_default();

        // Check for existing message
        let is_new_or_update = if let Some(existing) = sender_map.get(&message_id) {
            // Existing: only update if higher sequence
            if message.sequence <= existing.message.sequence {
                return false;
            }
            // This is an update, not a new message
            false
        } else {
            // New message
            true
        };

        let stored_at = now();
        sender_map.insert(
            message_id,
            StoredMessage {
                message,
                sender_public_key,
                stored_at,
            },
        );

        if is_new_or_update {
            self.total_count += 1;
        }

        // Return true for both new messages and updates
        true
    }

    /// Get a message by service, sender, and ID.
    pub fn get(
        &self,
        service_id: u32,
        sender_address: &[u8; 16],
        message_id: &[u8; 16],
    ) -> Option<&StoredMessage> {
        self.messages
            .get(&service_id)?
            .get(sender_address)?
            .get(message_id)
    }

    /// Get all messages from a sender in a service.
    pub fn by_sender(&self, service_id: u32, sender_address: &[u8; 16]) -> Vec<&StoredMessage> {
        self.messages
            .get(&service_id)
            .and_then(|s| s.get(sender_address))
            .map(|m| m.values().collect())
            .unwrap_or_default()
    }

    /// Get all messages in a service.
    pub fn by_service(&self, service_id: u32) -> Vec<&StoredMessage> {
        self.messages
            .get(&service_id)
            .map(|s| s.values().flat_map(|m| m.values()).collect())
            .unwrap_or_default()
    }

    /// Query messages with a predicate.
    pub fn query<F>(&self, service_id: u32, predicate: F) -> Vec<&StoredMessage>
    where
        F: Fn(&StoredMessage) -> bool,
    {
        self.by_service(service_id)
            .into_iter()
            .filter(|m| predicate(m))
            .collect()
    }

    /// Remove a specific message.
    pub fn remove(
        &mut self,
        service_id: u32,
        sender_address: &[u8; 16],
        message_id: &[u8; 16],
    ) -> Option<StoredMessage> {
        let result = self
            .messages
            .get_mut(&service_id)?
            .get_mut(sender_address)?
            .remove(message_id);

        if result.is_some() {
            self.total_count = self.total_count.saturating_sub(1);
        }

        result
    }

    /// Remove all messages from a sender.
    pub fn remove_sender(&mut self, service_id: u32, sender_address: &[u8; 16]) -> usize {
        let count = self
            .messages
            .get_mut(&service_id)
            .and_then(|s| s.remove(sender_address))
            .map(|m| m.len())
            .unwrap_or(0);

        self.total_count = self.total_count.saturating_sub(count);
        count
    }

    /// Prune expired messages.
    pub fn prune_expired(&mut self) -> usize {
        let now = now();
        let mut removed = 0;

        for service_map in self.messages.values_mut() {
            for sender_map in service_map.values_mut() {
                let expired: Vec<[u8; 16]> = sender_map
                    .iter()
                    .filter(|(_, m)| m.message.is_expired())
                    .map(|(id, _)| *id)
                    .collect();

                for id in expired {
                    sender_map.remove(&id);
                    removed += 1;
                }
            }
        }

        self.total_count = self.total_count.saturating_sub(removed);
        self.last_prune = now;
        removed
    }

    /// Get total message count.
    pub fn len(&self) -> usize {
        self.total_count
    }

    /// Check if empty.
    pub fn is_empty(&self) -> bool {
        self.total_count == 0
    }

    /// Get count by service.
    pub fn service_count(&self, service_id: u32) -> usize {
        self.messages
            .get(&service_id)
            .map(|s| s.values().map(|m| m.len()).sum())
            .unwrap_or(0)
    }

    /// Run prune if interval passed.
    fn maybe_prune(&mut self) {
        let now = now();
        if now.saturating_sub(self.last_prune) >= self.config.prune_interval_secs {
            self.prune_expired();
        }
    }

    /// Evict oldest message in a service.
    fn evict_oldest_in_service(&mut self, service_id: u32) {
        let Some(service_map) = self.messages.get_mut(&service_id) else {
            return;
        };

        let mut oldest: Option<([u8; 16], [u8; 16], u64)> = None;

        for (sender, msgs) in service_map.iter() {
            for (id, stored) in msgs.iter() {
                match oldest {
                    Some((_, _, ts)) if stored.message.timestamp < ts => {
                        oldest = Some((*sender, *id, stored.message.timestamp));
                    }
                    None => {
                        oldest = Some((*sender, *id, stored.message.timestamp));
                    }
                    _ => {}
                }
            }
        }

        if let Some((sender, id, _)) = oldest {
            if let Some(sender_map) = service_map.get_mut(&sender) {
                sender_map.remove(&id);
                self.total_count = self.total_count.saturating_sub(1);
            }
        }
    }

    /// Evict oldest message from a sender.
    fn evict_oldest_from_sender(&mut self, service_id: u32, sender_address: [u8; 16]) {
        let Some(sender_map) = self
            .messages
            .get_mut(&service_id)
            .and_then(|s| s.get_mut(&sender_address))
        else {
            return;
        };

        let oldest = sender_map
            .iter()
            .min_by_key(|(_, m)| m.message.timestamp)
            .map(|(id, _)| *id);

        if let Some(id) = oldest {
            sender_map.remove(&id);
            self.total_count = self.total_count.saturating_sub(1);
        }
    }
}

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

fn now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{identity::ServiceIdentity, message::ServiceMessage, service_ids::FAPP};

    fn make_message(id: &ServiceIdentity, seq: u64) -> ServiceMessage {
        ServiceMessage::announce(id, FAPP, b"test".to_vec(), seq)
    }

    #[test]
    fn store_and_retrieve() {
        let mut store = ServiceStore::new();
        let id = ServiceIdentity::generate();
        let msg = make_message(&id, 1);

        assert!(store.store(msg.clone(), id.public_key()));
        assert_eq!(store.len(), 1);

        let retrieved = store.get(FAPP, &id.address(), &msg.id);
        assert!(retrieved.is_some());
    }

    #[test]
    fn duplicate_rejected() {
        let mut store = ServiceStore::new();
        let id = ServiceIdentity::generate();
        let msg = make_message(&id, 1);

        assert!(store.store(msg.clone(), id.public_key()));
        assert!(!store.store(msg.clone(), id.public_key())); // Duplicate
        assert_eq!(store.len(), 1);
    }

    #[test]
    fn higher_sequence_updates() {
        let mut store = ServiceStore::new();
        let id = ServiceIdentity::generate();
        let msg1 = make_message(&id, 1);
        let mut msg2 = make_message(&id, 2);
        msg2.id = msg1.id; // Same ID

        store.store(msg1.clone(), id.public_key());
        assert!(store.store(msg2.clone(), id.public_key())); // Updates

        let retrieved = store.get(FAPP, &id.address(), &msg1.id).unwrap();
        assert_eq!(retrieved.message.sequence, 2);
    }

    #[test]
    fn query_by_sender() {
        let mut store = ServiceStore::new();
        let id1 = ServiceIdentity::generate();
        let id2 = ServiceIdentity::generate();

        store.store(make_message(&id1, 1), id1.public_key());
        store.store(make_message(&id1, 2), id1.public_key());
        store.store(make_message(&id2, 1), id2.public_key());

        let sender1_msgs = store.by_sender(FAPP, &id1.address());
        assert_eq!(sender1_msgs.len(), 2);

        let sender2_msgs = store.by_sender(FAPP, &id2.address());
        assert_eq!(sender2_msgs.len(), 1);
    }

    #[test]
    fn remove_sender() {
        let mut store = ServiceStore::new();
        let id = ServiceIdentity::generate();

        store.store(make_message(&id, 1), id.public_key());
        store.store(make_message(&id, 2), id.public_key());
        assert_eq!(store.len(), 2);

        let removed = store.remove_sender(FAPP, &id.address());
        assert_eq!(removed, 2);
        assert_eq!(store.len(), 0);
    }
}