//! VirtIO queue implementation.

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

use alloc::{boxed::Box, sync::Arc};
use device_api::dma::DmaAllocator;
use libk::{
    dma::{BusAddress, DmaBuffer, DmaSlice, DmaSliceMut},
    error::Error,
};
use libk_util::{
    event::OneTimeEvent, hash_table::DefaultHashTable, sync::IrqSafeSpinlock, waker::QueueWaker,
};

#[derive(Clone, Copy, Debug)]
#[repr(C)]
struct Descriptor {
    address: BusAddress,
    length: u32,
    flags: u16,
    next: u16,
}

struct AvailableRing {
    mapping: DmaBuffer<[u16]>,
    capacity: usize,
}

struct UsedRing {
    mapping: DmaBuffer<[u32]>,
    last_seen_used: u16,
    capacity: usize,
}

pub struct DescriptorTable {
    descriptors: DmaBuffer<[Descriptor]>,
    free_count: usize,
    first_free: Option<u16>,
    last_free: Option<u16>,
}

struct VqInner {
    descriptors: DescriptorTable,
    available: AvailableRing,
    used: UsedRing,
}

pub struct VirtQueue<N: VqNotificationMechanism = VqAsyncNotification> {
    inner: IrqSafeSpinlock<VqInner>,
    free_descriptor_notify: QueueWaker,

    #[allow(unused)]
    index: u16,
    capacity: usize,

    used_notify: N,
}

pub trait VqNotificationMechanism {
    type Token;

    fn notify_used(&self, head: u16, length: u32);
    fn create_token(&self, head: u16) -> Self::Token;
}

pub struct VqAsyncNotification {
    completions: IrqSafeSpinlock<DefaultHashTable<u16, Arc<OneTimeEvent<u32>>>>,
}

pub struct VqManualNotification;

pub struct VqCallbackNotification(Box<dyn Fn(u16, u32) + Sync + Send>);

impl VqNotificationMechanism for VqAsyncNotification {
    type Token = Arc<OneTimeEvent<u32>>;

    fn notify_used(&self, head: u16, length: u32) {
        let mut completions = self.completions.lock();
        if let Some(completion) = completions.remove(&head) {
            log::trace!("vq: completion #{head}");
            completion.signal(length);
        }
    }

    fn create_token(&self, head: u16) -> Self::Token {
        let mut completions = self.completions.lock();
        let token = Arc::new(OneTimeEvent::new());
        completions.insert(head, token.clone());
        token
    }
}

impl VqNotificationMechanism for VqCallbackNotification {
    type Token = u16;

    fn create_token(&self, head: u16) -> Self::Token {
        head
    }

    fn notify_used(&self, head: u16, length: u32) {
        (self.0)(head, length);
    }
}

impl VqNotificationMechanism for VqManualNotification {
    type Token = u16;

    fn create_token(&self, head: u16) -> Self::Token {
        head
    }

    fn notify_used(&self, _head: u16, _length: u32) {
        unreachable!()
    }
}

impl AvailableRing {
    const FLAGS: usize = 0;
    const IDX: usize = 1;
    const RING: usize = 2;

    pub fn with_capacity(
        dma: &dyn DmaAllocator,
        capacity: usize,
        no_interrupt: bool,
    ) -> Result<Self, Error> {
        // flags + idx + [ring] + used_event
        let mut mapping = DmaBuffer::new_slice(dma, 0u16, (capacity + 6) & !3)?;
        if no_interrupt {
            mapping[Self::FLAGS] |= 1 << 0;
        }
        Ok(Self { mapping, capacity })
    }

    pub fn push(&mut self, head: u16) -> u16 {
        log::trace!("enqueue #{head}");
        let idx = self.mapping[Self::IDX];
        let index = idx as usize % self.capacity;
        self.mapping[Self::RING + index] = head;
        self.mapping.cache_flush_element(Self::RING + index, true);
        let idx = idx.wrapping_add(1);
        self.mapping[Self::IDX] = idx;
        self.mapping.cache_flush_element(Self::IDX, true);
        idx
    }
}

