//! Command-wide execution budgets exclude the union of component compilation waits, including child spawns.
use std::{sync::Arc, time::Duration};
use tokio::{sync::watch, time::Instant};

#[derive(Clone, Copy)]
struct State {
    deadline: Instant,
    compiling: usize,
    paused_at: Instant,
}

pub(super) struct Deadline(watch::Sender<State>);

impl Deadline {
    pub fn new(duration: Duration) -> Arc<Self> {
        let now = Instant::now();
        Arc::new(Self(
            watch::channel(State {
                deadline: now + duration,
                compiling: 0,
                paused_at: now,
            })
            .0,
        ))
    }

    pub fn pause(self: &Arc<Self>) -> Compilation {
        self.0.send_modify(|state| {
            if state.compiling == 0 {
                state.paused_at = Instant::now();
            }
            state.compiling += 1;
        });
        Compilation(self.clone())
    }

    pub async fn expired(&self) {
        let mut state = self.0.subscribe();
        loop {
            let current = *state.borrow_and_update();
            if current.compiling == 0 {
                tokio::select! {
                    biased;
                    changed = state.changed() => { changed.unwrap(); }
                    _ = tokio::time::sleep_until(current.deadline) => return,
                }
            } else {
                state.changed().await.unwrap();
            }
        }
    }
}

pub(super) struct Compilation(Arc<Deadline>);

impl Drop for Compilation {
    fn drop(&mut self) {
        self.0.0.send_modify(|state| {
            state.compiling -= 1;
            if state.compiling == 0 {
                state.deadline += state.paused_at.elapsed();
            }
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn overlapping_compilations_share_one_pause_and_drop_resumes() {
        let deadline = Deadline::new(Duration::from_secs(2));
        let first = deadline.pause();
        let paused_at = deadline.0.borrow().paused_at;
        let second = deadline.pause();
        assert_eq!(deadline.0.borrow().compiling, 2);
        assert_eq!(deadline.0.borrow().paused_at, paused_at);
        let original = deadline.0.borrow().deadline;
        drop(first);
        assert_eq!(deadline.0.borrow().compiling, 1);
        assert_eq!(deadline.0.borrow().deadline, original);
        drop(second);
        assert_eq!(deadline.0.borrow().compiling, 0);
        assert!(deadline.0.borrow().deadline >= original);
    }
}