//! A multi-producer singe-consumer channel, every message sent in channel by a sender
//! is received by the reader.
//!
//!
//! The queue of the channel is internally represented as a linked-list, relying on atomic
//! only operations for concurrent lock-free access, and only the reader is responsible is tasked
//! with freeing up used memory. Writes are always lock free, while reads are wait free most of the
//! time and are 1 atomic operation to get message.
//!
//! ```
//! use qanat::mpsc;
//!
//! // Create new channel
//! let (tx, mut rx) = mpsc::channel(10.try_into().unwrap());
//!
//! // send a message
//! tx.try_send(1).unwrap();
//! // Or using async to await if channel is full
//! // tx.send(1).await.unwrap();
//!
//! // Reading from channel
//! assert_eq!(rx.try_recv().unwrap(), 1);
//!
//! tx.try_send(2).unwrap();
//!
//! // Creating a snapshot for channel
//! let snapshot = rx.snapshot();
//!
//! // Dropping channel
//! drop(tx);
//! drop(rx);
//!
//!
//! // Creating a new channel from snapshot of previous one
//! let (tx, mut rx) = mpsc::restore_channel_from_snapshot(snapshot)
//!     .unwrap();
//!
//! assert_eq!(rx.try_recv().unwrap(), 2);
//!
//! tx.try_send(3).unwrap();
//!
//! assert_eq!(rx.try_recv().unwrap(), 3);
//! ```

use std::{
    sync::{
        atomic::{Ordering, AtomicU64, AtomicI32, AtomicPtr, AtomicBool},
        Arc
    },
    num::NonZeroUsize,
    ops::Deref
};
use async_event::Event;
use diatomic_waker::{WakeSource, WakeSink};

/// Create a new channel with a given capacity
pub fn channel<T: Send + Sync>(capacity: NonZeroUsize) -> (Sender<T>, Receiver<T>) {
    Inner::create(capacity.get())
}

/// Create a new channel with a given capacity
pub fn unbounded_channel<T: Send + Sync>() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
    InnerUnbounded::create()
}

/// Create a new channel from a snapshot returned by [`Receiver::snapshot`]
/// None is returned if channel capacity is smaller than the number of message items
pub fn restore_channel_from_snapshot<T: Send + Sync + Clone>(snapshot: (Vec<T>, u64)) -> Option<(Sender<T>, Receiver<T>)> {
    // We first create channel with msg count as capacity
    let len   = snapshot.0.len();
    let inner = Inner::create(len);
    for msg in snapshot.0 {
        // We really can't do nothing if we fail here
        if inner.0.inner.push(msg).is_err() {
            return None;
        }
    }

    Some(inner)
}

/// Linked Node is a unit holding a message in a linked list data struct.
#[derive(Debug)]
struct LinkedNode<T> {
    /// The message itself.
    message: Option<T>,
    /// Next item in list
    next: AtomicPtr<LinkedNode<T>>
}

impl<T> LinkedNode<T> {
    #[inline(always)]
    fn new(message: T) -> Self {
        // Safety consideration
        // We check null pointer at every read
        LinkedNode {
            message: Some(message),
            next: AtomicPtr::new(std::ptr::null_mut())
        }
    }
}

/// Inner of unbounded channel
#[derive(Debug)]
pub struct InnerUnbounded<T> {
    /// Head of the linked list
    head: AtomicPtr<LinkedNode<T>>,
    /// Tail of linked list, exclusive access for reader only
    /// Here we used ptr of ptr to move tail ptr forward as we read from queue
    tail: *mut *mut LinkedNode<T>,
    /// Message count in queue
    len: AtomicU64,
    /// Number of active halfs
    active: AtomicI32,
    /// Flag to indicate if reader is active
    reader_active: AtomicBool,
    /// Reader waker
    waker: WakeSource
}

// Trust me!! I knew what I'm doing... probably...
unsafe impl<T> Send for InnerUnbounded<T> {}
unsafe impl<T> Sync for InnerUnbounded<T> {}

impl<T: Send + Sync> InnerUnbounded<T> {
    fn create() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
        let waker_sink = WakeSink::new();
        
