//! A *split virtqueue*: three arrays in guest RAM shared with the device.
//!
//! ```text
//!  descriptor table   [addr u64 | len u32 | flags u16 | next u16] x size
//!  available ring     flags u16, idx u16, ring[size] u16   (driver -> device)
//!  used ring          flags u16, idx u16, ring[size] {id u32, len u32}
//!                                                           (device -> driver)
//! ```
//!
//! To send a request the driver fills one or more descriptors (chained with
//! the NEXT flag; WRITE marks buffers the *device* writes into, like a
//! response), puts the first one's index in the available ring, bumps
//! `avail.idx` and rings the doorbell. When done, the device puts that index
//! (and how many bytes it wrote) into the used ring and bumps `used.idx`.
//! Both indexes only ever count up (wrapping at 65536); position in the ring
//! is `idx % size`. That's the whole protocol.

use crate::memory;
use alloc::vec::Vec;
use core::ptr::{read_volatile, write_volatile};
use core::sync::atomic::{Ordering, fence};
use x86_64::PhysAddr;

const DESC_NEXT: u16 = 1; // the chain continues at `next`
const DESC_WRITE: u16 = 2; // the device writes this buffer
const AVAIL_NO_INTERRUPT: u16 = 1; // "don't interrupt me when you use buffers"

#[repr(C)]
struct Descriptor {
    addr: u64,
    len: u32,
    flags: u16,
    next: u16,
}

pub struct Virtqueue {
    index: u16,
    size: u16,
    phys: u64,       // start of the descriptor table (all three parts are contiguous)
    avail_off: u64,  // byte offsets of the available and used rings
    used_off: u64,
    doorbell: u64,   // where to write `index` to notify the device
    free: Vec<u16>,  // descriptors not in use
    avail_idx: u16,  // our copy of avail.idx
    last_used: u16,  // used ring entries we've already processed
}

// The raw addresses are only touched under the driver's lock.
unsafe impl Send for Virtqueue {}

impl Virtqueue {
    /// Allocate and zero the rings: descriptors and the available ring
    /// first, then the used ring on its own page-aligned spot.
    pub fn new(index: u16, size: u16, doorbell: u64, interrupts: bool) -> Option<Virtqueue> {
        let n = size as u64;
        let avail_off = 16 * n;
        let used_off = (avail_off + 6 + 2 * n).next_multiple_of(4096);
        let bytes = used_off + (6 + 8 * n).next_multiple_of(4096);
        let phys = memory::alloc_contiguous_frames((bytes / 4096) as usize)?.as_u64();
        let queue = Virtqueue {
            index, size, phys, avail_off, used_off, doorbell,
            free: (0..size).rev().collect(),
            avail_idx: 0,
            last_used: 0,
        };
        if !interrupts {
            unsafe { write_volatile(queue.ptr::<u16>(avail_off), AVAIL_NO_INTERRUPT) };
        }
        Some(queue)
    }

    pub fn desc_phys(&self) -> u64 { self.phys }
    pub fn avail_phys(&self) -> u64 { self.phys + self.avail_off }
    pub fn used_phys(&self) -> u64 { self.phys + self.used_off }
    pub fn free_descriptors(&self) -> usize { self.free.len() }
    /// The descriptor the next one-buffer `push` will use (drivers that
    /// keep one data slot per descriptor want to know it beforehand).
    pub fn next_free(&self) -> Option<u16> { self.free.last().copied() }

    fn ptr<T>(&self, offset: u64) -> *mut T {
        memory::phys_to_virt(PhysAddr::new(self.phys + offset)).as_mut_ptr()
    }

    /// Queue one request made of `buffers` (physical address, length,
    /// device-writable?). Returns the head descriptor's index (the request's
    /// id), or None if the queue is full. Doesn't notify: call `notify`.
    pub fn push(&mut self, buffers: &[(u64, u32, bool)]) -> Option<u16> {
        // (A fixed array, not a Vec: the console calls this for every
        // kernel log line, so keep it cheap.)
        if buffers.is_empty() || buffers.len() > 4 || self.free.len() < buffers.len() {
            return None;
        }
        let mut ids = [0u16; 4];
        for id in ids.iter_mut().take(buffers.len()) {
            *id = self.free.pop().unwrap();
        }
        let ids = &ids[..buffers.len()];
        for (i, &(addr, len, device_writes)) in buffers.iter().enumerate() {
            let mut flags = if device_writes { DESC_WRITE } else { 0 };
            let next = if i + 1 < ids.len() { flags |= DESC_NEXT; ids[i + 1] } else { 0 };
            let desc = Descriptor { addr, len, flags, next };
            unsafe { write_volatile(self.ptr::<Descriptor>(16 * ids[i] as u64), desc) };
        }
        let slot = (self.avail_idx % self.size) as u64;
        unsafe { write_volatile(self.ptr::<u16>(self.avail_off + 4 + 2 * slot), ids[0]) };
        // The descriptors and the ring entry must be visible before the
        // device can see the new index.
        fence(Ordering::SeqCst);
        self.avail_idx = self.avail_idx.wrapping_add(1);
        unsafe { write_volatile(self.ptr::<u16>(self.avail_off + 2), self.avail_idx) };
        fence(Ordering::SeqCst);
        Some(ids[0])
    }

    /// Ring the doorbell: "queue `index` has new requests".
    pub fn notify(&self) {
        unsafe { write_volatile(self.doorbell as *mut u16, self.index) };
    }

    /// Take one finished request off the used ring: (head id, bytes the
    /// device wrote). Its descriptors become free again.
    pub fn pop_used(&mut self) -> Option<(u16, u32)> {
        fence(Ordering::SeqCst);
        let used_idx = unsafe { read_volatile(self.ptr::<u16>(self.used_off + 2)) };
        if used_idx == self.last_used {
            return None;
        }
        let slot = (self.last_used % self.size) as u64;
        let entry = self.used_off + 4 + 8 * slot;
        let (id, len) = unsafe { (read_volatile(self.ptr::<u32>(entry)), read_volatile(self.ptr::<u32>(entry + 4))) };
        self.last_used = self.last_used.wrapping_add(1);
        // Give the chain back to the free list.
        let mut d = id as u16;
        loop {
            self.free.push(d);
            let desc = unsafe { read_volatile(self.ptr::<Descriptor>(16 * d as u64)) };
            if desc.flags & DESC_NEXT == 0 {
                break;
            }
            d = desc.next;
        }
        Some((id as u16, len))
    }

    /// Wait (spinning) for the device to finish one request, at most
    /// `timeout_us` microseconds. For the polled devices (rng, gpu).
    pub fn wait_used(&mut self, timeout_us: u64) -> Option<(u16, u32)> {
        let deadline = crate::time::now_us() + timeout_us;
        // Also count spins, in case the clock isn't running (interrupts off
        // before the TSC is calibrated).
        for spins in 0u64.. {
            if let Some(done) = self.pop_used() {
                return Some(done);
            }
            if crate::time::now_us() > deadline || spins > 200_000_000 {
                return None;
            }
            core::hint::spin_loop();
        }
        None
    }
}