#![no_std]
#![forbid(unsafe_code)]

use hw_virtio::block::{self, DmaAddrs};
use hw_virtio::mmio::{MmioAccess, Registers};
use hw_virtio::queue::VirtioQueue;
use port_block_device::{BlockDevice, SectorId};
use types::address::{ByteOffset, PhysicalAddress, VirtualAddress};
use types::bounded::PageCount;

/// Drivers map their devices in the device window (0x01D0_0000 up, 64 KiB
/// each), clear of program images and heaps.
const MMIO_VA: VirtualAddress = VirtualAddress::from_raw(0x01D0_0000);
const VQUEUE_VA: VirtualAddress = VirtualAddress::from_raw(0x01D0_1000);
const REQ_HDR_VA: VirtualAddress = VirtualAddress::from_raw(0x01D0_3000);
const SECTOR_SIZE: usize = 512;

pub struct VirtioBlockDevice {
    mmio: Registers,
    queue: VirtioQueue,
    dma: DmaAddrs,
    sectors: usize,
}

impl VirtioBlockDevice {
    /// Takes over the virtio block device behind `mmio`. `None` if it is
    /// not one or any step fails.
    pub fn new(mmio: types::bounded::CapSlot) -> Option<Self> {
        let regs = Registers::new(libkos::device::map_mmio(mmio, MMIO_VA).ok()?);
        if regs.read(hw_virtio::mmio::REG_MAGIC) != hw_virtio::mmio::MAGIC
            || regs.read(hw_virtio::mmio::REG_DEVICE_ID) != hw_virtio::mmio::DEVICE_ID_BLOCK
        {
            return None;
        }
        let rings = libkos::device::dma(PageCount::new(2)?, VQUEUE_VA).ok()?;
        let req = libkos::device::dma(PageCount::new(1)?, REQ_HDR_VA).ok()?;
        hw_virtio::mmio::init_legacy_device(&regs, rings.physical, hw_virtio::queue::QUEUE_SIZE as u32);

        // Header, status byte and one sector share the request page.
        let at = |offset: usize| PhysicalAddress::from_raw(req.physical.raw() + offset);
        let dma = DmaAddrs {
            header_pa: at(0),
            header: req.window.sub(ByteOffset::new(0), 16),
            status_pa: at(16),
            status: req.window.sub(ByteOffset::new(16), 1),
            data_pa: at(32),
            data: req.window.sub(ByteOffset::new(32), SECTOR_SIZE),
        };
        let capacity = regs.read(block::CONFIG_CAPACITY) as u64
            | (regs.read(block::CONFIG_CAPACITY + 4) as u64) << 32;
        Some(Self {
            mmio: regs,
            queue: VirtioQueue::new(rings.window),
            dma,
            sectors: usize::try_from(capacity).ok()?,
        })
    }
}

impl BlockDevice for VirtioBlockDevice {
    fn sector_size(&self) -> usize { SECTOR_SIZE }
    fn sector_count(&self) -> usize { self.sectors }

    fn read_sector(&mut self, sector: SectorId, buf: &mut [u8]) -> bool {
        if buf.len() < SECTOR_SIZE || sector.raw() as usize >= self.sectors { return false; }
        let mut sector_buf = [0u8; SECTOR_SIZE];
        if !block::read_sector(&self.mmio, &self.queue, &self.dma, sector.raw(), &mut sector_buf) {
            return false;
        }
        buf[..SECTOR_SIZE].copy_from_slice(&sector_buf);
        true
    }

    fn write_sector(&mut self, sector: SectorId, buf: &[u8]) -> bool {
        if sector.raw() as usize >= self.sectors { return false; }
        block::write_sector(&self.mmio, &self.queue, &self.dma, sector.raw(), buf)
    }
}