        // Queue will always contain an element without being counted
        // This is used to always have a valid chain list without the need
        // to update tail pointer that is exclusive to reader
        let msg = Box::leak(Box::new(LinkedNode {
            message: None,
            next: AtomicPtr::new(std::ptr::null_mut())
        }));

        // Safety consideration
        // Null pointer are always checked
        let inner = Arc::new(InnerUnbounded {
            head: AtomicPtr::new(msg),
            tail: Box::leak(Box::new(msg as *mut LinkedNode<T>)),
            len: AtomicU64::new(0),
            active: AtomicI32::new(2),
            reader_active: AtomicBool::new(true),
            waker: waker_sink.source()
        });

        let writer = UnboundedSender { inner: inner.clone() };
        let reader = UnboundedReceiver { inner, waker: waker_sink };

        (writer, reader)
    }

    #[inline(always)]
    fn push_inner(&self, element: T) {
        // Create new element
        let msg = Box::leak(Box::new(LinkedNode::new(element)));
        // We first start by changing head pointer and retrieving previous pointer
        // This is safe to do because we are always sure that previous element won't get freed
        // by reader unless we update prev.next
        let prev = unsafe {
            self.head.swap(msg, Ordering::Relaxed).as_mut()
                .expect("Queue should always have one node")
        };

        prev.next.store(msg, Ordering::Relaxed);
        
        // Notify reader of new element
        self.waker.notify();
    }

    #[inline(always)]
    fn push(&self, element: T) -> Result<(), (SendError, T)> {
        if !self.reader_active.load(Ordering::Acquire) {
            return Err((SendError::Closed, element));
        }
    
        self.len.fetch_add(1, Ordering::Relaxed);
        self.push_inner(element);

        Ok(())
    }

    #[inline(always)]
    fn pop(&self) -> Result<T, TryRecvError> {
        // Safety: We never store an invalid tail pointer
        let tail_ref = unsafe {
            (*self.tail).as_mut()
                .expect("Tail pointer for mpsc should be valid")
        };

        let msg = tail_ref.next.load(Ordering::Relaxed);
        // If ptr is null that means there is no message to read
        if msg.is_null() {
            return if self.active.load(Ordering::Relaxed) == 1 {
                Err(TryRecvError::Closed)
            } else {
                Err(TryRecvError::Empty)
            }
        }

        // We move our tail pointer 1 step forward and we free old tail ref
        // Here box should take care of freeing memory
        let message;
        unsafe {
            _ = Box::from_raw(*(self.tail));
            (*self.tail) = msg;
            let tail_ref = msg.as_mut()
                .expect("Writer mpsc put invalid pointer as next element");
            message      = tail_ref.message
                .take()
                .expect("Writer mpsc put None as message value");
        };

        self.len.fetch_sub(1, Ordering::Relaxed);

        Ok(message)
    }

    fn cleanup(&self) {
        for _ in 0..self.len.load(Ordering::Relaxed) {
            _ = self.pop();
        };
        // Safety: We are freeing only valid pointers here
        unsafe {
            _ = Box::from_raw(*self.tail);
            _ = Box::from_raw(self.tail);
        }
    }

    /// Return number of active writers
    pub fn writers_count(&self) -> i32 {
        self.active.load(Ordering::Relaxed) - if self.reader_active.load(Ordering::Relaxed) { 1 } else { 0 }
    }

    /// Get currently present messages in channel
    pub fn len(&self) -> u64 {
        self.len.load(Ordering::Relaxed)
    }

    fn get_inner_ref(&self) -> &Self {
        self
    }

    /// In some cases, different channels halfs are stored together
    /// with no way to differentiate between them.
    /// this method is used to check if two halfs (either writer or reader)
    /// belong to the same channel
    pub fn same_channel<U: Deref<Target = Self>>(&self, other: &U) -> bool {
        std::ptr::eq(self.get_inner_ref(), other.get_inner_ref())
    }
}

/// Inner of channel
pub struct Inner<T> {
    inner: InnerUnbounded<T>,
    /// Capacity of the bound queue
    capacity: u64,
    /// Event to notify writers waiting until pushing new elements to queue is possible
    listener: Event
}

