use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::ops::Deref;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

// Represents a single producer, single consumer queue
pub struct Stream<T>
{
    buffer: Box<UnsafeCell<[MaybeUninit<T>]>>,
    capacity_mask: usize,

    head: CachePadded<AtomicUsize>,
    tail: CachePadded<AtomicUsize>,
}

#[repr(align(64))]
pub struct CachePadded<T>(T);

impl<T> CachePadded<T>
{
    pub fn new(element: T) -> Self
    {
        CachePadded(element)
    }
}

impl<T> Deref for CachePadded<T>
{
    type Target = T;

    fn deref(&self) -> &Self::Target
    {
        &self.0
    }
}

pub struct StreamProducer<T>
{
    inner: Arc<Stream<T>>,
}

pub struct StreamConsumer<T>
{
    inner: Arc<Stream<T>>,
}

unsafe impl<T: Send> Send for StreamProducer<T> {}
unsafe impl<T: Send> Sync for StreamProducer<T> {}

unsafe impl<T: Send> Send for StreamConsumer<T> {}
unsafe impl<T: Send> Sync for StreamConsumer<T> {}

#[repr(transparent)]
pub struct Takable<T>(MaybeUninit<T>);

impl<T> Takable<T>
{
    pub fn new(element: T) -> Self
    {
        Takable(MaybeUninit::new(element))
    }

    pub unsafe fn take(&mut self) -> T
    {
        unsafe { std::mem::replace(&mut self.0, MaybeUninit::uninit()).assume_init() }
    }

    pub unsafe fn peek(&self) -> &T
    {
        unsafe { self.0.assume_init_ref() }
    }

    pub unsafe fn peek_mut(&mut self) -> &mut T
    {
        unsafe { self.0.assume_init_mut() }
    }
}

unsafe fn takable_slice_from_maybe_uninitt<T>(slice: &mut [MaybeUninit<T>]) -> &mut [Takable<T>]
{
    unsafe { std::mem::transmute(slice) }
}

pub fn bounded_queue<T>(capacity: usize) -> (StreamProducer<T>, StreamConsumer<T>)
{
    // Require stream capacity to be a power of two for fast modulo computation
    assert!(
        capacity.count_ones() == 1,
        "Stream capacity must be a power of 2."
    );
    let slice = (0..capacity)
        .map(|_| MaybeUninit::uninit())
        .collect::<Vec<_>>()
        .into_boxed_slice();
    let cell = UnsafeCell::from_mut(Box::leak(slice));

    // SAFETY:
    // Cell was created in a box, and UnsafeCell is transparent,
    // the memory layout pointed to by cell is the one a box would thus
    // expect
    let buffer = unsafe { Box::from_raw(cell as *mut UnsafeCell<[MaybeUninit<T>]>) };

    let queue = Arc::new(Stream {
        buffer,
        capacity_mask: capacity - 1,
        head: CachePadded::new(0.into()),
        tail: CachePadded::new(0.into()),
    });

    (
        StreamProducer {
            inner: queue.clone(),
        },
        StreamConsumer {
            inner: queue.clone(),
        },
    )
}

pub struct StreamReader<'a, T>
{
    slices: (&'a mut [Takable<T>], &'a mut [Takable<T>]),

    // UNSAFE !
    consumer: &'a mut StreamConsumer<T>,
}

pub struct StreamWriter<'a, T>
{
    slices: (&'a mut [MaybeUninit<T>], &'a mut [MaybeUninit<T>]),

    // UNSAFE !
    producer: &'a mut StreamProducer<T>,
}

impl<'a, T> StreamReader<'a, T>
{
    pub fn slices(&self) -> (&[Takable<T>], &[Takable<T>])
    {
        (self.slices.0, self.slices.1)
    }

    pub fn slices_mut(&mut self) -> (&mut [Takable<T>], &mut [Takable<T>])
    {
        (self.slices.0, self.slices.1)
    }

    pub fn consume(self, amount: usize)
    {
        self.consumer.consume(amount);
    }
}

impl<'a, T> StreamWriter<'a, T>
{
    pub fn slices(&self) -> (&[MaybeUninit<T>], &[MaybeUninit<T>])
    {
        (self.slices.0, self.slices.1)
    }

    pub fn slices_mut(&mut self) -> (&mut [MaybeUninit<T>], &mut [MaybeUninit<T>])
    {
        (self.slices.0, self.slices.1)
    }

    pub fn produce(self, amount: usize)
    {
        self.producer.produce(amount);
    }
}

impl<T> StreamProducer<T>
{
    pub fn first_index(&self) -> usize
    {
        self.inner.head.load(Ordering::Relaxed)
    }

    pub fn produce(&mut self, written: usize)
    {
        // Advance head.
        let head = self.inner.head.load(Ordering::Relaxed);
        let tail = self.inner.tail.load(Ordering::Relaxed);

        // Check bounds
        assert!(head + written - tail <= (self.inner.capacity_mask + 1));

        // We want writes to the buffer to be visible when acquired in the pop side
        self.inner.head.store(head + written, Ordering::Release);
    }

