pub trait MmioAccess {
    fn read(&self, offset: usize) -> u32;
    fn write(&self, offset: usize, val: u32);
}

pub const MAGIC: u32 = 0x74726976;

pub const REG_MAGIC: usize = 0x000;
pub const REG_VERSION: usize = 0x004;
pub const REG_DEVICE_ID: usize = 0x008;
pub const REG_STATUS: usize = 0x070;
pub const REG_GUEST_PAGE_SIZE: usize = 0x028;
pub const REG_FEATURES_SEL: usize = 0x014;
pub const REG_FEATURES: usize = 0x010;
pub const REG_DRIVER_FEATURES_SEL: usize = 0x024;
pub const REG_DRIVER_FEATURES: usize = 0x020;
pub const REG_QUEUE_SEL: usize = 0x030;
pub const REG_QUEUE_NUM_MAX: usize = 0x034;
pub const REG_QUEUE_NUM: usize = 0x038;
pub const REG_QUEUE_PFN: usize = 0x040;
pub const REG_QUEUE_NOTIFY: usize = 0x050;

pub const STATUS_ACK: u32 = 1;
pub const STATUS_DRIVER: u32 = 2;
pub const STATUS_DRIVER_OK: u32 = 4;

pub const DEVICE_ID_BLOCK: u32 = 2;
pub const DEVICE_ID_GPU: u32 = 16;

pub const PAGE_SIZE: usize = 4096;

pub struct DirectMmio {
    base: usize,
}

impl DirectMmio {
    pub fn new(base: usize) -> Self {
        Self { base }
    }

    pub fn base(&self) -> usize {
        self.base
    }
}

impl MmioAccess for DirectMmio {
    fn read(&self, offset: usize) -> u32 {
        unsafe { core::ptr::read_volatile((self.base + offset) as *const u32) }
    }

    fn write(&self, offset: usize, val: u32) {
        unsafe { core::ptr::write_volatile((self.base + offset) as *mut u32, val) }
    }
}

pub fn init_legacy_device(mmio: &impl MmioAccess, queue_phys: usize, queue_size: u32) {
    mmio.write(REG_STATUS, 0);
    while mmio.read(REG_STATUS) != 0 {}

    mmio.write(REG_STATUS, STATUS_ACK);
    mmio.write(REG_STATUS, mmio.read(REG_STATUS) | STATUS_DRIVER);
    mmio.write(REG_GUEST_PAGE_SIZE, PAGE_SIZE as u32);

    mmio.write(REG_FEATURES_SEL, 0);
    mmio.write(REG_DRIVER_FEATURES_SEL, 0);
    mmio.write(REG_DRIVER_FEATURES, 0);

    mmio.write(REG_QUEUE_SEL, 0);
    let queue_max = mmio.read(REG_QUEUE_NUM_MAX);
    mmio.write(REG_QUEUE_NUM, queue_size.min(queue_max));
    mmio.write(REG_QUEUE_PFN, (queue_phys / PAGE_SIZE) as u32);

    mmio.write(REG_STATUS, mmio.read(REG_STATUS) | STATUS_DRIVER_OK);
}

pub fn probe_slot(base: usize, device_id: u32) -> Option<DirectMmio> {
    let mmio = DirectMmio::new(base);
    if mmio.read(REG_MAGIC) == MAGIC && mmio.read(REG_DEVICE_ID) == device_id {
        Some(mmio)
    } else {
        None
    }
}

pub const MMIO_FIRST: usize = 0x10001000;
pub const MMIO_STRIDE: usize = 0x1000;
pub const MMIO_SLOTS: usize = 8;

pub fn scan_for_device(device_id: u32) -> Option<DirectMmio> {
    for slot in 0..MMIO_SLOTS {
        if let Some(mmio) = probe_slot(MMIO_FIRST + slot * MMIO_STRIDE, device_id) {
            return Some(mmio);
        }
    }
    None
}