impl<T> std::fmt::Debug for Inner<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Inner")
    }
}

// Same thing goes here...
unsafe impl<T> Send for Inner<T> {}
unsafe impl<T> Sync for Inner<T> {}

impl<T: Send + Sync> Inner<T> {
    fn create(capacity: usize) -> (Sender<T>, Receiver<T>) {
        let waker_sink = WakeSink::new();
        
        // Queue will always contain an element without being counted
        // This is used to always have a valid chain list without the need
        // to update tail pointer that is exclusive to reader
        let msg = Box::leak(Box::new(LinkedNode {
            message: None,
            next: AtomicPtr::new(std::ptr::null_mut())
        }));

        // Safety consideration
        // Null pointer are always checked
        let inner = Arc::new(Inner {
            inner: InnerUnbounded {
                head: AtomicPtr::new(msg),
                tail: Box::leak(Box::new(msg as *mut LinkedNode<T>)),
                len: AtomicU64::new(0),
                active: AtomicI32::new(2),
                reader_active: AtomicBool::new(true),
                waker: waker_sink.source()
            },
            capacity: capacity as u64,
            listener: Event::new()
        });

        let writer = Sender { inner: inner.clone() };
        let reader = Receiver { inner, waker: waker_sink };

        (writer, reader)
    }

    #[inline(always)]
    fn push(&self, element: T) -> Result<(), (TrySendError, T)> {
        if !self.inner.reader_active.load(Ordering::Acquire) {
            return Err((TrySendError::Closed, element));
        }
        
        // Checking if queue is full before doing anything
        if self.inner.len.fetch_add(1, Ordering::Relaxed) >= self.capacity {
            self.inner.len.fetch_sub(1, Ordering::Relaxed);
            return Err((TrySendError::Full, element));
        }

        self.inner.push_inner(element);

        Ok(())
    }

    #[inline(always)]
    fn pop(&self) -> Result<T, TryRecvError> {
        let ret = self.inner.pop();

        // Notify waiting writers that a message has been read
        if ret.is_ok() {
            self.listener.notify_one();
        }

        ret
    }

    fn cleanup(&self) { self.inner.cleanup() }

    /// Return number of active writers
    pub fn writers_count(&self) -> i32 { self.inner.writers_count() }

    pub fn capacity(&self) -> u64 { self.capacity }

    /// Get currently present messages in channel
    pub fn len(&self) -> u64 { self.inner.len() }

    fn get_inner_ref(&self) -> &Self { self }

    /// In some cases, different channels halfs are stored together
    /// with no way to differentiate between them.
    /// this method is used to check if two halfs (either writer or reader)
    /// belong to the same channel
    pub fn same_channel<U: Deref<Target = Self>>(&self, other: &U) -> bool {
        std::ptr::eq(self.get_inner_ref(), other.get_inner_ref())
    }
}

impl<T: Send + Sync + Clone> InnerUnbounded<T> {
    #[inline(always)]
    fn snapshot(&self) -> Vec<T> {
        // Safety: No invalid pointer is ever stored at tail
        // Here we skip message at tail because it always contains None
        let mut iter = unsafe {
            (*self.tail).as_mut()
                .expect("Tail pointer for mpsc should be valid")
                .next
                .load(Ordering::Relaxed)
        };

        let mut messages = Vec::new();
        loop {
            if iter.is_null() {
                break;
            }
            // Safety: Only reader is ever allowed to free elements from queue
            let iter_ref = unsafe {
                iter.as_mut()
                    .expect("Message pointer in queue for mpsc should be valid")
            };
            // The only situation we ever have a None value is at tail which is already read
            // Here having a None is UB
            let msg = iter_ref.message
                .clone()
                .take()
                .expect("Message value in queue for mpsc shoud not be None");
            messages.push(msg);
            // Now we point to next element
            iter = iter_ref.next.load(Ordering::Relaxed);
        };

        messages
    }
}

impl<T: Send + Sync + Clone> Inner<T> {
    #[inline(always)]
    fn snapshot(&self) -> (Vec<T>, u64) {
        (self.inner.snapshot(), self.capacity)
    }
}