impl UsedRing {
    const FLAGS_IDX: usize = 0;
    const RING: usize = 1;

    pub fn with_capacity(
        dma: &dyn DmaAllocator,
        capacity: usize,
        no_notify: bool,
    ) -> Result<Self, Error> {
        // 2x u16 (flags + idx) + [ring x 2 x u32] + avail_event
        let mut mapping = DmaBuffer::new_slice(dma, 0, capacity * 2 + 1)?;
        if no_notify {
            mapping[Self::FLAGS_IDX] |= 1 << 0;
        }
        mapping.cache_flush_element(Self::FLAGS_IDX, true);
        Ok(Self {
            mapping,
            capacity,
            last_seen_used: 0,
        })
    }

    pub fn consume<F: FnMut(u16, u32)>(&mut self, mut handler: F) -> usize {
        self.mapping.cache_flush_element(Self::FLAGS_IDX, false);
        let idx = (self.mapping[Self::FLAGS_IDX] >> 16) as u16;
        let mut count = 0;
        while self.last_seen_used != idx {
            let index = self.last_seen_used as usize % self.capacity;
            self.mapping
                .cache_flush_range(Self::RING + index..Self::RING + index + 1, false);
            let head = self.mapping[Self::RING + index * 2] as u16;
            let len = self.mapping[Self::RING + index * 2 + 1];
            handler(head, len);
            count += 1;
            self.last_seen_used = self.last_seen_used.wrapping_add(1);
        }
        count
    }
}

impl Descriptor {
    pub const EMPTY: Self = Self {
        address: BusAddress::ZERO,
        length: 0,
        flags: 0,
        next: 0,
    };

    pub const F_NEXT: u16 = 1 << 0;
    pub const F_WRITE: u16 = 1 << 1;
}

impl DescriptorTable {
    const FREE_CHAIN_END: u16 = u16::MAX;

    pub fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
        if capacity >= Self::FREE_CHAIN_END as usize - 1 {
            return Err(Error::InvalidArgument);
        }
        let mut descriptors = DmaBuffer::new_slice(dma, Descriptor::EMPTY, capacity)?;
        for i in 0..capacity {
            if i == capacity - 1 {
                // Last descriptor of the free chain
                descriptors[i].next = Self::FREE_CHAIN_END;
            } else {
                descriptors[i].next = (i + 1) as u16;
            }
        }

