//! Driver for virtio-blk: the paravirtual disk QEMU/KVM offer
//! (`-drive file=disk.img,if=virtio`).
//!
//! Why virtio and not ATA PIO? ATA PIO is simpler (copy 256 words through an
//! I/O port per sector) but it's a 1980s interface that real machines barely
//! have any more. virtio is what VMs actually use, and it teaches the core
//! idea of every modern device: *shared memory queues + DMA*. We use the
//! **legacy** (virtio 0.9.5) interface, which QEMU's "transitional" virtio
//! devices still offer: its registers are a handful of I/O ports (BAR0), and
//! setting it up is ~40 lines. The modern (1.0) interface uses PCI
//! capabilities + MMIO; the queue format is the same.
//!
//! The virtqueue: three arrays in physical memory that driver and device share.
//!   * descriptor table: {physical addr, length, flags, next} -- buffers
//!   * available ring: driver writes "descriptor chain N is ready"
//!   * used ring: device writes "I finished chain N"
//! A block request is a chain of 3 descriptors:
//!   header {type: read/write, sector}  ->  data (512*n bytes)  ->  status byte
//! We put the chain in the available ring, write the queue number to the
//! "notify" port, and wait until the used ring's index moves.
//!
//! Simplifications: one request at a time; we *poll* the used ring instead of
//! using the interrupt (we even ask the device not to interrupt us), and we
//! copy data through one page-sized "bounce buffer", so a request moves at
//! most 8 sectors. Reference: virtio spec 1.2 sections 2.7 (split virtqueues),
//! 4.1.4.8 (legacy PCI interface), 5.2 (block device); wiki.osdev.org/Virtio

use super::{BlockDevice, SECTOR_SIZE, check_range};
use crate::errno::Errno;
use crate::memory;
use crate::pci::PciDevice;
use core::ptr::{read_volatile, write_volatile};
use core::sync::atomic::{Ordering, fence};
use spin::Mutex;
use x86_64::PhysAddr;
use x86_64::instructions::port::Port;

// Legacy virtio PCI registers (offsets into the I/O BAR).
const REG_DEVICE_FEATURES: u16 = 0x00;
const REG_GUEST_FEATURES: u16 = 0x04;
const REG_QUEUE_ADDRESS: u16 = 0x08; // physical page number of the queue
const REG_QUEUE_SIZE: u16 = 0x0c;
const REG_QUEUE_SELECT: u16 = 0x0e;
const REG_QUEUE_NOTIFY: u16 = 0x10;
const REG_DEVICE_STATUS: u16 = 0x12;
const REG_ISR_STATUS: u16 = 0x13;
const REG_CAPACITY: u16 = 0x14; // device config: disk size in sectors (u64)

// Device status bits (we tell the device how far initialization got).
const STATUS_ACKNOWLEDGE: u8 = 1; // "I see you"
const STATUS_DRIVER: u8 = 2; // "I know how to drive you"
const STATUS_DRIVER_OK: u8 = 4; // "ready to go"

// Descriptor flags.
const DESC_NEXT: u16 = 1; // chain continues at `next`
const DESC_WRITE: u16 = 2; // device writes into this buffer (vs. reads it)
const AVAIL_NO_INTERRUPT: u16 = 1; // "don't interrupt me, I'm polling"

const REQ_READ: u32 = 0; // VIRTIO_BLK_T_IN
const REQ_WRITE: u32 = 1; // VIRTIO_BLK_T_OUT

/// Max sectors per request = what fits in our one-page bounce buffer.
const MAX_SECTORS: usize = 4096 / SECTOR_SIZE;

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

#[derive(Debug)]
pub enum VirtioError {
    NoIoBar,
    QueueUnavailable,
    OutOfMemory,
}

/// PCI vendor 0x1af4 (Red Hat/virtio), device 0x1001 = transitional block device.
pub fn is_virtio_blk(dev: &PciDevice) -> bool {
    dev.vendor_id == 0x1af4 && dev.device_id == 0x1001
}

pub struct VirtioBlk {
    io_base: u16,
    sectors: u64,
    queue: Mutex<Queue>,
}

struct Queue {
    size: u16,
    /// Physical address of the queue memory (descriptors, then available
    /// ring, then -- at the next page boundary -- the used ring).
    phys: PhysAddr,
    used_offset: usize,
    /// Our copy of the used ring index: how many requests we've seen finish.
    last_used: u16,
    /// One page for the request header (offset 0) and status byte (offset 16).
    header_page: PhysAddr,
    /// One page bounce buffer for the data.
    data_page: PhysAddr,
}