/// A channel writer half
#[derive(Debug)]
pub struct Sender<T: Send + Sync> {
    inner: Arc<Inner<T>>
}

impl<T: Send + Sync> Sender<T> {
    /// Push a new message to channel, and wait if channel if full
    pub async fn send(&self, mut element: T) -> Result<(), (SendError, T)> {
        loop {
            match self.try_send(element) {
                Ok(_) => return Ok(()),
                Err((TrySendError::Closed, msg)) => return Err((SendError::Closed, msg)),
                Err((TrySendError::Full, msg)) => {
                    element = msg;
                    self.inner.listener.wait_until(|| {
                        Some(())
                    }).await;
                }
            }
        }
    }

    /// Try pushing a new message to channel without awaiting
    /// This returns element on error
    pub fn try_send(&self, element: T) -> Result<(), (TrySendError, T)> {
        self.inner.push(element)
    }
}

impl<T: Send + Sync> Deref for Sender<T> {
    type Target = Inner<T>;
    fn deref(&self) -> &Self::Target {
        self.inner.as_ref()
    }
}

impl<T: Send + Sync> Clone for Sender<T> {
    fn clone(&self) -> Self {
        let oth = Sender { inner: self.inner.clone() };
        self.inner.inner.active.fetch_add(1, Ordering::Relaxed);
        oth
    }
}

impl<T: Send + Sync> Drop for Sender<T> {
    fn drop(&mut self) {
        // We must notify active readers when all writers are dropped
        // Otherwise they will hang undefinitely
        // Or cleanup channel if there is no other reader or writer besides us
        if self.inner.inner.active.fetch_sub(1, Ordering::Release) <= 1 {
            self.inner.inner.cleanup();
        } else {
            self.inner.inner.waker.notify();
        }
    }
}

/// An unbounded channel writer half
#[derive(Debug)]
pub struct UnboundedSender<T: Send + Sync> {
    inner: Arc<InnerUnbounded<T>>
}

impl<T: Send + Sync> UnboundedSender<T> {
    /// Push a new message to channel
    /// This is non-blocking because we never wait
    pub fn send(&self, element: T) -> Result<(), (SendError, T)> {
        self.inner.push(element)
    }
}

impl<T: Send + Sync> Deref for UnboundedSender<T> {
    type Target = InnerUnbounded<T>;
    fn deref(&self) -> &Self::Target {
        self.inner.as_ref()
    }
}

impl<T: Send + Sync> Clone for UnboundedSender<T> {
    fn clone(&self) -> Self {
        let oth = UnboundedSender { inner: self.inner.clone() };
        self.inner.active.fetch_add(1, Ordering::Relaxed);
        oth
    }
}

impl<T: Send + Sync> Drop for UnboundedSender<T> {
    fn drop(&mut self) {
        // We must notify active readers when all writers are dropped
        // Otherwise they will hang undefinitely
        // Or cleanup channel if there is no other reader or writer besides us
        if self.inner.active.fetch_sub(1, Ordering::Release) <= 1 {
            self.inner.cleanup();
        } else {
            self.inner.waker.notify();
        }
    }
}

/// A channel reader half.
#[derive(Debug)]
pub struct Receiver<T: Send + Sync> {
    /// Inner reference
    inner: Arc<Inner<T>>,
    /// Waker used to awake receiver task when new element is present
    waker: WakeSink
}

impl<T: Send + Sync> Receiver<T> {
    /// Try to read a message from channel without awaiting if no message
    /// is currently available to be read.
    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
        self.inner.pop()
    }

    /// Asynchronously read a message from channel. This does not return until there is a message
    /// available to read.
    pub async fn recv(&mut self) -> Result<T, RecvError> {
        let inner = &self.inner;
        self.waker.wait_until(move || {
            match inner.pop() {
                Ok(msg) => Some(Ok(msg)),
                Err(TryRecvError::Closed) => Some(Err(RecvError::Closed)),
                Err(TryRecvError::Empty) => None
            }
        }).await
    }
}

impl<T: Send + Sync + Clone> Receiver<T> {
    /// Create a snapshot of channel that can be used to create an identical channel using
    /// [`restore_channel_from_snapshot`]
    pub fn snapshot(&self) -> (Vec<T>, u64) {
        self.inner.snapshot()
    }
}

