use crate::mmio::{MmioAccess, REG_QUEUE_NOTIFY};
use core::sync::atomic::{Ordering, fence};
use libkos::Window;
use types::address::ByteOffset;

pub const QUEUE_SIZE: usize = 8;
const DESC_SIZE: usize = 16;
const TIMEOUT: u32 = 10_000_000;
/// Where the legacy layout puts the available and used rings.
const AVAIL: ByteOffset = ByteOffset::new(QUEUE_SIZE * DESC_SIZE);
const USED: ByteOffset = ByteOffset::new(4096);

pub struct Descriptor {
    pub addr: u64,
    pub len: u32,
    pub writable: bool,
}

/// One virtqueue in the legacy layout, over the two pages its rings sit in.
pub struct VirtioQueue {
    rings: Window,
}

impl VirtioQueue {
    pub fn new(rings: Window) -> Self {
        Self { rings }
    }

    /// Hands the device a chain of `descs` and waits until it is used.
    pub fn submit(&self, mmio: &impl MmioAccess, descs: &[Descriptor]) -> bool {
        self.submit_written(mmio, descs).is_some()
    }

    /// Like `submit`; how many bytes the device wrote into the chain's
    /// writable buffers (it may write fewer than they hold).
    pub fn submit_written(&self, mmio: &impl MmioAccess, descs: &[Descriptor]) -> Option<u32> {
        if descs.is_empty() || descs.len() > QUEUE_SIZE {
            return None;
        }
        let r = &self.rings;
        for (i, d) in descs.iter().enumerate() {
            let at = ByteOffset::new(i * DESC_SIZE);
            let last = i == descs.len() - 1;
            let flags: u16 = if d.writable { 2 } else { 0 } | if !last { 1 } else { 0 };
            r.write(at, d.addr);
            r.write(at + 8, d.len);
            r.write(at + 12, flags);
            r.write(at + 14, if last { 0 } else { (i + 1) as u16 });
        }

        let avail_idx: u16 = r.read(AVAIL + 2);
        r.write(AVAIL + 4 + (avail_idx as usize % QUEUE_SIZE) * 2, 0u16);
        fence(Ordering::SeqCst);
        r.write(AVAIL + 2, avail_idx.wrapping_add(1));
        fence(Ordering::SeqCst);

        mmio.write(REG_QUEUE_NOTIFY, 0);

        let expected = avail_idx.wrapping_add(1);
        let mut timeout = TIMEOUT;
        while r.read::<u16>(USED + 2) != expected {
            core::hint::spin_loop();
            timeout -= 1;
            if timeout == 0 {
                return None;
            }
        }
        // The used element just added: the chain's head id, then its length.
        let element = USED + 4 + (avail_idx as usize % QUEUE_SIZE) * 8;
        Some(r.read::<u32>(element + 4))
    }
}