    pub fn write<'a>(&'a mut self) -> StreamWriter<'a, T>
    {
        // We need to claim the maximum amount of elements.
        let tail = self.inner.tail.load(Ordering::Acquire);
        let head = self.inner.head.load(Ordering::Relaxed);

        let wrapped_tail = tail & self.inner.capacity_mask;
        let wrapped_head = head & self.inner.capacity_mask;
        // While we have this information, the back can progress, but never in an unsafe way
        // Since it can only free up spots
        if head == tail
        {
            // Current configuration
            //
            // ▯▯▯▯▯▯▯▯▯▯▯▯▯
            //       |
            //  tail & head
            //    (empty)
            // ______.______
            // slice2 slice1

            // SAFETY:
            //
            // Since the StreamWriter object borrows the producer
            // we will return the only, mutually exclusive, mutable refences to the empty part of the buffer
            //
            // We know as per the previous diagram
            // That both slices are mutually exclusive
            // and that tail and head are indices of the slice
            unsafe {
                let k = &mut *self.inner.buffer.get();

                let (start_to_head, head_to_end) = k.split_at_mut_unchecked(wrapped_head);

                // This functions borrows the stream mutably. As such, only one instance
                // of these slices can exist for a given stream.
                StreamWriter {
                    slices: (start_to_head, head_to_end),
                    producer: self,
                }
            }
        }
        else
        {
            // We MUST have : tail < head
            if wrapped_tail < wrapped_head
            {
                // Current configuration :
                // ▯▯▯▮▮▮▮▮▮▯▯▯▯
                //    |     |
                // tail     head
                // ___      ____
                // slice2  slice1

                // SAFETY:
                //
                // Since the StreamWriter object borrows the producer
                // we will return the only, mutually exclusive, mutable refences to the empty part of the buffer
                //
                // We know as per the previous diagram
                // That both slices are mutually exclusive
                // and that tail and head are indices of the slice
                unsafe {
                    let k = &mut *self.inner.buffer.get();

                    let (start_to_head, head_to_end) = k.split_at_mut_unchecked(wrapped_head);
                    let (start_to_tail, _tail_to_head) =
                        start_to_head.split_at_mut_unchecked(wrapped_tail);

                    StreamWriter {
                        slices: (head_to_end, start_to_tail),
                        producer: self,
                    }
                }
            }
            else
            {
                // Current configuration :
                // ▮▮▮▯▯▯▯▯▯▮▮▮▮
                //    |     |
                // head     tail
                //     _____
                //     slice1
                //
                //      Or
                // ▮▮▮▮▮▮▮▮▮▮▮▮▮
                //       |
                //  tail & head
                //    (full)
                //       .
                //       slice1

                // SAFETY:
                //
                // Since the StreamWriter object borrows the producer
                // we will return the only mutable refence to the empty part of the buffer
                //
                // Head and tail are both indices of the slice
                unsafe {
                    let k = &mut *self.inner.buffer.get();
                    let (start_to_tail, _tail_to_end) = k.split_at_mut_unchecked(wrapped_tail);
                    let (_start_to_head, head_to_tail) = start_to_tail.split_at_mut_unchecked(wrapped_head);
                    let (empty_slice, head_to_tail) = head_to_tail.split_at_mut_unchecked(0);

                    StreamWriter {
                        slices: (head_to_tail, empty_slice),
                        producer: self,
                    }
                }
            }
        }
    }
}

impl<T> StreamConsumer<T>
{
    pub fn consume(&mut self, read: usize)
    {
        // Advance head.
        let head = self.inner.head.load(Ordering::Relaxed);
        let tail = self.inner.tail.load(Ordering::Relaxed);

        // Check bounds
        assert!(tail + read <= head);

        // We want writes to the buffer to be visible when acquired in the pop side
        self.inner.tail.store(tail + read, Ordering::Release);
    }

    pub fn first_index(&self) -> usize
    {
        self.inner.tail.load(Ordering::Relaxed)
    }

