use core::{
    future::poll_fn,
    sync::atomic::{AtomicU64, Ordering},
    task::Poll,
};

use futures_util::task::AtomicWaker;

pub struct EventBitmap {
    bitmap: AtomicU64,
    waker: AtomicWaker,
}

struct BitIter {
    mask: u64,
    index: usize,
    limit: usize,
}

impl Iterator for BitIter {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        while self.index < self.limit {
            let bit = self.index;
            self.index += 1;

            if self.mask & (1 << bit) != 0 {
                return Some(bit);
            }
        }
        None
    }
}

impl EventBitmap {
    pub fn new() -> Self {
        Self {
            bitmap: AtomicU64::new(0),
            waker: AtomicWaker::new(),
        }
    }

    pub fn signal(&self, bit: usize) {
        self.bitmap.fetch_or(1 << bit, Ordering::Release);
        self.waker.wake();
    }

    // pub async fn wait_specific(&self, bit: usize) {
    //     poll_fn(|cx| {
    //         let state = self.bitmap.fetch_and(!(1 << bit), Ordering::Acquire);
    //         if state & (1 << bit) != 0 {
    //             Poll::Ready(())
    //         } else {
    //             self.waker.register(cx.waker());
    //             Poll::Pending
    //         }
    //     })
    //     .await
    // }

    pub async fn wait_any(&self, max: usize) -> impl Iterator<Item = usize> {
        poll_fn(|cx| {
            let mask = self.bitmap.swap(0, Ordering::Acquire);
            if mask == 0 {
                self.waker.register(cx.waker());
                return Poll::Pending;
            }

            Poll::Ready(BitIter {
                mask,
                index: 0,
                limit: max,
            })
        })
        .await
    }
}