        Ok(Self {
            descriptors,
            free_count: capacity,
            first_free: Some(0),
            last_free: Some(capacity as u16 - 1),
        })
    }

    // Allocate a chain of descriptors
    fn alloc_descriptors(
        &mut self,
        h2d: &[DmaSlice<u8>],
        d2h: &[DmaSliceMut<MaybeUninit<u8>>],
    ) -> Result<u16, Error> {
        if d2h.len() + h2d.len() == 0 {
            // Empty transfer
            return Err(Error::InvalidArgument);
        }
        if self.free_count < d2h.len() + h2d.len() {
            // Not enough descriptor "slots" to place the buffers
            return Err(Error::WouldBlock);
        }
        // Implied by free_count
        debug_assert!(self.first_free.is_some());
        debug_assert!(self.last_free.is_some());
        let head = unsafe {
            self.follow_descriptor_chain(h2d.len() + d2h.len(), |i, desc| {
                if i < h2d.len() {
                    desc.address = h2d[i].bus_address();
                    desc.length = h2d[i].len() as u32;
                    desc.flags = 0;
                } else {
                    let i = i - h2d.len();
                    desc.address = d2h[i].bus_address();
                    desc.length = d2h[i].len() as u32;
                    desc.flags = Descriptor::F_WRITE;
                }
            })
        };
        Ok(head)
    }

    /// # Safety
    ///
    /// The following invariants must hold:
    ///
    /// * **There must actually be `count` free descriptors**.
    /// * **`count` should not be zero**.
    /// * The VQ must be in a consistent state: first_free/last_free/free_count must be consistent
    ///   in terms of tracking the current state of the descriptor allocation.
    /// * Free descriptors in the table must be chained.
    /// * Last free descriptor must have [Self::FREE_CHAIN_END] as its `next` field.
    unsafe fn follow_descriptor_chain<F: Fn(usize, &mut Descriptor)>(
        &mut self,
        count: usize,
        visitor: F,
    ) -> u16 {
        debug_assert_ne!(count, 0);
        let mut current = self.first_free.unwrap_unchecked();
        let head = current;
        for i in 0..count {
            let descriptor = &mut self.descriptors[current as usize];
            log::trace!("vq: alloc desc #{current}");
            visitor(i, descriptor);
            if i == count - 1 {
                debug_assert_eq!(descriptor.flags & Descriptor::F_NEXT, 0);
                current = descriptor.next;
                descriptor.next = 0;
            } else {
                current = descriptor.next;
                descriptor.flags |= Descriptor::F_NEXT;
                debug_assert_ne!(current, Self::FREE_CHAIN_END);
            };
        }
        if current == Self::FREE_CHAIN_END {
            // No free descriptors left
            debug_assert_eq!(self.free_count, count);
            self.first_free = None;
            self.last_free = None;
        } else {
            self.first_free = Some(current);
        }
        self.free_count -= count;

        head
    }

    fn add_free_descriptor(&mut self, idx: u16) {
        log::trace!("vq: free descriptor #{idx}");
        self.descriptors[idx as usize] = Descriptor {
            next: Self::FREE_CHAIN_END,
            ..Descriptor::EMPTY
        };

        if let Some(last) = self.last_free {
            // Implies first free is Some as well
            self.descriptors[last as usize].next = idx;
            // last -> idx
            self.last_free = Some(idx);
        } else {
            // Implies first free is None as well
            self.first_free = Some(idx);
            self.last_free = Some(idx);
        }

        self.free_count += 1;
    }

    fn free_descriptor_chain(&mut self, first: u16) -> usize {
        let mut current = first;
        let mut count = 0;
        loop {
            let descriptor = &self.descriptors[current as usize];
            let next = if descriptor.flags & Descriptor::F_NEXT != 0 {
                Some(descriptor.next)
            } else {
                None
            };

            count += 1;
            self.add_free_descriptor(current);

            if let Some(next) = next {
                current = next;
            } else {
                break;
            }
        }
        count
    }
}

impl VqInner {
    fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
        let descriptors = DescriptorTable::with_capacity(dma, capacity)?;
        let available = AvailableRing::with_capacity(dma, capacity, false)?;
        let used = UsedRing::with_capacity(dma, capacity, false)?;

        Ok(Self {
            descriptors,
            available,
            used,
        })
    }

    fn try_enqueue(
        &mut self,
        h2d: &[DmaSlice<u8>],
        d2h: &[DmaSliceMut<MaybeUninit<u8>>],
    ) -> Result<u16, Error> {
        let head = self.descriptors.alloc_descriptors(h2d, d2h)?;
        self.available.push(head);
        Ok(head)
    }

    fn consume<F: FnMut(u16, u32)>(&mut self, free_notify: &QueueWaker, mut handler: F) -> usize {
        self.used.consume(|head, len| {
            log::trace!("vq: used #{head}, len={len}");

            self.descriptors.free_descriptor_chain(head);
            free_notify.wake_all();

            handler(head, len);
        })
    }
}

impl VirtQueue<VqAsyncNotification> {
    pub fn with_capacity(
        dma: &dyn DmaAllocator,
        index: u16,
        capacity: usize,
    ) -> Result<Self, Error> {
        let used_notify = VqAsyncNotification::new();
        Self::with_capacity_and_notify(dma, index, capacity, used_notify)
    }

