//! A broadcast channel is where all readers receive every message pushed to channel.
//!
//! When channel reaches full capacity, oldest message is dropped and any reader that that
//! is not up to date will lag out and get moved to oldest new message in channel.
//! 
//!
//! The queue of the channel is internally represented as a linked-list, relying on atomic
//! only operations for concurrent lock-free access, and seize (hyaline) for memory reclamation
//! Reads are wait-free most of time, while writers have slow
//! performance degradation under contention.
//!
//! Note: while all read/write operations are atomic, there are other functionality that comes with
//! this channel like resizing channel and creating a snapshot for channel, if you need to use methods
//! related to that, you should be aware of your use case, in case you want to call resize methods
//! in different threads, you should always use the lock method version.
//!
//! ```
//! use qanat::broadcast;
//!
//! // Create new channel
//! let (mut tx, mut rx) = broadcast::channel(10.try_into().unwrap());
//!
//! // broadcast a message
//! tx.send(1);
//!
//! // Create new reader
//! let mut rx1 = rx.clone();
//!
//! // Reading from channel
//! assert_eq!(rx.try_recv().unwrap(), 1);
//! assert_eq!(rx1.try_recv().unwrap(), 1);
//!
//! // Changing channel capacity
//! tx.set_capacity(20.try_into().unwrap());
//! // Changing channel capacity when multithread call to set_capacity is needed
//! //tx.set_capacity_locked(20.try_into().unwrap()).await;
//!
//! tx.send(2);
//!
//! // Getting last index read
//! let resume_point = rx.read_position();
//!
//! // Creating a snapshot for channel
//! let snapshot = tx.snapshot();
//! // Same thing as set_capacity when multithread call to snapshot is needed
//! //let snapshot = tx.snapshot_locked().await;
//!
//! // Dropping channel
//! drop(tx);
//! drop(rx);
//! drop(rx1);
//!
//! // Creating a new channel from snapshot of previous one
//! let (mut tx, mut rx) = broadcast::restore_channel_from_snapshot(snapshot);
//!
//! // We can now restore and read from where we stopped
//! rx.restore(resume_point);
//! 
//! assert_eq!(rx.try_recv().unwrap(), 2);
//!
//!
//!
//! ```

use std::{
    sync::{
        atomic::{Ordering, AtomicU64, AtomicI32, AtomicPtr},
        Arc, Weak
    },
    num::NonZeroUsize, mem::ManuallyDrop,
    ops::Deref
};
use async_event::Event;
use seize::{reclaim, Collector, Guard, Linked};

use crate::lock::Lock;

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

/// Create a new channel from a snapshot returned by [`Sender::snapshot`] or [`Sender::snapshot_locked`]
pub fn restore_channel_from_snapshot<T: Clone + Send + Sync>(snapshot: (Vec<(u64, T)>, u64, u64)) -> (Sender<T>, Receiver<T>) {
    let inner = Inner::create(snapshot.1 as usize);
    for (id, msg) in snapshot.0 {
        inner.0.inner.push_snapshot(id, msg);
    }
    inner.0.inner.index.store(snapshot.2, Ordering::Relaxed);

    inner
}

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

impl<T: Clone> LinkedNode<T> {
    #[inline(always)]
    fn new(index: u64, message: T) -> Self {
        LinkedNode {
            index,
            message,
            next: ManuallyDrop::new(Arc::new(AtomicPtr::new(std::ptr::null_mut())))
        }
    }
}

/// Inner of channel
pub struct Inner<T: Clone> {
    collector: Collector,
    /// Index of latest message in the queue.
    /// Used as notification for readers to knew when there is a new message in the queue. And a
    /// way to synchronise indexes between writers.
    index: AtomicU64,
    /// Tail of the linked list
    tail: AtomicPtr<Linked<LinkedNode<T>>>,
    /// Head of the linked list
    head: AtomicPtr<Linked<LinkedNode<T>>>,
    /// Capacity of the bound queue
    pub(crate) capacity: AtomicU64,
    /// Message count in queue
    len: AtomicU64,
    /// Number of active writers
    writers: AtomicI32,
    /// Number of active readers
    readers: AtomicI32,
    /// Holding last popped message so reference of next message is not destroyed
    dead_msg: AtomicPtr<Linked<LinkedNode<T>>>,
    /// Event to notify readers when new element is ready
    listener: Event,
    /// Lock used to synchronise resize capacity transactions between different threads
    resize_lock: Lock
}

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

