path: root/src/message_subscriptions.rs

use std::collections::HashMap;

use jid::BareJID;
use tokio::sync::mpsc::{self, Receiver};
use uuid::Uuid;

use crate::message::{ArcMacawMessage, MacawMessage};

pub struct MessageSubscriptions {
    all: HashMap<Uuid, mpsc::Sender<(BareJID, ArcMacawMessage)>>,
    subset: HashMap<BareJID, HashMap<Uuid, mpsc::Sender<ArcMacawMessage>>>,
}

impl MessageSubscriptions {
    pub fn new() -> Self {
        Self {
            all: HashMap::new(),
            subset: HashMap::new(),
        }
    }

    pub async fn broadcast(&mut self, to: BareJID, message: ArcMacawMessage) {
        // subscriptions to all
        let mut removals = Vec::new();
        for (id, sender) in &self.all {
            match sender.send((to.clone(), message.clone())).await {
                Ok(_) => {}
                Err(_) => {
                    removals.push(*id);
                }
            }
        }
        for removal in removals {
            self.all.remove(&removal);
        }

        // subscriptions to specific chat
        if let Some(subscribers) = self.subset.get_mut(&to) {
            let mut removals = Vec::new();
            for (id, sender) in &*subscribers {
                match sender.send(message.clone()).await {
                    Ok(_) => {}
                    Err(_) => {
                        removals.push(*id);
                    }
                }
            }
            for removal in removals {
                subscribers.remove(&removal);
            }
            if subscribers.is_empty() {
                self.subset.remove(&to);
            }
        }
    }

    pub fn subscribe_all(&mut self) -> (Uuid, Receiver<(BareJID, ArcMacawMessage)>) {
        let (send, recv) = mpsc::channel(10);
        let id = Uuid::new_v4();
        self.all.insert(id, send);
        (id, recv)
    }

    pub fn subscribe_chat(&mut self, chat: BareJID) -> (Uuid, Receiver<ArcMacawMessage>) {
        let (send, recv) = mpsc::channel(10);
        let id = Uuid::new_v4();
        if let Some(chat_subscribers) = self.subset.get_mut(&chat) {
            chat_subscribers.insert(id, send);
        } else {
            let hash_map = HashMap::from([(id, send)]);
            self.subset.insert(chat, hash_map);
        }
        (id, recv)
    }

    pub fn unsubscribe_all(&mut self, sub_id: Uuid) {
        self.all.remove(&sub_id);
    }

    pub fn unsubscribe_chat(&mut self, sub_id: Uuid, chat: BareJID) {
        if let Some(chat_subs) = self.subset.get_mut(&chat) {
            chat_subs.remove(&sub_id);
        }
    }
}