linch 0.6.3

use std::{
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::Duration,
};

use crate::error::{RecvError, RecvTimeoutError, TryRecvError};
use futures::{task::AtomicWaker, Stream};
use std::future::Future;

use crate::{inner::Inner, util::async_recv};

// static RECV_ID: AtomicUsize = AtomicUsize::new(1);

/// The receiving half of a channel.
///
/// This struct allows receiving values both synchronously and asynchronously.
/// It can be cloned to create multiple receivers for the same channel.
///
/// # Examples
///
/// ```rust
/// use linch::bounded;
///
/// let (sender, receiver) = bounded(5);
/// sender.send(42).unwrap();
///
/// // Receive synchronously
/// let value = receiver.recv().unwrap();
/// assert_eq!(value, 42);
///
/// // Clone the receiver
/// let receiver2 = receiver.clone();
/// ```
pub struct Receiver<T> {
    pub(crate) rx: crossbeam_channel::Receiver<T>,
    inner: Arc<Inner>,
}

impl<T> Clone for Receiver<T> {
    fn clone(&self) -> Self {
        self.inner.inc_rx();

        Self {
            rx: self.rx.clone(),
            inner: self.inner.clone(),
        }
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        if self.inner.dec_rx() == 1 {
            let mut signal_queues = self.inner.signal_queues();
            while let Some(waker) = signal_queues.pop_recv() {
                waker.as_ref().wake();
            }

            while let Some(waker) = signal_queues.pop_send() {
                waker.as_ref().wake();
            }
        }
    }
}

impl<T> Receiver<T> {
    /// Creates a new receiver from the underlying crossbeam channel and inner state.
    ///
    /// This is typically not called directly by users, but rather through [`channel`](crate::channel).
    pub(crate) fn new(rx: crossbeam_channel::Receiver<T>, inner: Arc<Inner>) -> Self {
        Self { rx, inner }
    }

    /// Receives a value synchronously.
    ///
    /// This method blocks until a value is available in the channel or until
    /// all senders have been dropped.
    ///
    /// # Returns
    ///
    /// * `Ok(value)` if a value was received successfully
    /// * `Err(RecvError)` if all senders have been dropped and the channel is empty
    ///
    /// # Examples
    ///
    /// ```rust
    /// use linch::bounded;
    ///
    /// let (sender, receiver) = bounded(1);
    /// sender.send(42).unwrap();
    /// assert_eq!(receiver.recv().unwrap(), 42);
    /// ```
    pub fn recv(&self) -> Result<T, RecvError> {
        let res = self.rx.recv();
        if res.is_ok() {
            self.signal_send();
        }

        Ok(res?)
    }

    /// Receives a value synchronously with a timeout.
    ///
    /// This method blocks until a value is available in the channel, the timeout
    /// expires, or all senders have been dropped.
    ///
    /// # Arguments
    ///
    /// * `timeout` - The maximum duration to wait
    ///
    /// # Returns
    ///
    /// * `Ok(value)` if a value was received successfully
    /// * `Err(RecvTimeoutError::Timeout)` if the timeout expired
    /// * `Err(RecvTimeoutError::Disconnected)` if all senders have been dropped
    ///
    /// # Examples
    ///
    /// ```rust
    /// use linch::bounded;
    /// use std::time::Duration;
    ///
    /// let (_sender, receiver) = bounded::<i32>(1);
    ///
    /// // This will timeout since no values are sent
    /// let result = receiver.recv_timeout(Duration::from_millis(10));
    /// assert!(result.is_err());
    /// ```
    pub fn recv_timeout(&self, timeout: Duration) -> Result<T, RecvTimeoutError> {
        let res = self.rx.recv_timeout(timeout);
        if res.is_ok() {
            self.signal_send();
        }

        Ok(res?)
    }