impl<T: Clone + Send + Sync> Inner<T> {
    fn create(capacity: usize) -> (Sender<T>, Receiver<T>) {
        // Safety consideration
        // Null pointer are always checked
        let inner = Arc::new(Inner {
            collector: Collector::new(),
            index: AtomicU64::new(0),
            tail: AtomicPtr::new(std::ptr::null_mut()),
            head: AtomicPtr::new(std::ptr::null_mut()),
            capacity: AtomicU64::new(capacity as u64),
            len: AtomicU64::new(0),
            writers: AtomicI32::new(1),
            readers: AtomicI32::new(1),
            dead_msg: AtomicPtr::new(std::ptr::null_mut()),
            listener: Event::new(),
            resize_lock: Lock::new()
        });

        let writer = Sender {
            inner: inner.clone(),
        };

        let reader = Receiver {
            inner,
            pos: LinkedNodeReadRef { index: 0, refe: LinkedNodeReadPtr::None }
        };

        (writer, reader)
    }

    #[inline(always)]
    fn find_next_read_pos(&self, node: &LinkedNode<T>, pos: &mut LinkedNodeReadPtr<T>) {
        *pos = LinkedNodeReadPtr::NextMsgRef(Arc::downgrade(&node.next));
    }

    #[inline(always)]
    #[allow(dangerous_implicit_autorefs)]
    fn push(&self, element: T) {
        let (index, _) = self.index.fetch_add(1, Ordering::Relaxed).overflowing_add(1);
        let message    = self.collector.link_boxed(LinkedNode::new(index, element));

        // Safety: We are not puting any invalid pointer in this operation
        unsafe {
            let changed = self.head.swap(message, Ordering::Relaxed);
            if changed.is_null() {
                self.tail.store(message, Ordering::Relaxed);
            } else {
                (*changed).next.store(message, Ordering::Relaxed);
            }
        }
        self.listener.notify_all();

        let len = self.len.fetch_add(1, Ordering::Relaxed);
        // Each new allocation release 1 element from tail
        if len >= self.capacity.load(Ordering::Relaxed) {
            self.pop(false);
        }
    }

    #[allow(dangerous_implicit_autorefs)]
    fn push_snapshot(&self, id: u64, msg: T) {
        let message = self.collector.link_boxed(LinkedNode::new(id, msg));
        // Safety: We are not puting any invalid pointer in this CAS operation
        unsafe {
            let changed = self.head.swap(message, Ordering::Relaxed);
            if changed.is_null() {
                self.tail.store(message, Ordering::Relaxed);
            } else {
                (*changed).next.store(message, Ordering::Relaxed);
            }
        }
        self.len.fetch_add(1, Ordering::Relaxed);
    }