impl VirtioBlk {
    pub fn new(pci: &PciDevice) -> Result<Self, VirtioError> {
        let bar0 = pci.bar(0);
        if bar0 & 1 == 0 {
            return Err(VirtioError::NoIoBar); // legacy interface lives in I/O space
        }
        pci.enable_io_and_bus_master(); // bus master = allowed to DMA into our RAM
        let io_base = (bar0 & !0x3) as u16;
        let dev = VirtioBlk { io_base, sectors: 0, queue: Mutex::new(Queue::empty()) };

        // 1. Reset, then announce ourselves.
        dev.write8(REG_DEVICE_STATUS, 0);
        dev.write8(REG_DEVICE_STATUS, STATUS_ACKNOWLEDGE);
        dev.write8(REG_DEVICE_STATUS, STATUS_ACKNOWLEDGE | STATUS_DRIVER);
        // 2. Feature negotiation: the device lists what it can do; we accept
        //    nothing optional (plain read/write is always supported).
        let _offered = dev.read32(REG_DEVICE_FEATURES);
        dev.write32(REG_GUEST_FEATURES, 0);
        // 3. Set up queue 0 (block devices have one request queue).
        dev.write16(REG_QUEUE_SELECT, 0);
        let size = dev.read16(REG_QUEUE_SIZE); // the device decides (QEMU: 256)
        if size == 0 {
            return Err(VirtioError::QueueUnavailable);
        }
        // Legacy layout: descriptors (16 bytes each) + available ring
        // (2+2+2*size+2), then the used ring (2+2+8*size+2) page-aligned.
        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).ok_or(VirtioError::OutOfMemory)?;
        let header_page = memory::alloc_contiguous_frames(1).ok_or(VirtioError::OutOfMemory)?;
        let data_page = memory::alloc_contiguous_frames(1).ok_or(VirtioError::OutOfMemory)?;
        dev.write32(REG_QUEUE_ADDRESS, (phys.as_u64() / 4096) as u32);
        // 4. Read the disk size from the device config area, then go.
        let lo = dev.read32(REG_CAPACITY) as u64;
        let hi = dev.read32(REG_CAPACITY + 4) as u64;
        dev.write8(REG_DEVICE_STATUS, STATUS_ACKNOWLEDGE | STATUS_DRIVER | STATUS_DRIVER_OK);

        let queue = Queue { size, phys, used_offset, last_used: 0, header_page, data_page };
        unsafe { write_volatile(queue.avail_flags(), AVAIL_NO_INTERRUPT) };
        Ok(VirtioBlk { io_base, sectors: hi << 32 | lo, queue: Mutex::new(queue) })
    }

    /// Do one request of up to MAX_SECTORS. For writes, the data must
    /// already be in the bounce buffer; for reads, it's there afterwards.
    fn request(&self, q: &mut Queue, kind: u32, lba: u64, sectors: usize) -> Result<(), Errno> {
        let header = memory::phys_to_virt(q.header_page).as_mut_ptr::<u8>();
        unsafe {
            write_volatile(header as *mut u32, kind); // type
            write_volatile(header.add(4) as *mut u32, 0); // reserved
            write_volatile(header.add(8) as *mut u64, lba); // first sector
            write_volatile(header.add(16), 0xff); // status: device sets 0 = OK
        }
        // The chain: descriptors 0 -> 1 -> 2 (always the same three, since we
        // only ever have one request in flight).
        let data_flags = if kind == REQ_READ { DESC_WRITE } else { 0 };
        q.set_desc(0, q.header_page.as_u64(), 16, DESC_NEXT, 1);
        q.set_desc(1, q.data_page.as_u64(), (sectors * SECTOR_SIZE) as u32, data_flags | DESC_NEXT, 2);
        q.set_desc(2, q.header_page.as_u64() + 16, 1, DESC_WRITE, 0);

        // Publish chain 0 in the available ring, then bump its index.
        unsafe {
            let idx = read_volatile(q.avail_idx());
            write_volatile(q.avail_ring(idx % q.size), 0);
            fence(Ordering::SeqCst); // ring entry visible before the index
            write_volatile(q.avail_idx(), idx.wrapping_add(1));
            fence(Ordering::SeqCst);
        }
        self.write16(REG_QUEUE_NOTIFY, 0); // "queue 0 has work" -> QEMU does the I/O

        // Poll the used ring until the device reports completion.
        while unsafe { read_volatile(q.used_idx()) } == q.last_used {
            core::hint::spin_loop();
        }
        fence(Ordering::SeqCst);
        q.last_used = q.last_used.wrapping_add(1);
        let _ = self.read8(REG_ISR_STATUS); // reading acknowledges (clears) it

        match unsafe { read_volatile(header.add(16)) } {
            0 => Ok(()),
            _ => Err(Errno::EIO),
        }
    }

    fn read8(&self, reg: u16) -> u8 {
        unsafe { Port::<u8>::new(self.io_base + reg).read() }
    }
    fn write8(&self, reg: u16, v: u8) {
        unsafe { Port::<u8>::new(self.io_base + reg).write(v) }
    }
    fn read16(&self, reg: u16) -> u16 {
        unsafe { Port::<u16>::new(self.io_base + reg).read() }
    }
    fn write16(&self, reg: u16, v: u16) {
        unsafe { Port::<u16>::new(self.io_base + reg).write(v) }
    }
    fn read32(&self, reg: u16) -> u32 {
        unsafe { Port::<u32>::new(self.io_base + reg).read() }
    }
    fn write32(&self, reg: u16, v: u32) {
        unsafe { Port::<u32>::new(self.io_base + reg).write(v) }
    }
}

