use async_wakers::{Registration, Wakers};

#[tokio::test]
async fn works() {
    use std::future::Future;
    use std::pin::Pin;
    use std::task::{Context, Poll};
    use std::time::Duration;
    use tokio::task;
    use tokio::time;

    struct WorksFut {
        reg: Option<Registration<'static>>,
        wakers: &'static Wakers,
    }

    impl Future for WorksFut {
        type Output = bool;

        fn poll(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> {
            if self.reg.is_none() {
                match self.wakers.register_waker(ctx.waker().clone()) {
                    Some(reg) => {
                        self.reg = Some(reg);
                        return Poll::Pending;
                    }
                    None => return Poll::Ready(false),
                }
            }

            let reg = self.reg.as_mut().unwrap();
            if reg.is_notified() {
                return Poll::Ready(true);
            }

            reg.update_waker(ctx.waker().clone());
            Poll::Pending
        }
    }

    static WAKERS: Wakers = Wakers::new();

    let task = task::spawn(time::timeout(
        Duration::from_secs(2),
        WorksFut {
            reg: None,
            wakers: &WAKERS,
        },
    ));

    while WAKERS.notifiable() == 0 {
        time::delay_for(Duration::from_millis(100)).await;
    }

    WAKERS.notify_one();

    assert!(task.await.unwrap().unwrap());
    assert_eq!(WAKERS.notifiable(), 0);
}

#[tokio::test]
async fn works_multi() {
    use std::cmp;
    use std::future::Future;
    use std::mem;
    use std::pin::Pin;
    use std::task::{Context, Poll};
    use std::time::Duration;
    use tokio::task;
    use tokio::time;

    struct WorksFut {
        reg: Option<Registration<'static>>,
        wakers: &'static Wakers,
    }

    impl Future for WorksFut {
        type Output = bool;

        // FIXME: contention; allow to lock and do all operation
        fn poll(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> {
            if self.reg.is_none() {
                match self.wakers.register_waker(ctx.waker().clone()) {
                    Some(reg) => {
                        self.reg = Some(reg);
                        return Poll::Pending;
                    }
                    None => return Poll::Ready(false),
                }
            }

            let reg = self.reg.as_mut().unwrap();
            if reg.is_notified() {
                self.reg.take();
                self.wakers.notify_any();
                self.reg = self.wakers.register_waker(ctx.waker().clone());
                //return Poll::Pending;
                return Poll::Ready(true);
            }

            reg.update_waker(ctx.waker().clone());
            Poll::Pending
        }
    }

    static WAKERS: Wakers = Wakers::new();

    let mut tasks = Vec::with_capacity(128);
    for _ in 0..tasks.capacity() {
        tasks.push(task::spawn(time::timeout(
            Duration::from_secs(20),
            WorksFut {
                reg: None,
                wakers: &WAKERS,
            },
        )));
    }

    while WAKERS.notifiable() == 0 {
        time::delay_for(Duration::from_millis(100)).await;
    }

    WAKERS.notify_one();

    let n_tasks = tasks.len();
    let mut succeeded = 0;
    for task in tasks {
        if task.await.unwrap().unwrap() {
            succeeded += 1;
        }
    }

    assert_eq!(cmp::min(8 * mem::size_of::<usize>(), n_tasks), succeeded);
    assert_eq!(WAKERS.notifiable(), 0);
}