const VIRTIO_MAGIC: u32 = 0x74726976;
const VIRTIO_BLK_DEVICE_ID: u32 = 2;
const VIRTIO_MMIO_FIRST: usize = 0x10001000;
const VIRTIO_MMIO_STRIDE: usize = 0x1000;
const VIRTIO_MMIO_SLOTS: usize = 8;

const VIRTIO_STATUS_ACK: u32 = 1;
const VIRTIO_STATUS_DRIVER: u32 = 2;
const VIRTIO_STATUS_FEATURES_OK: u32 = 8;
const VIRTIO_STATUS_DRIVER_OK: u32 = 4;

const QUEUE_SIZE: usize = 8;
const SECTOR_SIZE: usize = 512;
const PAGE_SIZE: usize = 4096;

// Legacy VirtIO layout: used ring must start at page boundary after avail ring.
// descs: 8 * 16 = 128 bytes, avail: 22 bytes → total 150 → pad to 4096.
#[repr(C, align(4096))]
struct VirtqMem {
    descs: [VirtqDesc; QUEUE_SIZE],
    avail: VirtqAvail,
    _pad: [u8; 3946],
    used: VirtqUsed,
}

#[repr(C)]
#[derive(Clone, Copy)]
struct VirtqDesc {
    addr: u64,
    len: u32,
    flags: u16,
    next: u16,
}

#[repr(C)]
struct VirtqAvail {
    flags: u16,
    idx: u16,
    ring: [u16; QUEUE_SIZE],
    used_event: u16,
}

#[repr(C)]
struct VirtqUsed {
    flags: u16,
    idx: u16,
    ring: [VirtqUsedElem; QUEUE_SIZE],
}

#[repr(C)]
#[derive(Clone, Copy)]
struct VirtqUsedElem {
    id: u32,
    len: u32,
}

#[repr(C)]
struct VirtioBlkReqHeader {
    req_type: u32,
    reserved: u32,
    sector: u64,
}

static mut VQUEUE: VirtqMem = VirtqMem {
    descs: [const { VirtqDesc { addr: 0, len: 0, flags: 0, next: 0 } }; QUEUE_SIZE],
    avail: VirtqAvail { flags: 0, idx: 0, ring: [0; QUEUE_SIZE], used_event: 0 },
    _pad: [0; 3946],
    used: VirtqUsed {
        flags: 0, idx: 0,
        ring: [const { VirtqUsedElem { id: 0, len: 0 } }; QUEUE_SIZE],
    },
};

static mut REQ_HEADER: VirtioBlkReqHeader = VirtioBlkReqHeader {
    req_type: 0, reserved: 0, sector: 0,
};

static mut REQ_STATUS: u8 = 0xFF;

fn mmio_read(base: usize, offset: usize) -> u32 {
    unsafe { core::ptr::read_volatile((base + offset) as *const u32) }
}

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

pub struct VirtioBlk {
    base: usize,
    legacy: bool,
}

impl VirtioBlk {
    pub fn probe() -> Option<Self> {
        for slot in 0..VIRTIO_MMIO_SLOTS {
            let base = VIRTIO_MMIO_FIRST + slot * VIRTIO_MMIO_STRIDE;
            let magic = mmio_read(base, 0x000);
            let device_id = mmio_read(base, 0x008);
            if magic == VIRTIO_MAGIC && device_id == VIRTIO_BLK_DEVICE_ID {
                return Self::init_at(base);
            }
        }
        None
    }

    fn init_at(base: usize) -> Option<Self> {
        let version = mmio_read(base, 0x004);
        let legacy = version == 1;

        // Reset
        mmio_write(base, 0x070, 0);
        while mmio_read(base, 0x070) != 0 {}

        // Acknowledge + Driver
        mmio_write(base, 0x070, VIRTIO_STATUS_ACK);
        mmio_write(base, 0x070, mmio_read(base, 0x070) | VIRTIO_STATUS_DRIVER);

        if legacy {
            // Legacy: set guest page size
            mmio_write(base, 0x028, PAGE_SIZE as u32);
        }

        // Feature negotiation
        mmio_write(base, 0x014, 0);
        let _features = mmio_read(base, 0x010);
        mmio_write(base, 0x024, 0);
        mmio_write(base, 0x020, 0);

        if !legacy {
            mmio_write(base, 0x070, mmio_read(base, 0x070) | VIRTIO_STATUS_FEATURES_OK);
            if mmio_read(base, 0x070) & VIRTIO_STATUS_FEATURES_OK == 0 {
                return None;
            }
        }

        // Queue setup
        mmio_write(base, 0x030, 0);
        let queue_max = mmio_read(base, 0x034);
        if queue_max == 0 { return None; }
        let queue_num = (QUEUE_SIZE as u32).min(queue_max);
        mmio_write(base, 0x038, queue_num);

        unsafe {
            let vq_addr = core::ptr::addr_of!(VQUEUE) as usize;

            if legacy {
                mmio_write(base, 0x040, (vq_addr / PAGE_SIZE) as u32);
            } else {
                let desc_addr = vq_addr as u64;
                let avail_addr = (vq_addr + core::mem::offset_of!(VirtqMem, avail)) as u64;
                let used_addr = (vq_addr + core::mem::offset_of!(VirtqMem, used)) as u64;

                mmio_write(base, 0x080, desc_addr as u32);
                mmio_write(base, 0x084, (desc_addr >> 32) as u32);
                mmio_write(base, 0x090, avail_addr as u32);
                mmio_write(base, 0x094, (avail_addr >> 32) as u32);
                mmio_write(base, 0x0A0, used_addr as u32);
                mmio_write(base, 0x0A4, (used_addr >> 32) as u32);
                mmio_write(base, 0x044, 1);
            }
        }

        mmio_write(base, 0x070, mmio_read(base, 0x070) | VIRTIO_STATUS_DRIVER_OK);

        Some(Self { base, legacy })
    }

    pub fn read_sector(&mut self, sector: u32, buf: &mut [u8; SECTOR_SIZE]) -> bool {
        unsafe {
            REQ_HEADER.req_type = 0;
            REQ_HEADER.reserved = 0;
            REQ_HEADER.sector = sector as u64;
            REQ_STATUS = 0xFF;

            VQUEUE.descs[0] = VirtqDesc {
                addr: core::ptr::addr_of!(REQ_HEADER) as u64,
                len: 16,
                flags: 1, // NEXT
                next: 1,
            };
            VQUEUE.descs[1] = VirtqDesc {
                addr: buf.as_ptr() as u64,
                len: SECTOR_SIZE as u32,
                flags: 2 | 1, // WRITE | NEXT
                next: 2,
            };
            VQUEUE.descs[2] = VirtqDesc {
                addr: core::ptr::addr_of!(REQ_STATUS) as u64,
                len: 1,
                flags: 2, // WRITE
                next: 0,
            };

            let avail_idx = core::ptr::read_volatile(core::ptr::addr_of!(VQUEUE.avail.idx));
            VQUEUE.avail.ring[(avail_idx as usize) % QUEUE_SIZE] = 0;
            core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
            core::ptr::write_volatile(
                core::ptr::addr_of_mut!(VQUEUE.avail.idx),
                avail_idx.wrapping_add(1),
            );
            core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);

            mmio_write(self.base, 0x050, 0);

            let expected = avail_idx.wrapping_add(1);
            let mut timeout = 10_000_000u32;
            while core::ptr::read_volatile(core::ptr::addr_of!(VQUEUE.used.idx)) != expected {
                core::hint::spin_loop();
                timeout -= 1;
                if timeout == 0 { return false; }
            }

            REQ_STATUS == 0
        }
    }
}