postvan 0.0.2

use std::{
    collections::HashSet,
    fmt::Debug,
    mem::{Discriminant, discriminant},
    sync::Arc,
};

use anyhow::Result;
use dashmap::DashMap;
use tokio::sync::{
    broadcast::{self, Receiver, Sender},
    mpsc::{self, UnboundedReceiver, UnboundedSender},
};
use uuid::Uuid;

pub enum CommandMessage<M> {
    Route(Arc<M>),
    Subscribe(Address, Discriminant<M>),
    Unregister(Address),
}

pub trait Message: Clone + Debug + Send + Sync + 'static {}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Address(pub Uuid);

#[derive(Debug)]
pub struct Letterbox<M: Clone> {
    sender: Sender<Arc<M>>,
    receiver: Receiver<Arc<M>>,
    address: Option<Address>,
    post_tx: Option<UnboundedSender<CommandMessage<M>>>,
}
impl<M: Message> Letterbox<M> {
    pub fn new(capacity: usize) -> Self {
        let (sender, receiver) = broadcast::channel(capacity);
        Self {
            sender,
            receiver,
            address: None,
            post_tx: None,
        }
    }
    pub fn post(&mut self, message: M) -> Result<()> {
        let arc_msg = Arc::new(message);
        self.post_tx
            .as_ref()
            .expect("Cannot post anything using an unregistered letterbox")
            .send(CommandMessage::Route(arc_msg))?;
        Ok(())
    }
    pub fn check(&mut self) -> Vec<Arc<M>> {
        let mut messages = Vec::new();
        loop {
            match self.receiver.try_recv() {
                Ok(message) => messages.push(message),
                Err(broadcast::error::TryRecvError::Empty) => break,
                Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
                    // You missed `skipped` number of messages.
                    // Log this or handle it depending on your requirements,
                    // but keep looping to get the fresh data.
                }
                Err(broadcast::error::TryRecvError::Closed) => break,
            }
        }
        messages
    }
    pub fn subscribe(&mut self, message: M) -> Result<()> {
        let (post_tx, address) = self
            .post_tx
            .as_ref()
            .zip(self.address.clone())
            .expect("Cannot subscribe using an unregistered letterbox");
        post_tx.send(CommandMessage::Subscribe(address, discriminant(&message)))?;
        Ok(())
    }
}

#[derive(Debug)]
pub struct Postoffice<M: Message> {
    subscriptions: DashMap<Discriminant<M>, HashSet<Address>>,
    registry: DashMap<Address, Sender<Arc<M>>>,
    post_tx: UnboundedSender<CommandMessage<M>>,
    post_rx: UnboundedReceiver<CommandMessage<M>>,
}
impl<M: Message> Postoffice<M> {
    pub fn new() -> Self {
        let (post_tx, post_rx) = mpsc::unbounded_channel();
        Self {
            subscriptions: DashMap::new(),
            registry: DashMap::new(),
            post_tx,
            post_rx,
        }
    }
    pub fn register(&self, mailbox: &mut Letterbox<M>) {
        let address = Address(Uuid::new_v4());
        self.registry
            .insert(address.clone(), mailbox.sender.clone());
        mailbox.address = Some(address);
        mailbox.post_tx = Some(self.post_tx.clone());
    }
    pub fn tick(&mut self) -> Option<()> {
        if let Ok(cmd_message) = self.post_rx.try_recv() {
            self.handle_command(cmd_message);
            Some(())
        } else {
            None
        }
    }
    pub async fn tick_async(&mut self) -> Option<()> {
        if let Some(cmd_message) = self.post_rx.recv().await {
            self.handle_command(cmd_message);
            Some(())
        } else {
            None
        }
    }
    fn handle_command(&mut self, cmd_message: CommandMessage<M>) {
        match cmd_message {
            CommandMessage::Route(message) => {
                let target_discriminant = discriminant(&*message);
                if let Some(addresses) = self.subscriptions.get(&target_discriminant) {
                    for address in addresses.iter() {
                        if let Some(sender) = self.registry.get(address) {
                            let _ = sender.send(message.clone());
                        }
                    }
                }
            }
            CommandMessage::Subscribe(address, d) => {
                self.subscriptions
                    .entry(d)
                    .or_default()
                    .value_mut()
                    .insert(address);
            }
            CommandMessage::Unregister(address) => {
                self.registry.remove(&address);
                for mut entry in self.subscriptions.iter_mut() {
                    entry.value_mut().remove(&address);
                }
            }
        }
    }
}
impl<M: Message> Default for Postoffice<M> {
    fn default() -> Self {
        Self::new()
    }
}