    #[inline(always)]
    #[allow(dangerous_implicit_autorefs)]
    fn pop(&self, cleanup: bool) {
        let guard = self.collector.enter();
        loop {
            // We can safely unwrap here becasue inserting to queue starts first with modifying
            // next, there will never be a case where next will be none here
            let tail = guard.protect(&self.tail, Ordering::Acquire);
            // Here we don't need to safe load next because we already have a guard for tail
            let next = unsafe { (*tail).next.load(Ordering::Relaxed) };
            if next.is_null() && !cleanup {
                // Next is not set yet
                continue;
            }
            // Safety consideration: We are completely safe, tail is guarded so next ptr is
            // also valid
            let changed = self.tail.compare_exchange(tail, next,
                                                     Ordering::Release, Ordering::Relaxed);
            if let Ok(ptr) = changed {
                // we successfully removed tail
                self.len.fetch_sub(1, Ordering::Relaxed);

                if !ptr.is_null() {
                    loop {
                        let old = self.dead_msg.load(Ordering::Acquire);
                        if old.is_null() {
                            if self.dead_msg.compare_exchange(std::ptr::null_mut(), ptr,
                                                              Ordering::Release, Ordering::Relaxed).is_ok() {
                                break;
                            }
                        } else {
                            // If dead_msg.next == tail then we swap
                            let next = unsafe { (*old).next.load(Ordering::Relaxed) };
                            if std::ptr::eq(next, ptr) {
                                // This should always succeed, because other writers need to wait
                                let old = self.dead_msg.swap(ptr, Ordering::Relaxed);
                                if !old.is_null() {
                                    // Safety: Here we free Arc ref so lagging readers will knew that they did
                                    // fall behind, while also being sure it won't be freed again because we
                                    // have ManuallyDrop.
                                    // And after that we retire object, it is also safe because we only ever
                                    // retire object once
                                    unsafe {
                                        ManuallyDrop::drop(&mut (*old).next);
                                        self.collector.retire(old, reclaim::boxed::<Linked<LinkedNode<T>>>);
                                    }
                                }
                                break;
                            }
                        }
                    }
                }
                break;
            }
            // otherwise we retry
        }
    }

    #[inline(always)]
    fn try_recv(&self, pos: &mut LinkedNodeReadRef<T>) -> Result<T, TryRecvError> {
        let guard = self.collector.enter();
        match &pos.refe {
            LinkedNodeReadPtr::NextMsgRef(node_ref) => {
                match node_ref.upgrade() {
                    Some(node_ref) => {
                        let node = guard.protect(&node_ref, Ordering::Acquire);
                        if node.is_null() {
                            if self.writers.load(Ordering::Relaxed) == 0 {
                                Err(TryRecvError::Closed)
                            } else {
                                Err(TryRecvError::Empty)
                            }
                        } else {
                            // Node is guarded so it should be safe to access
                            let node = unsafe { &*node };
                            self.find_next_read_pos(node, &mut pos.refe);
                            pos.index = node.index;
                            Ok(node.message.clone())
                        }
                    },
                    None => {
                        // We're too late to read, moving to tail of the queue
                        pos.refe = LinkedNodeReadPtr::None;
                        Err(TryRecvError::Lagged)
                    }
                }
            },
            LinkedNodeReadPtr::None => {
                // First time to read or we lagged before this
                // Fetching tail
                let node = guard.protect(&self.tail, Ordering::Acquire);
                if node.is_null() {
                    if self.writers.load(Ordering::Relaxed) == 0 {
                        Err(TryRecvError::Closed)
                    } else {
                        Err(TryRecvError::Empty)
                    }
                } else {
                    // Same thing here
                    let node  = unsafe { &*node };
                    self.find_next_read_pos(node, &mut pos.refe);
                    pos.index = node.index;
                    Ok(node.message.clone())
                }
            }
        }
    }

    #[inline]
    fn set_capacity(&self, capacity: NonZeroUsize) {
        let new_cap      = Into::<usize>::into(capacity) as u64;
        let len          = self.len.load(Ordering::Relaxed);
        let old_capacity = self.capacity.swap(new_cap, Ordering::Relaxed);

        if new_cap >= old_capacity {
            // When new capaciry is bigger or equal we do nothing
            return;
        }

        // We will remove exactly len - new_cap message from queue tail
        // That is if current length of queue is bigger than new_cap
        if len > new_cap {
            for _ in 0..(len - new_cap) {
                self.pop(false);
            }
        }
    }