impl BlockDevice for VirtioBlk {
    fn sector_count(&self) -> u64 {
        self.sectors
    }

    fn read_sectors(&self, lba: u64, buf: &mut [u8]) -> Result<(), Errno> {
        check_range(self, lba, buf.len())?;
        // Interrupts off while we hold the queue lock (see thread.rs locking
        // rule): a request takes microseconds in QEMU.
        x86_64::instructions::interrupts::without_interrupts(|| {
            let mut q = self.queue.lock();
            for (i, chunk) in buf.chunks_mut(MAX_SECTORS * SECTOR_SIZE).enumerate() {
                let sectors = chunk.len() / SECTOR_SIZE;
                self.request(&mut q, REQ_READ, lba + (i * MAX_SECTORS) as u64, sectors)?;
                let src = memory::phys_to_virt(q.data_page).as_ptr::<u8>();
                unsafe { core::ptr::copy_nonoverlapping(src, chunk.as_mut_ptr(), chunk.len()) };
            }
            Ok(())
        })
    }

    fn write_sectors(&self, lba: u64, buf: &[u8]) -> Result<(), Errno> {
        check_range(self, lba, buf.len())?;
        x86_64::instructions::interrupts::without_interrupts(|| {
            let mut q = self.queue.lock();
            for (i, chunk) in buf.chunks(MAX_SECTORS * SECTOR_SIZE).enumerate() {
                let dst = memory::phys_to_virt(q.data_page).as_mut_ptr::<u8>();
                unsafe { core::ptr::copy_nonoverlapping(chunk.as_ptr(), dst, chunk.len()) };
                self.request(&mut q, REQ_WRITE, lba + (i * MAX_SECTORS) as u64, chunk.len() / SECTOR_SIZE)?;
            }
            Ok(())
        })
    }
}

impl Queue {
    const fn empty() -> Self {
        Queue {
            size: 0,
            phys: PhysAddr::zero(),
            used_offset: 0,
            last_used: 0,
            header_page: PhysAddr::zero(),
            data_page: PhysAddr::zero(),
        }
    }
    fn base(&self) -> *mut u8 {
        memory::phys_to_virt(self.phys).as_mut_ptr()
    }
    fn set_desc(&self, i: usize, addr: u64, len: u32, flags: u16, next: u16) {
        let desc = unsafe { (self.base() as *mut Descriptor).add(i) };
        unsafe { write_volatile(desc, Descriptor { addr, len, flags, next }) };
    }
    fn avail_base(&self) -> *mut u16 {
        unsafe { self.base().add(16 * self.size as usize) as *mut u16 }
    }
    fn avail_flags(&self) -> *mut u16 {
        self.avail_base()
    }
    fn avail_idx(&self) -> *mut u16 {
        unsafe { self.avail_base().add(1) }
    }
    fn avail_ring(&self, i: u16) -> *mut u16 {
        unsafe { self.avail_base().add(2 + i as usize) }
    }
    fn used_idx(&self) -> *const u16 {
        unsafe { (self.base().add(self.used_offset) as *const u16).add(1) }
    }
}

fn align_up(x: usize, align: usize) -> usize {
    (x + align - 1) & !(align - 1)
}