//! Driver for virtio-net: QEMU's paravirtual network card
//! (`-device virtio-net-pci`).
//!
//! It works like virtio-blk (block/virtio_blk.rs: same legacy I/O-port
//! registers, same virtqueues) but with *two* queues:
//!   * queue 0, receive: we hand the card empty buffers in advance; when a
//!     frame arrives it fills one, puts it in the used ring and (because we
//!     asked for it) raises an interrupt.
//!   * queue 1, transmit: we fill a buffer with a frame and hand it over.
//! Every buffer starts with a 10-byte `virtio_net_hdr` (checksum offload,
//! segmentation offload...). We negotiate none of those features, so on
//! send it's all zeros and on receive we skip it. After the header comes a
//! plain Ethernet frame: destination MAC, source MAC, type, payload.
//!
//! The only feature we accept is MAC (bit 5): "the device config area holds
//! a MAC address", which is how we learn ours (QEMU: 52:54:00:12:34:56).
//!
//! This file knows nothing about IP: it just moves Ethernet frames. The
//! `smoltcp::phy::Device` impl at the bottom is the glue that lets smoltcp
//! (net/mod.rs) send and receive through it.
//! Reference: virtio spec 1.2 section 5.1 (network device).

use crate::memory;
use crate::pci::PciDevice;
use alloc::vec::Vec;
use core::ptr::{read_volatile, write_volatile};
use core::sync::atomic::{Ordering, fence};
use smoltcp::phy::{self, DeviceCapabilities, Medium};
use smoltcp::time::Instant;
use x86_64::PhysAddr;
use x86_64::instructions::port::Port;

// Legacy virtio PCI registers (offsets into the I/O BAR), as in virtio_blk.
const REG_DEVICE_FEATURES: u16 = 0x00;
const REG_GUEST_FEATURES: u16 = 0x04;
const REG_QUEUE_ADDRESS: u16 = 0x08;
const REG_QUEUE_SIZE: u16 = 0x0c;
const REG_QUEUE_SELECT: u16 = 0x0e;
const REG_QUEUE_NOTIFY: u16 = 0x10;
const REG_DEVICE_STATUS: u16 = 0x12;
pub const REG_ISR_STATUS: u16 = 0x13; // reading it acknowledges the interrupt
const REG_MAC: u16 = 0x14; // device config: 6-byte MAC address

const STATUS_ACKNOWLEDGE: u8 = 1;
const STATUS_DRIVER: u8 = 2;
const STATUS_DRIVER_OK: u8 = 4;
const FEATURE_MAC: u32 = 1 << 5;

const DESC_WRITE: u16 = 2; // device writes into this buffer
const AVAIL_NO_INTERRUPT: u16 = 1;

const QUEUE_RX: u16 = 0;
const QUEUE_TX: u16 = 1;
/// Buffers per queue, and their size: a 10-byte header + a max-size
/// Ethernet frame (1514 bytes) fits in 2 KiB, so 2 buffers per page.
const BUFFERS: usize = 32;
const BUF_SIZE: usize = 2048;
const HEADER_LEN: usize = 10;
/// Largest Ethernet frame without the 4-byte checksum (the card adds it).
pub const MTU: usize = 1514;

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

/// PCI vendor 0x1af4 (virtio), device 0x1000 = transitional network card.
pub fn is_virtio_net(dev: &PciDevice) -> bool {
    dev.vendor_id == 0x1af4 && dev.device_id == 0x1000
}

/// One virtqueue plus its buffers. Descriptor `i` always points at
/// buffer `i`, which keeps the bookkeeping trivial.
struct Virtqueue {
    size: u16,
    phys: PhysAddr, // descriptors + available ring, then (page aligned) used ring
    used_offset: usize,
    last_used: u16, // how far we've read the used ring
    buffers: PhysAddr, // BUFFERS * BUF_SIZE bytes
}

