use crate::mmu::{phys_to_virt, virt_to_phys};
use hw_virtio::block::{self, DmaAddrs};
use hw_virtio::mmio::{self, DirectMmio};
use hw_virtio::queue::{QUEUE_SIZE, QueueAddrs, VirtioQueue};
use types::address::{PhysicalAddress, VirtualAddress};

const SECTOR_SIZE: usize = 512;

#[repr(C, align(4096))]
struct VirtqMem {
    _data: [u8; 8192],
}

static mut VQUEUE: VirtqMem = VirtqMem { _data: [0; 8192] };
static mut REQ_HEADER: block::BlockReqHeader = block::BlockReqHeader {
    req_type: 0,
    reserved: 0,
    sector: 0,
};
static mut REQ_STATUS: u8 = 0xFF;
static mut SECTOR_BUF: [u8; SECTOR_SIZE] = [0; SECTOR_SIZE];

pub struct VirtioBlk {
    mmio: DirectMmio,
    queue: VirtioQueue,
    dma: DmaAddrs,
}

impl VirtioBlk {
    pub fn probe() -> Option<Self> {
        let mmio = mmio::scan_for_device(phys_to_virt(mmio::MMIO_FIRST), mmio::DEVICE_ID_BLOCK)?;

        // The device sees physical addresses; the CPU uses kernel ones.
        let vq_va = core::ptr::addr_of!(VQUEUE) as usize;
        mmio::init_legacy_device(&mmio, virt_to_phys(vq_va), QUEUE_SIZE as u32);

        let queue = VirtioQueue::new(QueueAddrs::from_base(vq_va));

        let pair = |va: usize| {
            (
                PhysicalAddress::from_raw(virt_to_phys(va)),
                VirtualAddress::from_raw(va),
            )
        };
        let (header_pa, header_va) = pair(core::ptr::addr_of_mut!(REQ_HEADER) as usize);
        let (status_pa, status_va) = pair(core::ptr::addr_of_mut!(REQ_STATUS) as usize);
        let (data_pa, data_va) = pair(core::ptr::addr_of_mut!(SECTOR_BUF) as usize);
        let dma = DmaAddrs {
            header_pa,
            header_va,
            status_pa,
            status_va,
            data_pa,
            data_va,
        };

        Some(Self { mmio, queue, dma })
    }

    /// Physical address of the device's MMIO window.
    pub fn base(&self) -> usize {
        virt_to_phys(self.mmio.base())
    }

    pub fn read_sector(&mut self, sector: u32, buf: &mut [u8; SECTOR_SIZE]) -> bool {
        block::read_sector(&self.mmio, &self.queue, &self.dma, sector, buf)
    }
}