//! Ring buffer implementation

use core::{
    future::poll_fn,
    mem::MaybeUninit,
    task::{Context, Poll},
};

use alloc::boxed::Box;
use yggdrasil_abi::error::Error;

use crate::{
    sync::{IrqSafeSpinlock, IrqSafeSpinlockGuard},
    waker::QueueWaker,
};

// TODO use lockless ring data structure

/// Ring buffer base
pub struct RingBuffer<T: Copy> {
    rd: usize,
    wr: usize,
    capacity: usize,
    data: Option<Box<[MaybeUninit<T>]>>,
}

pub struct LossyRingQueue<T: Copy> {
    ring: IrqSafeSpinlock<RingBuffer<T>>,
    read_notify: QueueWaker,
}

pub struct BlockingRingQueue<T: Copy> {
    ring: IrqSafeSpinlock<RingBuffer<T>>,
    write_notify: QueueWaker,
    read_notify: QueueWaker,
}

impl<T: Copy> RingBuffer<T> {
    /// Constructs an empty [RingBuffer].
    ///
    /// NOTE does not allocate until the first write operation. Any read accesses
    ///      without checking readable count will cause a panic.
    pub const fn with_capacity(capacity: usize) -> Self {
        Self {
            rd: 0,
            wr: 0,
            capacity,
            data: None,
        }
    }

    /// Constructs an empty [RingBuffer], trying to allocate its data
    pub fn try_with_capacity(capacity: usize) -> Result<Self, Error> {
        let data = Box::try_new_uninit_slice(capacity).map_err(|_| Error::OutOfMemory)?;
        Ok(Self {
            rd: 0,
            wr: 0,
            capacity,
            data: Some(data),
        })
    }

    pub fn reserve(&mut self) -> Result<(), Error> {
        if self.data.is_none() {
            let buffer =
                Box::try_new_uninit_slice(self.capacity).map_err(|_| Error::OutOfMemory)?;
            self.data = Some(buffer);
        }
        Ok(())
    }

    const fn readable_count_at(&self, at: usize) -> usize {
        if at <= self.wr {
            self.wr - at
        } else {
            self.wr + (self.capacity - at)
        }
    }

    /// Returns `true` if the ring has data that can be read
    #[inline]
    pub const fn is_readable(&self) -> bool {
        self.is_readable_at(self.rd)
    }

    const fn is_readable_at(&self, at: usize) -> bool {
        if at <= self.wr {
            (self.wr - at) > 0
        } else {
            (self.wr + (self.capacity - at)) > 0
        }
    }

    /// Returns `true` if the ring is not full
    #[inline]
    pub fn is_writable(&self) -> bool {
        (self.wr + 1) % self.capacity != self.rd
    }

    /// Returns `true` if the ring contains given element
    pub fn contains(&self, t: &T) -> bool
    where
        T: PartialEq,
    {
        let data = self.data.as_ref().unwrap();
        let count = self.readable_count_at(self.rd);

        for i in 0..count {
            if unsafe { data[(self.rd + i) % self.capacity].assume_init_ref() == t } {
                return true;
            }
        }

        false
    }

    /// Reads a single value from the ring without checking if it's empty.
    ///
    /// # Safety
    ///
    /// The caller must perform the necessary checks to avoid reading beyond the write head.
    #[inline]
    pub unsafe fn read_single_unchecked(&mut self) -> T {
        let data = self.data.as_ref().unwrap();
        let res = data[self.rd].assume_init_read();
        self.rd = (self.rd + 1) % self.capacity;
        res
    }

    /// Reads all entries available from `pos` to the write head.
    ///
    /// # Safety
    ///
    /// The caller must perform the necessary checks to avoid reading beyond the write head.
    pub unsafe fn read_all_static(&mut self, pos: usize, buffer: &mut [T]) -> usize
    where
        T: Copy,
    {
        let data = self.data.as_ref().unwrap();
        let mut pos = (self.rd + pos) % self.capacity;
        let mut off = 0;
        while off < buffer.len() && self.is_readable_at(pos) {
            buffer[off] = data[pos].assume_init();
            pos += 1;
            off += 1;
        }
        off
    }

    /// Writes a single entry to the buffer
    #[inline]
    pub fn write(&mut self, ch: T) {
        let data = self
            .data
            .get_or_insert_with(|| Box::new_uninit_slice(self.capacity));

        data[self.wr].write(ch);
        self.wr = (self.wr + 1) % self.capacity;
    }

    /// Returns [Some] with a value if it could be read, [None] otherwise
    pub fn try_read(&mut self) -> Option<T> {
        if self.is_readable() {
            Some(unsafe { self.read_single_unchecked() })
        } else {
            None
        }
    }
}