impl Virtqueue {
    fn new(io: u16, index: u16) -> Option<Virtqueue> {
        unsafe { Port::<u16>::new(io + REG_QUEUE_SELECT).write(index) };
        let size = unsafe { Port::<u16>::new(io + REG_QUEUE_SIZE).read() };
        if (size as usize) < BUFFERS {
            return None;
        }
        let n = size as usize;
        let used_offset = align_up(16 * n + 6 + 2 * n, 4096);
        let total = used_offset + align_up(6 + 8 * n, 4096);
        let phys = memory::alloc_contiguous_frames(total / 4096)?;
        let buffers = memory::alloc_contiguous_frames(BUFFERS * BUF_SIZE / 4096)?;
        unsafe {
            core::ptr::write_bytes(memory::phys_to_virt(phys).as_mut_ptr::<u8>(), 0, total);
            Port::<u32>::new(io + REG_QUEUE_ADDRESS).write((phys.as_u64() / 4096) as u32);
        }
        Some(Virtqueue { size, phys, used_offset, last_used: 0, buffers })
    }

    fn base(&self) -> *mut u8 {
        memory::phys_to_virt(self.phys).as_mut_ptr()
    }

    /// Our view of buffer `i`.
    fn buffer(&mut self, i: u16) -> &mut [u8] {
        let virt = memory::phys_to_virt(self.buffers + i as u64 * BUF_SIZE as u64);
        unsafe { core::slice::from_raw_parts_mut(virt.as_mut_ptr(), BUF_SIZE) }
    }

    /// Hand buffer `i` (its first `len` bytes) to the device.
    fn push(&mut self, i: u16, len: usize, flags: u16) {
        let n = self.size as usize;
        unsafe {
            let desc = self.base() as *mut Descriptor;
            let addr = (self.buffers + i as u64 * BUF_SIZE as u64).as_u64();
            write_volatile(desc.add(i as usize), Descriptor { addr, len: len as u32, flags, next: 0 });
            // Available ring: flags (u16), idx (u16), ring[n] (u16 each).
            let avail = self.base().add(16 * n) as *mut u16;
            let idx = read_volatile(avail.add(1));
            write_volatile(avail.add(2 + (idx as usize % n)), i);
            fence(Ordering::SeqCst); // the entry must be visible before the index moves
            write_volatile(avail.add(1), idx.wrapping_add(1));
        }
    }

    /// The next buffer the device has finished with: (buffer index, bytes
    /// it wrote).
    fn pop_used(&mut self) -> Option<(u16, usize)> {
        let n = self.size as usize;
        unsafe {
            // Used ring: flags (u16), idx (u16), ring[n] of {id: u32, len: u32}.
            let used = self.base().add(self.used_offset);
            let idx = read_volatile(used.add(2) as *const u16);
            if idx == self.last_used {
                return None;
            }
            fence(Ordering::SeqCst);
            let entry = used.add(4 + 8 * (self.last_used as usize % n)) as *const u32;
            let (id, len) = (read_volatile(entry), read_volatile(entry.add(1)));
            self.last_used = self.last_used.wrapping_add(1);
            Some((id as u16, len as usize))
        }
    }

    fn set_no_interrupt(&mut self) {
        unsafe { write_volatile(self.base().add(16 * self.size as usize) as *mut u16, AVAIL_NO_INTERRUPT) };
    }
}

/// Packet and byte counters, shown by `ifconfig`.
#[derive(Default, Clone, Copy)]
pub struct Stats {
    pub rx_packets: u64,
    pub rx_bytes: u64,
    pub tx_packets: u64,
    pub tx_bytes: u64,
}

pub struct VirtioNet {
    pub io_base: u16,
    pub mac: [u8; 6],
    pub irq: u8,
    pub stats: Stats,
    rx: Virtqueue,
    tx: Virtqueue,
    tx_free: Vec<u16>, // transmit buffers the device has given back
}

impl VirtioNet {
    pub fn new(pci: &PciDevice) -> Option<VirtioNet> {
        let bar0 = pci.bar(0);
        if bar0 & 1 == 0 {
            return None;
        }
        pci.enable_io_and_bus_master();
        let io = (bar0 & !0x3) as u16;
        let status = |v: u8| unsafe { Port::<u8>::new(io + REG_DEVICE_STATUS).write(v) };
        // Same dance as virtio-blk: reset, acknowledge, negotiate, queues, go.
        status(0);
        status(STATUS_ACKNOWLEDGE);
        status(STATUS_ACKNOWLEDGE | STATUS_DRIVER);
        let offered = unsafe { Port::<u32>::new(io + REG_DEVICE_FEATURES).read() };
        unsafe { Port::<u32>::new(io + REG_GUEST_FEATURES).write(offered & FEATURE_MAC) };
        let mut mac = [0u8; 6];
        for (i, b) in mac.iter_mut().enumerate() {
            *b = unsafe { Port::<u8>::new(io + REG_MAC + i as u16).read() };
        }
        let rx = Virtqueue::new(io, QUEUE_RX)?;
        let mut tx = Virtqueue::new(io, QUEUE_TX)?;
        tx.set_no_interrupt(); // we notice finished sends when we next send
        let mut dev = VirtioNet {
            io_base: io,
            mac,
            irq: pci.interrupt_line(),
            stats: Stats::default(),
            rx,
            tx,
            tx_free: (0..BUFFERS as u16).collect(),
        };
        // Give the device every receive buffer up front.
        for i in 0..BUFFERS as u16 {
            dev.rx.push(i, BUF_SIZE, DESC_WRITE);
        }
        status(STATUS_ACKNOWLEDGE | STATUS_DRIVER | STATUS_DRIVER_OK);
        dev.notify(QUEUE_RX);
        Some(dev)
    }

