use core::{
    future::poll_fn,
    task::{Context, Poll},
};

use crossbeam_queue::{ArrayQueue, SegQueue};

use crate::waker::QueueWaker;

pub type BoundedQueue<T> = ArrayQueue<T>;
pub type UnboundedQueue<T> = SegQueue<T>;

pub struct BoundedMpmcQueue<T> {
    queue: BoundedQueue<T>,
    read_notify: QueueWaker,
    write_notify: QueueWaker,
}

pub struct UnboundedMpmcQueue<T> {
    queue: UnboundedQueue<T>,
    read_notify: QueueWaker,
}

impl<T> BoundedMpmcQueue<T> {
    pub fn new(capacity: usize) -> Self {
        Self {
            queue: BoundedQueue::new(capacity),
            read_notify: QueueWaker::new(),
            write_notify: QueueWaker::new(),
        }
    }

    pub fn try_push_back(&self, value: T) -> Result<(), T> {
        self.queue.push(value)?;
        self.read_notify.wake_one();
        Ok(())
    }

    pub fn try_pop_front(&self) -> Option<T> {
        let value = self.queue.pop()?;
        self.write_notify.wake_one();
        Some(value)
    }

    pub async fn pop_front(&self) -> T {
        poll_fn(|cx| {
            self.read_notify.register(cx.waker());
            match self.try_pop_front() {
                Some(value) => {
                    self.read_notify.remove(cx.waker());
                    Poll::Ready(value)
                }
                None => Poll::Pending,
            }
        })
        .await
    }

    pub async fn push_back(&self, value: T) {
        let mut value = Some(value);
        poll_fn(|cx| {
            self.write_notify.register(cx.waker());
            match self.try_push_back(value.take().unwrap()) {
                Ok(()) => {
                    // Succeeded
                    self.write_notify.remove(cx.waker());
                    Poll::Ready(())
                }
                Err(back) => {
                    // Failed, put the value back
                    value.replace(back);
                    Poll::Pending
                }
            }
        })
        .await
    }

    pub fn poll_not_empty(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.read_notify.register(cx.waker());
        if !self.queue.is_empty() {
            self.read_notify.remove(cx.waker());
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

impl<T> UnboundedMpmcQueue<T> {
    pub const fn new() -> Self {
        Self {
            queue: SegQueue::new(),
            read_notify: QueueWaker::new(),
        }
    }

    pub fn push_back(&self, value: T) {
        self.queue.push(value);
        self.read_notify.wake_one();
    }

    pub fn try_pop_front(&self) -> Option<T> {
        self.queue.pop()
    }

    pub async fn pop_front(&self) -> T {
        poll_fn(|cx| {
            self.read_notify.register(cx.waker());
            if let Some(value) = self.try_pop_front() {
                self.read_notify.remove(cx.waker());
                Poll::Ready(value)
            } else {
                Poll::Pending
            }
        })
        .await
    }
}
