use crate::util::BoxFuture;
use std::error;
use tokio_util::sync::CancellationToken;

pub struct TaskStack<E>
where
    E: error::Error + Send + 'static,
{
    wait_group: Vec<BoxFuture<'static, Result<(), E>>>,
}

impl<E> TaskStack<E>
where
    E: error::Error + Send + 'static,
{
    pub fn new() -> TaskStack<E> {
        TaskStack {
            wait_group: Vec::new(),
        }
    }

    /// push adds the given Future to the stack, to be executed once stop is called.
    pub fn push<Fut>(&mut self, f: Fut)
    where
        Fut: futures::Future<Output = Result<(), E>> + Send + 'static,
    {
        self.wait_group.push(Box::pin(f))
    }

    /// push_spawn will spawn the given closure in a tokio task. Once the CancellationToken is
    /// cancelled the closure is expected to return.
    pub fn push_spawn<F, Fut>(&mut self, mut f: F)
    where
        Fut: futures::Future<Output = Result<(), E>> + Send + 'static,
        F: FnMut(CancellationToken) -> Fut,
    {
        let canceller = CancellationToken::new();
        let handle = tokio::spawn(f(canceller.clone()));
        self.push(async move {
            canceller.cancel();
            handle.await.expect("failed to join task")
        });
    }

    /// stop will process all operations which have been pushed onto the stack in the reverse order
    /// they were pushed.
    pub async fn stop(mut self) -> Result<(), E> {
        // reverse wait_group in place, so we stop the most recently added first. Since this method
        // consumes self this is fine.
        self.wait_group.reverse();

        for fut in self.wait_group {
            fut.await?;
        }
        Ok(())
    }
}

impl<E> Default for TaskStack<E>
where
    E: error::Error + Send + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}