    #[allow(dangerous_implicit_autorefs)]
    fn cleanup(&self) {
        let len = self.len.load(Ordering::Relaxed);
        for _ in 0..len {
            self.pop(true);
        }
        // Safety: We are checking if it is a null pointer
        // We need to retire head to reclaim last message in queue
        let last = self.head.swap(std::ptr::null_mut(), Ordering::Relaxed);
        if !last.is_null() {
            unsafe {
                ManuallyDrop::drop(&mut (*last).next);
                self.collector.retire(last, reclaim::boxed::<Linked<LinkedNode<T>>>);
            }
        }
        // No need to reclaim from self.dead_msg because it's holding same message as head
    }

    /// Return last message index inserted to channel
    pub fn last_index(&self) -> u64 {
        self.index.load(Ordering::Relaxed)
    }

    /// Return number of active readers
    pub fn readers_count(&self) -> i32 {
        self.readers.load(Ordering::Relaxed)
    }

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

    /// Get channel capacity
    pub fn capacity(&self) -> u64 {
        self.capacity.load(Ordering::Relaxed)
    }

    /// 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())
    }
}


#[derive(Debug)]
pub struct Sender<T: Clone + Send + Sync> {
    inner: Arc<Inner<T>>
}

impl<T: Clone + Send + Sync> Sender<T> {
    /// Push a new message to channel
    pub fn send(&mut self, element: T) {
        self.inner.push(element);
    }
    
    /// Change channel capacity, this uses a lock to synchronise concurrent capacity changing
    /// requests from multiple threads.
    ///
    /// Warning: If you need to change capacity from multiple threads, you must use THIS instead of
    /// [`Sender::set_capacity`].
    pub async fn set_capacity_locked(&mut self, capacity: NonZeroUsize) {
        let guard = self.inner.resize_lock.lock().await;
        self.inner.set_capacity(capacity);
        drop(guard);
    }

    /// Change channel capacity in a single thread context.
    ///
    /// Warning: This method must not be called to from multiple threads as it leads to indefined
    /// behavior.
    pub fn set_capacity(&mut self, capacity: NonZeroUsize) {
        self.inner.set_capacity(capacity)
    }

    /// Create a snapshot of channel that can be used to create an identical channel
    ///
    /// Warning: If you need to take snapshots from different threads you must use THIS method
    /// insted of [`Sender::snapshot`].
    ///
    /// return type (Vec<(index, message)>, channel_capacity, last_index)
    pub async fn snapshot_locked(&self) -> (Vec<(u64, T)>, u64, u64) {
        let guard = self.inner.resize_lock.lock().await;
        let ret   = self.snapshot();
        drop(guard);
        ret
    }

    /// Create a snapshot of channel that can be used to create an identical channel using
    /// [`restore_channel_from_snapshot`]
    ///
    /// Warning: Doesn't support multithreaded access, and [`Sender::snapshot_locked`] must be used
    /// instead.
    ///
    /// return type (Vec<(index, message)>, channel_capacity, last_index)
    pub fn snapshot(&self) -> (Vec<(u64, T)>, u64, u64) {
        let mut reader = Receiver {
            inner: self.inner.clone(),
            pos: LinkedNodeReadRef { index: 0, refe: LinkedNodeReadPtr::None }
        };
        let mut messages = Vec::new();
        let mut last_ind = 0;
        loop {
            match reader.try_recv() {
                Ok(msg) => {
                    if last_ind < reader.read_position() {
                        last_ind = reader.read_position();
                    }
                    messages.push((reader.read_position(), msg));
                },
                // Here only 2 types of errors may ever happen either lagged or empty
                // Because we are holding into a writer ref
                Err(TryRecvError::Lagged) => (),
                Err(_) => break
            }
        };

        (
            messages,
            self.inner.capacity.load(Ordering::Relaxed),
            last_ind
        )
    }
}

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

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

impl<T: Clone + 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.writers.fetch_sub(1, Ordering::Release) <= 1 {
            if self.inner.readers.load(Ordering::Acquire) == 0 {
                self.inner.cleanup();
            } else {
                self.inner.listener.notify_all();
            }
        }
    }
}

#[derive(Debug)]
pub struct LinkedNodeReadRef<T: Clone> {
    index: u64,
    refe: LinkedNodeReadPtr<T>
}