    fn notify(&self, queue: u16) {
        unsafe { Port::<u16>::new(self.io_base + REG_QUEUE_NOTIFY).write(queue) };
    }

    /// Take one received frame (copied out; its buffer goes straight back
    /// to the device).
    fn receive_frame(&mut self) -> Option<Vec<u8>> {
        let (i, len) = self.rx.pop_used()?;
        let len = len.clamp(HEADER_LEN, BUF_SIZE);
        let frame = self.rx.buffer(i)[HEADER_LEN..len].to_vec();
        self.rx.push(i, BUF_SIZE, DESC_WRITE);
        self.notify(QUEUE_RX);
        self.stats.rx_packets += 1;
        self.stats.rx_bytes += frame.len() as u64;
        Some(frame)
    }

    /// Collect transmit buffers the device has finished sending.
    fn reclaim_tx(&mut self) {
        while let Some((i, _)) = self.tx.pop_used() {
            self.tx_free.push(i);
        }
    }

    /// Send one frame: `fill` writes its `len` bytes.
    fn send_frame<R>(&mut self, len: usize, fill: impl FnOnce(&mut [u8]) -> R) -> R {
        let len = len.min(MTU);
        self.reclaim_tx();
        // QEMU sends synchronously when notified, so a buffer is always
        // free in practice; just in case, wait for one.
        let i = loop {
            if let Some(i) = self.tx_free.pop() {
                break i;
            }
            core::hint::spin_loop();
            self.reclaim_tx();
        };
        let buf = self.tx.buffer(i);
        buf[..HEADER_LEN].fill(0); // no offloads: all-zero virtio_net_hdr
        let result = fill(&mut buf[HEADER_LEN..HEADER_LEN + len]);
        self.tx.push(i, HEADER_LEN + len, 0);
        self.notify(QUEUE_TX);
        self.stats.tx_packets += 1;
        self.stats.tx_bytes += len as u64;
        result
    }
}

fn align_up(x: usize, align: usize) -> usize {
    x.div_ceil(align) * align
}

// --- smoltcp glue -------------------------------------------------------------
//
// smoltcp asks a `Device` for *tokens*: an RxToken is "a received frame you
// may look at", a TxToken is "permission to build one frame to send". This
// lets smoltcp work without copying in zero-copy drivers; ours copies the
// received frame out (simple, and frames are small).

pub struct RxToken(Vec<u8>);
pub struct TxToken<'a>(&'a mut VirtioNet);

impl phy::Device for VirtioNet {
    type RxToken<'a> = RxToken;
    type TxToken<'a> = TxToken<'a>;

    fn receive(&mut self, _now: Instant) -> Option<(RxToken, TxToken<'_>)> {
        let frame = self.receive_frame()?;
        Some((RxToken(frame), TxToken(self)))
    }

    fn transmit(&mut self, _now: Instant) -> Option<TxToken<'_>> {
        Some(TxToken(self))
    }

    fn capabilities(&self) -> DeviceCapabilities {
        let mut caps = DeviceCapabilities::default();
        caps.medium = Medium::Ethernet;
        caps.max_transmission_unit = MTU;
        caps
    }
}

impl phy::RxToken for RxToken {
    fn consume<R, F: FnOnce(&[u8]) -> R>(self, f: F) -> R {
        f(&self.0)
    }
}

impl phy::TxToken for TxToken<'_> {
    fn consume<R, F: FnOnce(&mut [u8]) -> R>(self, len: usize, f: F) -> R {
        self.0.send_frame(len, f)
    }
}