    pub fn read_takable<'a>(&'a mut self) -> StreamReader<'a, T>
    {
        // We need to claim the maximum amount of elements.
        let head = self.inner.head.load(Ordering::Acquire);
        let tail = self.inner.tail.load(Ordering::Relaxed);

        let wrapped_tail = tail & self.inner.capacity_mask;
        let wrapped_head = head & self.inner.capacity_mask;

        // While we have this information, the back can progress, but never in an unsafe way
        // Since it can only free up spots

        if tail == head
        {
            // Buffer is empty. Return empty slice
            unsafe {
                let k = &mut *self.inner.buffer.get();
                let empty = &mut k[0..0];
                let (empty_1, empty_2) = empty.split_at_mut_unchecked(0);
                StreamReader {
                    slices: (
                        takable_slice_from_maybe_uninitt(empty_1),
                        takable_slice_from_maybe_uninitt(empty_2),
                    ),
                    consumer: self,
                }
            }
        }
        else
        {
            // Necessarly:  wrapped_tail <= wrapped_head
            // Two cases : The buffer overlaps the wrapping or not
            if wrapped_tail < wrapped_head
            {
                // Current configuration :
                // ▯▯▯▮▮▮▮▮▮▯▯▯▯
                //    |     |
                // tail     head
                //    _______
                //    slice1

                // SAFETY:
                //
                // Since the StreamWriter object borrows the producer
                // we will return the only mutable refence to the empty part of the buffer
                //
                // Head and tail are both indices of the slice
                unsafe {
                    let k = &mut (&mut *self.inner.buffer.get())[wrapped_tail..wrapped_head];
                    let (tail_to_head, empty_slice) =
                        k.split_at_mut_unchecked(wrapped_head - wrapped_tail);
                    assert_eq!(empty_slice.len(), 0);

                    StreamReader {
                        slices: (
                            takable_slice_from_maybe_uninitt(tail_to_head),
                            takable_slice_from_maybe_uninitt(empty_slice),
                        ),
                        consumer: self,
                    }
                }
            }
            else
            {
                // Current configuration :
                // ▮▮▮▯▯▯▯▯▯▮▮▮▮
                //    |     |
                // head     tail
                // ___      ____
                // slice2   slice1
                //
                //      Or
                // ▮▮▮▮▮▮▮▮▮▮▮▮▮
                //       |
                //  tail & head
                //    (full)
                // ______.______
                // slice2 slice1

                // SAFETY:
                //
                // Since the StreamWriter object borrows the producer
                // we will return the only, mutually exclusive, mutable refences to the empty part of the buffer
                //
                // We know as per the previous diagram
                // That both slices are mutually exclusive
                // and that tail and head are indices of the slice
                unsafe {
                    let k = &mut *self.inner.buffer.get();

                    let (start_to_tail, tail_to_end) = k.split_at_mut_unchecked(wrapped_tail);
                    let (start_to_head, _head_to_tail) =
                        start_to_tail.split_at_mut_unchecked(wrapped_head);

                    StreamReader {
                        slices: (
                            takable_slice_from_maybe_uninitt(tail_to_end),
                            takable_slice_from_maybe_uninitt(start_to_head),
                        ),
                        consumer: self,
                    }
                }
            }
        }
    }
}

// impl<T: Copy> StreamConsumer<T>
// {
//     pub fn read(&mut self) -> (&[T], &[T])
//     {
//         let (slice_1, slice_2) = self.read_takable();
//         unsafe { (std::mem::transmute(slice_1), std::mem::transmute(slice_2)) }
//     }
// }

mod test
{
    #[allow(unused_imports)]
    use std::mem::MaybeUninit;

    #[allow(unused_imports)]
    use crate::stream::bounded_queue;

    // Test push and pop on single thread
    #[test]
    pub fn stream_simple_push_pop()
    {
        let (mut tx, mut rx) = bounded_queue::<usize>(4);

        {
            let mut writer = tx.write();
            let (a, b) = writer.slices_mut();
            assert_eq!(a.len(), 4);
            assert_eq!(b.len(), 0);

            a[0] = MaybeUninit::new(0);
            a[1] = MaybeUninit::new(1);
            a[2] = MaybeUninit::new(2);
            a[3] = MaybeUninit::new(3);

            tx.produce(4);
        }

        {
            let mut reader = rx.read_takable();
            let (a, b) = reader.slices_mut();
            assert_eq!(a.len(), 4);
            assert_eq!(b.len(), 0);

            unsafe {
                assert_eq!(a[0].take(), 0);
                assert_eq!(a[1].take(), 1);
                assert_eq!(a[2].take(), 2);
                assert_eq!(a[3].take(), 3);
            }

            rx.consume(4);
        }

        // Put stream into weird situation
        {
            let mut writer = tx.write();
            let (a, b) = writer.slices_mut();
            assert_eq!(a.len(), 4);
            assert_eq!(b.len(), 0);

            a[0] = MaybeUninit::new(0);
            a[1] = MaybeUninit::new(1);
            a[2] = MaybeUninit::new(2);

            tx.produce(3);
        }

        {
            let mut reader = rx.read_takable();
            let (a, b) = reader.slices_mut();
            assert_eq!(a.len(), 3);
            assert_eq!(b.len(), 0);

            unsafe 
            {
                assert_eq!(a[0].take(), 0);
                assert_eq!(a[1].take(), 1);
                assert_eq!(a[2].take(), 2);
            }

            rx.consume(1);
        }

        {
            let writer = tx.write();
            let (a, b) = writer.slices();
            assert_eq!(a.len(), 1);
            assert_eq!(b.len(), 1);
        }

        {
            let reader = rx.read_takable();
            let (a, b) = reader.slices();
            assert_eq!(a.len(), 2);
            assert_eq!(b.len(), 0);
        }
    }
}