impl<T: Clone> Clone for LinkedNodeReadRef<T> {
    fn clone(&self) -> Self {
        LinkedNodeReadRef {
            index: self.index,
            refe: match &self.refe {
                LinkedNodeReadPtr::None => LinkedNodeReadPtr::None,
                LinkedNodeReadPtr::NextMsgRef(weak) => LinkedNodeReadPtr::NextMsgRef(weak.clone())
            }
        }
    }
}

/// An enum for different state next element may be in
#[derive(Debug)]
enum LinkedNodeReadPtr<T: Clone> {
    /// Weak ref of next message to read
    NextMsgRef(Weak<AtomicPtr<Linked<LinkedNode<T>>>>),
    /// This is the initial when nothing is read yet
    /// Or when lag happens
    None
}

/// A channel reader half.
#[derive(Debug)]
pub struct Receiver<T: Clone + Send + Sync> {
    /// Inner reference
    inner: Arc<Inner<T>>,
    /// We keep track of next element to read by a weak ref
    pos: LinkedNodeReadRef<T>,
}

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

impl<T: Clone + 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.try_recv(&mut self.pos)
    }

    /// 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> {
        loop {
            match self.try_recv() {
                Ok(msg) => return Ok(msg),
                Err(TryRecvError::Closed) => return Err(RecvError::Closed),
                Err(TryRecvError::Lagged) => return Err(RecvError::Lagged),
                Err(TryRecvError::Empty) => {
                    self.inner.listener.wait_until(|| {
                        if self.pos.index != self.inner.index.load(Ordering::Relaxed)
                            || self.inner.writers.load(Ordering::Relaxed) == 0 {
                            Some(())
                        } else {
                            None
                        }
                    }).await;
                }
            }
        }
    }

    /// Return last message index read
    ///
    /// Handy when restoring channel from a snapshot while preserving read position
    pub fn read_position(&self) -> u64 {
        self.pos.index
    }

    /// Go back to exactly last position read before when using snapshot/restore, this method must
    /// not be called on newly created channel that is not restored from a snapshot
    pub fn restore(&mut self, index: u64) {
        // Set to default in case this is a clone of another receiver
        self.pos.index = 0;
        self.pos.refe  = LinkedNodeReadPtr::None;

        loop {
            match self.try_recv() {
                Ok(_) => {
                    if self.pos.index == index {
                        break;
                    }
                },
                Err(TryRecvError::Lagged) => {
                    // Happens when we first read from channel
                    continue
                },
                Err(TryRecvError::Empty) => {
                    // We read pretty much all the messages in the channel but didn't find our
                    // wanted index, that means we either lagged or last index is exactly
                    // the index of last removed message from queue
                    self.pos.index = index;
                    self.pos.refe  = LinkedNodeReadPtr::None;
                    break;
                },
                Err(TryRecvError::Closed) => {
                    // Nothing to do here
                    break;
                },
            }
        };
    }
}

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

impl<T: Clone + 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.readers.fetch_sub(1, Ordering::Release) <= 1
            && self.inner.writers.load(Ordering::Acquire) == 0 {
            self.inner.cleanup();
        }
    }
}

/// An error that may be returned by calling [`Receiver::try_recv`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum TryRecvError {
    /// This received lagged too far behind.
    /// Trying to receive after this will return last message in queue.
    ///
    /// To avoid more synchronisation under writers high contention this may return a negative number if
    /// capacity was just changed.
    Lagged,
    /// 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"),
            TryRecvError::Lagged => write!(f, "receiving operation observed lost messages")
        }
    }
}

/// An error that may be returned by calling [`Receiver::recv`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum RecvError {
    /// This received lagged too far behind.
    /// Trying to receive after this will return last message in queue.
    ///
    /// To avoid more synchronisation under writers high contention this may return a negative number if
    /// capacity was just changed.
    Lagged,
    /// 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"),
            RecvError::Lagged => write!(f, "receiving operation observed lost messages")
        }
    }
}