impl<T: Send + Sync> Deref for Receiver<T> {
    type Target = Inner<T>;
    fn deref(&self) -> &Self::Target {
        self.inner.as_ref()
    }
}

impl<T: Send + Sync> Drop for Receiver<T> {
    fn drop(&mut self) {
        // Cleanup channel if there is no other reader or writer besides us
        if self.inner.inner.active.fetch_sub(1, Ordering::Release) <= 1 {
            self.inner.cleanup();
        } else {
            self.inner.inner.reader_active.store(false, Ordering::Release);
            self.listener.notify_all();
        }
    }
}

/// An unbounded channel reader half.
#[derive(Debug)]
pub struct UnboundedReceiver<T: Send + Sync> {
    inner: Arc<InnerUnbounded<T>>,
    /// Waker used to awake receiver task when new element is present
    waker: WakeSink
}

impl<T: Send + Sync> UnboundedReceiver<T> {
    /// Try to read a message from channel without awaiting if no message
    /// is currently available to be read.
    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
        self.inner.pop()
    }

    /// Asynchronously read a message from channel. This does not return until there is a message
    /// available to read.
    pub async fn recv(&mut self) -> Result<T, RecvError> {
        let inner = &self.inner;
        self.waker.wait_until(move || {
            match inner.pop() {
                Ok(msg) => Some(Ok(msg)),
                Err(TryRecvError::Closed) => Some(Err(RecvError::Closed)),
                Err(TryRecvError::Empty) => None
            }
        }).await
    }
}

impl<T: Send + Sync + Clone> UnboundedReceiver<T> {
    /// Create a snapshot of channel that can be used to create an identical channel using
    /// [`restore_channel_from_snapshot`]
    pub fn snapshot(&self) -> Vec<T> {
        self.inner.snapshot()
    }
}

impl<T: Send + Sync> Deref for UnboundedReceiver<T> {
    type Target = InnerUnbounded<T>;
    fn deref(&self) -> &Self::Target {
        self.inner.as_ref()
    }
}

impl<T: Send + Sync> Drop for UnboundedReceiver<T> {
    fn drop(&mut self) {
        // Cleanup channel if there is no other reader or writer besides us
        if self.inner.active.fetch_sub(1, Ordering::Release) <= 1 {
            self.inner.cleanup();
        } else {
            self.inner.reader_active.store(false, Ordering::Release);
        }
    }
}

/// An error that may be returned by calling [`Sender::try_send`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum TrySendError {
    /// Channel is full, cannot push new message to channel unless receiver reads older messages.
    Full,
    /// Reader is no longer active
    Closed
}

impl std::error::Error for TrySendError {}

impl std::fmt::Display for TrySendError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match *self {
            TrySendError::Full => write!(f, "sending to a full channel"),
            TrySendError::Closed => write!(f, "sending to a closed channel")
        }
    }
}

/// An error that may be returned by calling [`Sender::send`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum SendError {
    /// Reader is no longer active
    Closed
}

impl std::error::Error for SendError {}

impl std::fmt::Display for SendError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match *self {
            SendError::Closed => write!(f, "sending to a closed channel")
        }
    }
}

/// An error that may be returned by calling [`Receiver::try_recv`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum TryRecvError {
    /// Queue is empty or we read all available messages.
    Empty,
    /// No writer is active, implying no messages will ever be pushed to queue.
    /// This only happens when all messages in queue are read.
    Closed
}

impl std::error::Error for TryRecvError {}

impl std::fmt::Display for TryRecvError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match *self {
            TryRecvError::Empty => write!(f, "receiving from an empty channel"),
            TryRecvError::Closed => write!(f, "receiving from an empty and closed channel")
        }
    }
}

/// An error that may be returned by calling [`Receiver::recv`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum RecvError {
    /// No writer is active, implying no messages will ever be pushed to queue.
    /// This only happens when all messages in queue are read.
    Closed
}

impl std::error::Error for RecvError {}

impl std::fmt::Display for RecvError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match *self {
            RecvError::Closed => write!(f, "receiving from an empty and closed channel")
        }
    }
}