    /// Attempts to receive a value without blocking.
    ///
    /// This method will either receive a value immediately if one is available
    /// in the channel, or return an error without blocking.
    ///
    /// # Returns
    ///
    /// * `Ok(value)` if a value was received successfully
    /// * `Err(TryRecvError::Empty)` if the channel is empty but not closed
    /// * `Err(TryRecvError::Disconnected)` if all senders have been dropped and the channel is empty
    ///
    /// # Examples
    ///
    /// ```rust
    /// use linch::bounded;
    /// use linch::TryRecvError;
    ///
    /// let (sender, receiver) = bounded(1);
    ///
    /// // Channel is empty
    /// match receiver.try_recv() {
    ///     Err(TryRecvError::Empty) => {}
    ///     _ => panic!("Expected Empty error"),
    /// }
    ///
    /// // After sending, we can receive
    /// sender.send(42).unwrap();
    /// assert_eq!(receiver.try_recv().unwrap(), 42);
    ///
    /// // Channel is empty again
    /// match receiver.try_recv() {
    ///     Err(TryRecvError::Empty) => {}
    ///     _ => panic!("Expected Empty error"),
    /// }
    /// ```
    pub fn try_recv(&self) -> Result<T, TryRecvError> {
        let res = self.rx.try_recv();
        if res.is_ok() {
            self.signal_send();
        }

        Ok(res?)
    }

    #[inline(always)]
    pub(crate) fn signal_send(&self) {
        if let Some(waker) = { self.inner.signal_queues().pop_send() } {
            waker.as_ref().wake();
        }
    }

    /// Receives a value asynchronously.
    ///
    /// This method returns a future that will complete when a value is available
    /// in the channel or when all senders have been dropped.
    ///
    /// # Returns
    ///
    /// A [`RecvFut`] future that resolves to:
    /// * `Ok(value)` if a value was received successfully
    /// * `Err(RecvError)` if all senders have been dropped and the channel is empty
    ///
    /// # Examples
    ///
    /// ```rust
    /// use linch::bounded;
    ///
    /// # tokio_test::block_on(async {
    /// let (sender, receiver) = bounded(1);
    /// sender.send(42).unwrap();
    /// assert_eq!(receiver.recv_async().await.unwrap(), 42);
    /// # });
    /// ```
    pub fn recv_async(&self) -> RecvFut<'_, T> {
        RecvFut {
            rx: &self.rx,
            inner: &self.inner,
            poll_cnt: 0,
            // recv_id: RECV_ID.fetch_add(2, Ordering::Release),
            waker: AtomicWaker::new(),
        }
    }

    /// Converts the receiver into a stream.
    ///
    /// This allows using the receiver with async stream utilities and for iteration.
    ///
    /// # Returns
    ///
    /// A [`RecvStream`] that yields values from the channel
    ///
    /// # Examples
    ///
    /// ```rust
    /// use linch::bounded;
    /// use futures::StreamExt;
    ///
    /// # tokio_test::block_on(async {
    /// let (sender, receiver) = bounded(3);
    /// sender.send(42).unwrap();
    /// sender.send(43).unwrap();
    /// drop(sender); // Close the channel
    ///
    /// let mut stream = receiver.into_stream();
    /// let values: Vec<_> = stream.collect().await;
    /// assert_eq!(values, vec![42, 43]);
    /// # });
    /// ```
    pub fn into_stream(self) -> RecvStream<T> {
        RecvStream {
            inner: Box::new(RecvStreamInner {
                rx: self,
                poll_cnt: 0,
                waker: AtomicWaker::new(),
            }),
        }
    }
}

/// A future representing an asynchronous receive operation.
///
/// This future is created by the [`recv_async`](Receiver::recv_async) method and will
/// complete when a value is available or when all senders have been dropped.
///
/// # Examples
///
/// ```rust
/// use linch::bounded;
///
/// # tokio_test::block_on(async {
/// let (sender, receiver) = bounded(1);
/// sender.send(42).unwrap();
/// let recv_fut = receiver.recv_async();
/// assert_eq!(recv_fut.await.unwrap(), 42);
/// # });
/// ```
pub struct RecvFut<'a, T> {
    rx: &'a crossbeam_channel::Receiver<T>,
    inner: &'a Arc<Inner>,
    poll_cnt: u32,
    waker: AtomicWaker,
}

impl<'a, T> Unpin for RecvFut<'a, T> {}

