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

use futures_util::task::AtomicWaker;

use crate::{sync::spin_rwlock::IrqSafeRwLock, waker::QueueWaker};

pub trait EventNotify: Sync {
    fn subscribe(&self, waker: &Waker);
    fn unsubscribe(&self, waker: &Waker);
    fn notify_one(&self);
    fn notify_all(&self);
}

impl EventNotify for AtomicWaker {
    fn subscribe(&self, waker: &Waker) {
        self.register(waker);
    }
    fn unsubscribe(&self, _waker: &Waker) {}

    fn notify_one(&self) {
        self.wake();
    }
    fn notify_all(&self) {
        self.wake();
    }
}

pub struct OneTimeEvent<T> {
    // TODO lockless like OneTimeInit?
    value: IrqSafeRwLock<Option<T>>,
    notify: QueueWaker,
}

pub struct BitmapEvent<N: EventNotify> {
    value: AtomicU64,
    notify: N,
}

pub struct BoolEvent {
    state: AtomicBool,
    notify: QueueWaker,
}

impl<T> OneTimeEvent<T> {
    /// Constructs an [OneTimeEvent]
    pub const fn new() -> Self {
        Self {
            value: IrqSafeRwLock::new(None),
            notify: QueueWaker::new(),
        }
    }

    /// Signals a value into the event. Returns an [Err] with the value if the event
    /// has already been signalled before.
    pub fn try_signal(&self, value: T) -> Result<(), T> {
        let mut lock = self.value.write();
        if lock.is_some() {
            return Err(value);
        }
        lock.replace(value);
        self.notify.wake_all();
        Ok(())
    }

    /// Signals a value into the event.
    ///
    /// # Panics
    ///
    /// Will panic if the event has already been signal. Use this function only when
    /// the event is known to only be signalled once.
    pub fn signal(&self, value: T) {
        if self.try_signal(value).is_err() {
            panic!("OneTimeEvent<T> signalled more than once");
        }
    }

    /// Returns `true` if a value has been signalled into the event
    pub fn is_signalled(&self) -> bool {
        self.value.read().is_some()
    }

    /// Returns [Some] with the copy of the value if the event has been signalled,
    /// [None] otherwise
    pub fn try_read_copy(&self) -> Option<T>
    where
        T: Copy,
    {
        *self.value.read()
    }

    pub fn poll(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.notify.register(cx.waker());
        if self.is_signalled() {
            self.notify.remove(cx.waker());
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    /// Waits for event to happen without returning anything
    pub async fn wait(&self) {
        poll_fn(|cx| self.poll(cx)).await
    }

    /// Waits for event to happen and copies its data as a return value
    pub async fn wait_copy(&self) -> T
    where
        T: Copy,
    {
        poll_fn(|cx| {
            self.notify.register(cx.waker());
            if let Some(value) = self.try_read_copy() {
                self.notify.remove(cx.waker());
                Poll::Ready(value)
            } else {
                Poll::Pending
            }
        })
        .await
    }
}

impl BoolEvent {
    pub const fn new() -> Self {
        Self {
            state: AtomicBool::new(false),
            notify: QueueWaker::new(),
        }
    }

    pub fn is_signalled(&self) -> bool {
        self.state.load(Ordering::Acquire)
    }

    pub fn signal(&self) {
        if self.try_signal().is_err() {
            panic!("OneTimeBoolEvent signalled more than once");
        }
    }

    pub fn signal_saturating(&self) {
        if !self.state.swap(true, Ordering::Release) {
            self.notify.wake_all();
        }
        // self.try_signal().ok();
    }

    pub fn try_signal(&self) -> Result<(), bool> {
        self.state
            .compare_exchange(false, true, Ordering::Release, Ordering::Relaxed)?;
        self.notify.wake_all();
        Ok(())
    }

    pub fn poll_reset(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.notify.register(cx.waker());
        if self
            .state
            .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            // Event set and immediately reset
            self.notify.remove(cx.waker());
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    pub fn poll(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.notify.register(cx.waker());
        if self.is_signalled() {
            self.notify.remove(cx.waker());
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    pub async fn wait_reset(&self) {
        poll_fn(|cx| self.poll_reset(cx)).await
    }

    pub async fn wait(&self) {
        poll_fn(|cx| self.poll(cx)).await
    }
}

impl<N: EventNotify> BitmapEvent<N> {
    pub const fn new(notify: N) -> Self {
        Self {
            value: AtomicU64::new(0),
            notify,
        }
    }

    pub fn signal(&self, mask: u64) {
        self.value.fetch_or(mask, Ordering::Release);
        self.notify.notify_all();
    }

    pub async fn wait(&self) -> u64 {
        poll_fn(|cx| {
            let state = self.value.swap(0, Ordering::AcqRel);
            if state != 0 {
                self.notify.unsubscribe(cx.waker());
                Poll::Ready(state)
            } else {
                self.notify.subscribe(cx.waker());
                Poll::Pending
            }
        })
        .await
    }
}