    pub async fn enqueue_wait<F: FnOnce()>(
        &self,
        h2d: &[DmaSlice<'_, u8>],
        d2h: &[DmaSliceMut<'_, MaybeUninit<u8>>],
        notify_queue: F,
    ) -> Result<u32, Error> {
        let completion = self.enqueue(h2d, d2h).await?;
        notify_queue();
        let result = completion.wait_copy().await;
        Ok(result)
    }
}

impl VirtQueue<VqCallbackNotification> {
    pub fn with_capacity_and_callback(
        dma: &dyn DmaAllocator,
        index: u16,
        capacity: usize,
        used_callback: Box<dyn Fn(u16, u32) + Sync + Send>,
    ) -> Result<Self, Error> {
        let used_notify = VqCallbackNotification(used_callback);
        Self::with_capacity_and_notify(dma, index, capacity, used_notify)
    }
}

impl VirtQueue<VqManualNotification> {
    pub fn with_capacity_manual(
        dma: &dyn DmaAllocator,
        index: u16,
        capacity: usize,
    ) -> Result<Self, Error> {
        let used_notify = VqManualNotification;
        Self::with_capacity_and_notify(dma, index, capacity, used_notify)
    }

    pub fn handle_notify_manual<F: FnMut(u16, u32)>(&self, handler: F) -> usize {
        self.inner
            .lock()
            .consume(&self.free_descriptor_notify, handler)
    }

    // Used when queue does not support device-side used notification
    pub fn enqueue_blocking<F: FnOnce()>(
        &self,
        h2d: &[DmaSlice<u8>],
        d2h: &[DmaSliceMut<MaybeUninit<u8>>],
        notify_queue: F,
    ) -> Result<u32, Error> {
        let token = self.try_enqueue(h2d, d2h)?;
        let mut length = 0;
        notify_queue();
        loop {
            self.handle_notify_manual(|head, len| {
                assert_eq!(head, token);
                length = len;
            });
            if length != 0 {
                break;
            }
            core::hint::spin_loop();
        }
        Ok(length)
    }
}

impl<N: VqNotificationMechanism> VirtQueue<N> {
    pub fn with_capacity_and_notify(
        dma: &dyn DmaAllocator,
        index: u16,
        capacity: usize,
        used_notify: N,
    ) -> Result<Self, Error> {
        let inner = VqInner::with_capacity(dma, capacity)?;
        Ok(Self {
            inner: IrqSafeSpinlock::new(inner),

            index,
            capacity,
            used_notify,
            free_descriptor_notify: QueueWaker::new(),
        })
    }

    pub fn descriptor_table_base(&self) -> BusAddress {
        self.inner.lock().descriptors.descriptors.bus_address()
    }

    pub fn available_ring_base(&self) -> BusAddress {
        self.inner.lock().available.mapping.bus_address()
    }

    pub fn used_ring_base(&self) -> BusAddress {
        self.inner.lock().used.mapping.bus_address()
    }

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

    pub fn handle_notify(&self) -> usize {
        self.inner
            .lock()
            .consume(&self.free_descriptor_notify, |head, len| {
                self.used_notify.notify_used(head, len);
            })
    }

    pub fn try_enqueue(
        &self,
        h2d: &[DmaSlice<'_, u8>],
        d2h: &[DmaSliceMut<'_, MaybeUninit<u8>>],
    ) -> Result<N::Token, Error> {
        let head = self.inner.lock().try_enqueue(h2d, d2h)?;
        Ok(self.used_notify.create_token(head))
    }

    pub async fn enqueue(
        &self,
        h2d: &[DmaSlice<'_, u8>],
        d2h: &[DmaSliceMut<'_, MaybeUninit<u8>>],
    ) -> Result<N::Token, Error> {
        poll_fn(|cx| match self.try_enqueue(h2d, d2h) {
            Err(Error::WouldBlock) => {
                self.free_descriptor_notify.register(cx.waker());
                Poll::Pending
            }
            result => {
                self.free_descriptor_notify.remove(cx.waker());
                Poll::Ready(result)
            }
        })
        .await
    }
}

impl VqAsyncNotification {
    pub fn new() -> Self {
        Self {
            completions: IrqSafeSpinlock::new(DefaultHashTable::new()),
        }
    }
}