impl<'a, T> RecvFut<'a, T> {
    fn poll(
        rx: &crossbeam_channel::Receiver<T>,
        inner: &Arc<Inner>,
        poll_cnt: &mut u32,
        waker: &AtomicWaker,
        cx: &mut Context<'_>,
    ) -> Poll<Result<T, RecvError>> {
        for _ in 0..1 {
            match async_recv(rx) {
                Ok(value) => {
                    let mut signal_queues = inner.signal_queues();

                    if *poll_cnt > 0 {
                        signal_queues.remove_recv(waker as *const AtomicWaker as usize);
                    }

                    if let Some(waker_ptr) = signal_queues.pop_send() {
                        drop(signal_queues);
                        waker_ptr.as_ref().wake();
                    } else {
                        drop(signal_queues);
                    }

                    return Poll::Ready(Ok(value));
                }
                Err(TryRecvError::Empty) => {}
                Err(TryRecvError::Disconnected) => {
                    let mut signal_queues = inner.signal_queues();

                    while let Some(waker_ptr) = signal_queues.pop_send() {
                        waker_ptr.as_ref().wake();
                    }

                    while let Some(waker_ptr) = signal_queues.pop_recv() {
                        waker_ptr.as_ref().wake();
                    }

                    drop(signal_queues);
                    return Poll::Ready(Err(RecvError));
                }
            }
        }

        // if this.poll_cnt < 2 {
        //     cx.waker().wake_by_ref();
        //     return Poll::Pending;
        // }

        waker.register(cx.waker());

        let mut signal_queues = inner.signal_queues();

        match rx.try_recv() {
            Ok(value) => {
                if *poll_cnt > 0 {
                    signal_queues.remove_recv(waker as *const AtomicWaker as usize);
                }

                if let Some(waker_ptr) = signal_queues.pop_send() {
                    drop(signal_queues);
                    waker_ptr.as_ref().wake();
                }

                return Poll::Ready(Ok(value));
            }
            Err(crossbeam_channel::TryRecvError::Empty) => {}
            Err(crossbeam_channel::TryRecvError::Disconnected) => {
                while let Some(waker_ptr) = signal_queues.pop_send() {
                    waker_ptr.as_ref().wake();
                }

                while let Some(waker_ptr) = signal_queues.pop_recv() {
                    waker_ptr.as_ref().wake();
                }

                drop(signal_queues);
                return Poll::Ready(Err(RecvError));
            }
        }

        let waker_ptr = waker as *const AtomicWaker as usize;
        if *poll_cnt > 0 {
            signal_queues.remove_recv(waker_ptr);
        }
        *poll_cnt += 1;
        signal_queues.add_recv(waker_ptr);

        if let Some(waker_ptr) = signal_queues.pop_send() {
            drop(signal_queues);
            waker_ptr.as_ref().wake();
        } else {
            drop(signal_queues);
        }

        Poll::Pending
    }
}

impl<'a, T> Future for RecvFut<'a, T> {
    type Output = Result<T, RecvError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();

        RecvFut::poll(this.rx, this.inner, &mut this.poll_cnt, &mut this.waker, cx)
    }
}

impl<'a, T> Drop for RecvFut<'a, T> {
    fn drop(&mut self) {
        if self.poll_cnt >= 1 {
            let mut signal_queues = self.inner.signal_queues();
            signal_queues.remove_recv(&self.waker as *const AtomicWaker as usize);
        }
    }
}

pub struct RecvStreamInner<T> {
    rx: Receiver<T>,
    poll_cnt: u32,
    waker: AtomicWaker,
}

pub struct RecvStream<T> {
    inner: Box<RecvStreamInner<T>>,
}

impl<T> Unpin for RecvStream<T> {}

impl<T> Stream for RecvStream<T> {
    type Item = T;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();

        match RecvFut::poll(
            &this.inner.rx.rx,
            &this.inner.rx.inner,
            &mut this.inner.poll_cnt,
            &this.inner.waker,
            cx,
        ) {
            Poll::Ready(Ok(item)) => {
                this.inner.poll_cnt = 0;
                // this.inner.waker.take();

                Poll::Ready(Some(item))
            }
            Poll::Ready(Err(_e)) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<T> Drop for RecvStream<T> {
    fn drop(&mut self) {
        if self.inner.poll_cnt >= 1 {
            let mut signal_queues = self.inner.rx.inner.signal_queues();
            signal_queues.remove_recv(&self.inner.waker as *const AtomicWaker as usize);
        }
    }
}