use crate::work::WorkHandle;
use eyre::bail;
use std::cell::RefCell;
use tracing::error;
use tracing::info;

#[derive(Debug, Default)]
pub struct WorkTracker {
    // We use interior mutability here to avoid annoyances when passing around context when doing immediate mode UI stuff
    remaining_work: RefCell<Vec<WorkHandle>>,
}

impl WorkTracker {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn track(&self, work: WorkHandle) {
        let mut remaining_work = self.remaining_work.borrow_mut();
        remaining_work.push(work);
    }
    pub fn prune(&self) {
        let mut remaining_work = self.remaining_work.borrow_mut();
        remaining_work.retain(|work| !work.join_handle.is_finished());
    }
    pub async fn finish(self) -> eyre::Result<()> {
        let mut remaining_work = self.remaining_work.take();
        if remaining_work.is_empty() {
            return Ok(());
        }
        let mut pass = true;
        while let Some(work) = remaining_work.pop() {
            match work.is_err_if_discarded {
                true => {
                    info!(
                        "Waiting for task {:?} to finish, {} remain...",
                        work.description,
                        remaining_work.len()
                    );
                    match work.join_handle.await? {
                        Ok(()) => {}
                        Err(e) => {
                            pass = false;
                            error!(
                                "Error encountered when waiting for background work to finish: {e:?}"
                            );
                        }
                    }
                }
                false => {
                    info!(
                        "Discarding work {:?} because it is not marked as important, {} remain...",
                        work.description,
                        remaining_work.len()
                    );
                    work.join_handle.abort();
                }
            }
        }
        if pass {
            Ok(())
        } else {
            bail!("Failed to complete all work. See above for details.");
        }
    }
}
