//! Background tasks handling. You must first await a clone of `Tasks` (i.e., spawn it). When
//! you then await `Tasks.finish()`, it will await all blocking tasks. If they push more
//! blocking tasks meanwhile, they will be awaited, too. After all blocking tasks are
//! finished, non-blocking tasks will be awaited for some additional time. If they're not
//! finished in that time, they will be dropped. A non-blocking task cannot push a blocking
//! task (it will panic).

use std::{future::Future, pin::Pin, rc::Rc, time::Duration};
use std::cell::{Cell, OnceCell, RefCell};
use std::task::{Context, Poll, Waker};

#[derive(Clone, derive_more::Deref)]
pub(crate) struct Tasks(Rc<TasksInner>);

pub(crate) struct TasksInner {
    blocking: RefCell<Vec<Pin<Box<dyn Future<Output = ()>>>>>,
    non_blocking: RefCell<Vec<Pin<Box<dyn Future<Output = ()>>>>>,
    timer: RefCell<Option<Pin<Box<futures_timer::Delay>>>>,
    run_waker: Cell<Option<Waker>>,
    finish_waker: Cell<Option<Waker>>,
    finishing: OnceCell<()>,
    finished: Cell<bool>,
}

pub(crate) struct TasksFinish<'a> {
    tasks: &'a Tasks,
}

impl Tasks {
    pub(crate) fn new() -> Self {
        // Cannot derive Default for Tasks, see e.g.
        // https://users.rust-lang.org/t/why-cant-vec-t-implement-default-when-t-isnt-default/45868
        Self(Rc::new(TasksInner {
            blocking: RefCell::new(Vec::new()),
            non_blocking: RefCell::new(Vec::new()),
            timer: RefCell::new(None),
            run_waker: Cell::new(None),
            finish_waker: Cell::new(None),
            finishing: OnceCell::new(),
            finished: Cell::new(false),
        }))
    }

    /// Add the task to the list of blocking tasks.
    pub(crate) fn push_blocking(&self, task: impl Future<Output = ()> + 'static) {
        assert!(self.timer.borrow().is_none());
        self.blocking.borrow_mut().push(Box::pin(task));
    }

    /// Add the task to the list of non-blocking tasks.
    pub(crate) fn push_non_blocking(&self, task: impl Future<Output = ()> + 'static) {
        self.non_blocking.borrow_mut().push(Box::pin(task));
    }

    /// Return a Future that waits for the tasks. Can be called only once.
    pub(crate) fn finish(&self) -> TasksFinish<'_> {
        self.finishing.set(()).unwrap();
        if let Some(run_waker) = self.run_waker.take() {
            // wake the task running loop; since we set self.finishing, this will set up the
            // timeout inside the loop
            run_waker.wake();
        }
        TasksFinish {
            tasks: self,
        }
    }
}

impl Future for Tasks {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // First, poll the blocking tasks.
        let mut remaining = Vec::new();
        loop {
            let tasks = self.blocking.take();
            if tasks.is_empty() {
                break;
            }
            for mut task in tasks {
                match task.as_mut().poll(cx) {
                    Poll::Ready(_) => {}   // fnished, don't add it back
                    Poll::Pending => remaining.push(task),
                }
            }
            // We need to also check any tasks that might have been pushed meanwhile.
        }
        let has_blocking = !remaining.is_empty();
        self.blocking.replace(remaining);

        // Start the timeout if there are no more blocking tasks and we're finishing.
        if !has_blocking && self.finishing.get().is_some() && self.timer.borrow().is_none() {
            self.timer.replace(Some(
                Box::pin(futures_timer::Delay::new(Duration::from_secs(1)))
            ));
        }

        // Now, poll the non-blocking tasks.
        let mut remaining = Vec::new();
        loop {
            let tasks = self.non_blocking.take();
            if tasks.is_empty() {
                break;
            }
            for mut task in tasks {
                match task.as_mut().poll(cx) {
                    Poll::Ready(_) => {}   // fnished, don't add it back
                    Poll::Pending => remaining.push(task),
                }
            }
            // We need to also check any tasks that might have been pushed meanwhile.
        }
        let has_non_blocking = !remaining.is_empty();
        self.non_blocking.replace(remaining);

        // Poll the timer.
        let mut pending = true;
        if let Some(timer) = &mut *self.timer.borrow_mut() {
            // Note the timer is set (and thus we can get here) only if we're finishing and
            // don't have any blocking tasks. We can finish either when there are no more
            // tasks (it's enough to check non-blocking tasks, since there are no blocking
            // tasks at this point), or when the timeout elapses.
            pending = has_non_blocking;
            if timer.as_mut().poll(cx) == Poll::Ready(()) {
                // Timeout; we're done. (Note the timer is not set if we still have blocking
                // tasks.)
                pending = false;
            }
        }

        if pending {
            self.run_waker.set(Some(cx.waker().clone()));
            Poll::Pending
        } else {
            self.run_waker.set(None);
            self.finished.set(true);
            if let Some(finish_waker) = self.finish_waker.take() {
                finish_waker.wake();
            }
            Poll::Ready(())
        }
    }
}

impl<'a> Future for TasksFinish<'a> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.tasks.finished.get() {
            self.tasks.finish_waker.set(None);
            Poll::Ready(())
        } else {
            self.tasks.finish_waker.set(Some(cx.waker().clone()));
            Poll::Pending
        }
    }
}