impl<T: Copy> LossyRingQueue<T> {
    pub const fn with_capacity(capacity: usize) -> Self {
        Self {
            ring: IrqSafeSpinlock::new(RingBuffer::with_capacity(capacity)),
            read_notify: QueueWaker::new(),
        }
    }

    pub fn try_with_capacity(capacity: usize) -> Result<Self, Error> {
        Ok(Self {
            ring: IrqSafeSpinlock::new(RingBuffer::try_with_capacity(capacity)?),
            read_notify: QueueWaker::new(),
        })
    }

    pub fn write(&self, value: T) {
        self.ring.lock().write(value);
        self.read_notify.wake_all();
    }

    pub fn write_multiple(&self, values: &[T]) {
        let mut lock = self.ring.lock();
        for &value in values {
            lock.write(value);
        }
        self.read_notify.wake_all();
    }

    pub fn try_read(&self) -> Option<T> {
        // SAFETY If lock is acquired, then there's at least one readable element
        self.try_read_lock()
            .map(|mut f| unsafe { f.read_single_unchecked() })
    }

    #[inline]
    pub fn try_read_lock(&self) -> Option<IrqSafeSpinlockGuard<'_, RingBuffer<T>>> {
        let lock = self.ring.lock();
        lock.is_readable().then_some(lock)
    }

    pub async fn read(&self) -> T {
        let mut lock = self.read_lock().await;
        // SAFETY If lock is acquired, then there's at least one readable element
        unsafe { lock.read_single_unchecked() }
    }

    pub async fn read_lock(&self) -> IrqSafeSpinlockGuard<'_, RingBuffer<T>> {
        poll_fn(|cx| self.poll_lock(cx)).await
    }

    #[inline]
    pub fn poll_lock(&self, cx: &mut Context<'_>) -> Poll<IrqSafeSpinlockGuard<'_, RingBuffer<T>>> {
        self.read_notify.register(cx.waker());
        if let Some(lock) = self.try_read_lock() {
            self.read_notify.remove(cx.waker());
            Poll::Ready(lock)
        } else {
            Poll::Pending
        }
    }

    pub fn poll_readable(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.poll_lock(cx).map(|_| ())
    }

    pub fn notify_all(&self) {
        self.read_notify.wake_all();
    }
}

impl<T: Copy> BlockingRingQueue<T> {
    pub const fn with_capacity(capacity: usize) -> Self {
        Self {
            ring: IrqSafeSpinlock::new(RingBuffer::with_capacity(capacity)),
            write_notify: QueueWaker::new(),
            read_notify: QueueWaker::new(),
        }
    }

    pub fn try_with_capacity(capacity: usize) -> Result<Self, Error> {
        Ok(Self {
            ring: IrqSafeSpinlock::new(RingBuffer::try_with_capacity(capacity)?),
            write_notify: QueueWaker::new(),
            read_notify: QueueWaker::new(),
        })
    }

    pub async fn write(&self, data: T) {
        poll_fn(|cx| {
            if self.try_write(data) {
                self.notify_readers();
                self.write_notify.remove(cx.waker());
                Poll::Ready(())
            } else {
                self.write_notify.register(cx.waker());
                Poll::Pending
            }
        })
        .await
    }

    pub async fn write_all(&self, data: &[T]) {
        for &ch in data {
            self.write(ch).await;
        }
    }

    pub fn try_write(&self, data: T) -> bool {
        let mut lock = self.ring.lock();
        if !lock.is_writable() {
            return false;
        }
        lock.write(data);
        true
    }

    pub async fn read_lock(&self) -> IrqSafeSpinlockGuard<'_, RingBuffer<T>> {
        poll_fn(|cx| self.poll_read_lock(cx)).await
    }

    #[inline]
    pub fn poll_read_lock(
        &self,
        cx: &mut Context<'_>,
    ) -> Poll<IrqSafeSpinlockGuard<'_, RingBuffer<T>>> {
        if let Some(lock) = self.try_read_lock() {
            self.read_notify.remove(cx.waker());
            Poll::Ready(lock)
        } else {
            self.read_notify.register(cx.waker());
            Poll::Pending
        }
    }

    #[inline]
    pub fn try_read_lock(&self) -> Option<IrqSafeSpinlockGuard<'_, RingBuffer<T>>> {
        let lock = self.ring.lock();
        lock.is_readable().then_some(lock)
    }

    pub fn poll_readable(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.poll_read_lock(cx).map(|_| ())
    }

    pub fn notify_readers(&self) {
        self.read_notify.wake_all();
    }

    pub fn notify_writers(&self) {
        self.write_notify.wake_all();
    }
}
