use tokio::sync::broadcast;

pub(crate) struct Shutdown {
    shutdown: bool,
    notify: broadcast::Receiver<()>,
    sender: broadcast::Sender<()>,
}

impl Shutdown {
    pub(crate) fn new() -> Shutdown {
        let (sender, notify) = broadcast::channel(1);
        Shutdown {
            shutdown: false,
            notify,
            sender,
        }
    }

    pub(crate) fn shutdown(&mut self) {
        if self.shutdown {
            return;
        }
        let _ = self.sender.send(());
        self.shutdown = true;
    }

    pub(crate) fn receiver(&self) -> ShutdownReceiver {
        ShutdownReceiver::new(self.notify.resubscribe())
    }
}

#[derive(Debug)]
pub(crate) struct ShutdownReceiver {
    shutdown: bool,
    notify: broadcast::Receiver<()>,
}

impl ShutdownReceiver {
    pub(crate) fn new(notify: broadcast::Receiver<()>) -> ShutdownReceiver {
        ShutdownReceiver {
            shutdown: false,
            notify,
        }
    }

    pub(crate) async fn recv(&mut self) {
        if self.shutdown {
            return;
        }
        let _ = self.notify.recv().await;
        self.shutdown = true;
    }
}

impl Clone for ShutdownReceiver {
    fn clone(&self) -> Self {
        ShutdownReceiver {
            shutdown: self.shutdown,
            notify: self.notify.resubscribe(),
        }
    }
}