//! A single-producer singe-consumer channel used to send single message or signal.
//! 
//! ```
//! use qanat::oneshot::{self, TryRecvError};
//!
//! // Create new channel
//! let (tx, mut rx) = oneshot::channel();
//!
//! /// Checking if message is already sent
//! assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
//!
//! /// Sending message
//! tx.send(1).unwrap();
//!
//! /// Receiving message
//! assert_eq!(rx.try_recv(), Ok(1));
//! /// Retrying
//! assert_eq!(rx.try_recv(), Err(TryRecvError::Closed));
//! ```
use std::{
    cell::UnsafeCell,
    sync::{atomic::{Ordering, AtomicU8}, Arc}
};

use diatomic_waker::{WakeSink, WakeSource};

/// Create a new oneshot channel where only a single message is sent
/// in the lifetime of the channel
pub fn channel<T: Send + Sync>() -> (Sender<T>, Receiver<T>) {
    Inner::create()
}

/// Container holding message
#[derive(Debug)]
struct Inner<T> {
    /// The message itself.
    message: UnsafeCell<Option<T>>,
    /// Channel state:
    /// 0 = Nothing happened yet
    /// 1 = Channel has message and write half is dropped
    /// 2 = Either halfs was dropped
    channel_state: AtomicU8,
    /// Source for notfying reader when message is in channel
    waker: WakeSource
}

// ...
unsafe impl<T> Send for Inner<T> {}
unsafe impl<T> Sync for Inner<T> {}

impl<T: Send + Sync> Inner<T> {
    fn create() -> (Sender<T>, Receiver<T>) {
        let waker_sink = WakeSink::new();

        let inner = Arc::new(Inner {
            message: UnsafeCell::new(None),
            channel_state: AtomicU8::new(0),
            waker: waker_sink.source()
        });

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

        (writer, reader)
    }

    #[inline(always)]
    fn push(&self, element: T) -> Result<(), (SendError, T)> {
        if self.channel_state.load(Ordering::Acquire) == 2 {
            return Err((SendError::Closed, element));
        }

        // Safety: we are guarding this with channel_state so it should be safe
        unsafe {
            *self.message.get() = Some(element);
        }

        // We omit case here where reader drops right after we did write message
        self.channel_state.swap(1, Ordering::Release);
        
        // Also notify if reader is waiting
        self.waker.notify();

        Ok(())
    }

    #[inline(always)]
    fn pop(&self) -> Result<T, TryRecvError> {
        match self.channel_state.load(Ordering::Acquire) {
            0 => return Err(TryRecvError::Empty),
            1 => (),
            2 => return Err(TryRecvError::Closed),
            _ => unreachable!()
        };

        let message;
        unsafe {
            message = (*self.message.get())
                .take()
                .expect("Oneshot message should not be None at this point");
        }

        Ok(message)
    }
}

/// Channel writer half
#[derive(Debug)]
pub struct Sender<T: Send + Sync> {
    inner: Arc<Inner<T>>,
    /// Did we write yet?
    written: bool
}

impl<T: Send + Sync> Sender<T> {
    pub fn send(mut self, element: T) -> Result<(), (SendError, T)> {
        self.written = true;
        self.inner.push(element)
    }
}

impl<T: Send + Sync> Drop for Sender<T> {
    fn drop(&mut self) {
        if !self.written {
            self.inner.channel_state.store(2, Ordering::Release);
        }
        // Notify if reader is waiting, we don't want to deadlock
        self.inner.waker.notify();
    }
}

/// Channel reader half
#[derive(Debug)]
pub struct Receiver<T: Send + Sync> {
    inner: Arc<Inner<T>>,
    /// Notifyer when message is pushed to channel
    waker: WakeSink,
    /// Flag not to read again
    read: bool
}

impl<T: Send + Sync> Receiver<T> {
    /// Try to read message from channel without awaiting if message is not yet in channel
    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
        if self.read {
            return Err(TryRecvError::Closed);
        }
        let ret = self.inner.pop();

        match ret.as_ref() {
            Ok(_) | Err(TryRecvError::Closed) => self.read = true,
            _ => ()
        }

        ret
    }
    
    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> Drop for Receiver<T> {
    fn drop(&mut self) {
        self.inner.channel_state.store(2, Ordering::Release);
    }
}

/// 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 awaiting [`Receiver`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum RecvError {
    /// Writer is no longer active
    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")
        }
    }
}

/// An error that may be returned by calling [`Receiver::try_recv`]
///
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum TryRecvError {
    /// Message still not sent yet
    Empty,
    /// Writer is no longer active
    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")
        }
    }
}
