//! Bounded, read-only virtio block transport shared by TurkFS shell probes.
//! A read is accepted only after the device completes its descriptor chain.
use super::*;
use turkos_kernel::{virtio_mmio, virtio_ring};
use os_virtio_descriptors::{Descriptor as OwnedDescriptor, Error as OwnerError, Pool, NEXT, WRITE};
use os_virtio_indices::{Error as IndexError, Queue as IndexQueue};
use os_virtio_completions::{Error as CompletionError, Table as CompletionTable};

pub(super) struct RoDisk {
    base: u64,
    capacity: u64,
    frames: [u64; 4],
    frequency: u64,
    next: u16,
    free_before: u64,
    lease: virtio_ring::DmaLease,
    descriptor_owners: Pool<8, 1>,
    indices: IndexQueue<8>,
    completions: CompletionTable<8>,
    get_marker: &'static str,
}

unsafe fn read32(base: u64, offset: u64) -> u32 {
    unsafe { read_volatile((base + offset) as *const u32) }
}

unsafe fn write32(base: u64, offset: u64, value: u32) {
    unsafe { write_volatile((base + offset) as *mut u32, value) }
}

unsafe fn write64(base: u64, offset: u64, address: u64) {
    unsafe {
        write32(base, offset, address as u32);
        write32(base, offset + 4, (address >> 32) as u32);
    }
}

fn counter() -> u64 {
    let value: u64;
    unsafe {
        asm!("mrs {value}, CNTPCT_EL0", value = out(reg) value, options(nomem, nostack, preserves_flags))
    };
    value
}

impl RoDisk {
    pub(super) fn open(
        boot: &BootData,
        allocator: &mut memory::FrameAllocator,
        get_marker: &'static str,
    ) -> Result<Self, &'static str> {
        let mut block = None;
        let blob = unsafe {
            core::slice::from_raw_parts(boot.dtb_address as *const u8, boot.dtb_size)
        };
        let nodes = dtb::discover_virtio_mmio(blob).map_err(|_| "block DTB nodes")?;
        for node in nodes.as_slice() {
            let base = node.base;
            let (magic, version, id) =
                unsafe { (read32(base, 0), read32(base, 4), read32(base, 8)) };
            if magic != virtio_mmio::MAGIC || version != 2 {
                return Err("modern transport identity");
            }
            if id == 2 && block.replace(base).is_some() {
                return Err("duplicate block device");
            }
            if id != 0 && id != 1 && id != 2 && id != 4 {
                return Err("unexpected device");
            }
        }
        let base = block.ok_or("block device missing")?;
        let capacity =
            unsafe { u64::from(read32(base, 0x100)) | u64::from(read32(base, 0x104)) << 32 };
        if capacity < 18 {
            return Err("disk too small");
        }
        let free_before = allocator.remaining_pages();
        let mut lease = virtio_ring::DmaLease::allocate(
            allocator,
            boot.hardware.ram_base,
            boot.hardware.ram_size,
        )
        .map_err(|_| "DMA lease")?;
        let frames = lease.frames().map_err(|_| "DMA frames")?;
        let [desc, avail, used, _request] = frames;
        virtio_ring::split_ring_pages(
            desc,
            avail,
            used,
            8,
            boot.hardware.ram_base,
            boot.hardware.ram_size,
        )
        .map_err(|_| "ring geometry")?;
        unsafe {
            for frame in frames {
                core::ptr::write_bytes(frame as *mut u8, 0, memory::PAGE_SIZE as usize);
            }
            write32(base, 0x70, 0);
            if read32(base, 0x70) != 0 {
                return Err("reset");
            }
            write32(base, 0x70, 3);
            write32(base, 0x14, 0);
            let low = read32(base, 0x10);
            write32(base, 0x14, 1);
            let high = read32(base, 0x10);
            let (accepted_low, accepted_high) =
                virtio_mmio::select_modern_read_only_features(low, high)
                    .map_err(|_| "requires RO modern device")?;
            write32(base, 0x24, 0);
            write32(base, 0x20, accepted_low);
            write32(base, 0x24, 1);
            write32(base, 0x20, accepted_high);
            write32(base, 0x70, 0x0b);
            if read32(base, 0x70) != 0x0b {
                return Err("FEATURES_OK");
            }
            write32(base, 0x30, 0);
            if read32(base, 0x34) < 8 || read32(base, 0x44) != 0 {
                return Err("queue capacity");
            }
            write32(base, 0x38, 8);
            write64(base, 0x80, desc);
            write64(base, 0x90, avail);
            write64(base, 0xa0, used);
            asm!("dmb oshst", options(nostack, preserves_flags));
            write32(base, 0x44, 1);
            if read32(base, 0x44) != 1 {
                return Err("QueueReady");
            }
            lease.mark_ready().map_err(|_| "lease ready")?;
            write32(base, 0x70, 0x0f);
            if read32(base, 0x70) != 0x0f {
                return Err("DRIVER_OK");
            }
        }
        let frequency: u64;
        unsafe {
            asm!("mrs {value}, CNTFRQ_EL0", value = out(reg) frequency, options(nomem, nostack, preserves_flags))
        };
        if frequency == 0 {
            return Err("zero counter frequency");
        }
        Ok(Self {
            base,
            capacity,
            frames,
            frequency,
            next: 1,
            free_before,
            lease,
            descriptor_owners: Pool::new(8).map_err(|_| "descriptor pool")?,
            indices: IndexQueue::new(8).map_err(|_| "index queue")?,
            completions: CompletionTable::new(8).map_err(|_| "completion table")?,
            get_marker,
        })
    }

    pub(super) const fn capacity(&self) -> u64 {
        self.capacity
    }

    pub(super) fn reject_cycle_without_publish(&mut self) -> Result<(), &'static str> {
        let before = self.descriptor_owners.snapshot();
        let first = OwnedDescriptor {address: self.frames[3] + 512,
            length: 16, flags: NEXT, next: 0};
        let last = OwnedDescriptor {address: self.frames[3] + 528,
            length: 16, flags: WRITE, next: 0};
        if self.descriptor_owners.claim(0, &[(0, first), (0, last)])
            != Err(OwnerError::InvalidChain)
            || self.descriptor_owners.snapshot() != before {
            return Err("self-cycle changed descriptor ownership");
        }
        Ok(())
    }

    pub(super) fn reject_stale_used_without_publish(&mut self)
        -> Result<(), &'static str> {
        let used = self.frames[2];
        let observed = unsafe {
            u16::from_le(core::ptr::read_volatile((used + 2) as *const u16))
        };
        let before = self.indices.snapshot();
        if observed != before.used_index
            || self.indices.observe_used(observed, 0)
                != Err(IndexError::StaleObservation)
            || self.indices.observe_used(observed.wrapping_add(2), 0)
                != Err(IndexError::DeviceJump)
            || self.indices.snapshot() != before {
            return Err("stale used index changed ownership");
        }
        Ok(())
    }

    pub(super) fn reject_duplicate_completion(&mut self)
        -> Result<(), &'static str> {
        let before = self.completions.snapshot();
        if self.completions.complete(7, 513, 0) != Err(CompletionError::UnknownId)
            || self.completions.complete(0, 513, 0)
                != Err(CompletionError::UnknownId)
            || self.completions.snapshot() != before {
            return Err("foreign or replayed completion changed DMA ownership");
        }
        Ok(())
    }

    pub(super) fn read(
        &mut self,
        lba: u64,
        destination: &mut [u8; 512],
    ) -> Result<(), &'static str> {
        if virtio_mmio::checked_read_span(lba, 1, self.capacity).is_err() {
            return Err("LBA outside device");
        }
        if self.next > 8 {
            return Err("queue exhausted");
        }
        let [desc, avail, used, _request] = self.frames;
        let header = self
            .lease
            .request_subregion(0, 16)
            .map_err(|_| "header range")?;
        let status = self
            .lease
            .request_subregion(16, 1)
            .map_err(|_| "status range")?;
        let data = self
            .lease
            .request_subregion(512, 512)
            .map_err(|_| "data range")?;
        let chain = virtio_ring::read_descriptor_chain(header, data, status);
        self.lease
            .validate_read_chain(&chain)
            .map_err(|_| "descriptor ownership")?;
        let owned = chain.map(|entry| OwnedDescriptor {
            address: u64::from_le(entry.addr_le),
            length: u32::from_le(entry.len_le),
            flags: u16::from_le(entry.flags_le),
            next: u16::from_le(entry.next_le),
        });
        let owner = self.descriptor_owners.claim(0,
            &[(0, owned[0]), (1, owned[1]), (2, owned[2])])
            .map_err(|_| "descriptor chain claim")?;
        let completion_owner = self.completions.submit(0, data, 513)
            .map_err(|_| "completion owner claim")?;
        let sequence = self.next;
        unsafe {
            for (index, descriptor) in chain.into_iter().enumerate() {
                core::ptr::write_volatile(
                    (desc as *mut virtio_ring::Descriptor).add(index),
                    descriptor,
                );
            }
            core::ptr::write_volatile(
                header as *mut virtio_ring::BlockReadHeader,
                virtio_ring::block_read_header(lba),
            );
            core::ptr::write_bytes(data as *mut u8, 0xa5, 512);
            core::ptr::write_volatile(status as *mut u8, 0xff);
            core::ptr::write_volatile(
                (avail + 4 + u64::from(sequence - 1) * 2) as *mut u16,
                0u16.to_le(),
            );
            asm!("dmb oshst", options(nostack, preserves_flags));
            let ticket = self.indices.publish(0)
                .map_err(|_| "index publish")?;
            if ticket.avail_index != sequence {
                return Err("index publish sequence");
            }
            core::ptr::write_volatile((avail + 2) as *mut u16, sequence.to_le());
            asm!("dmb oshst", options(nostack, preserves_flags));
            self.descriptor_owners.expose(owner)
                .map_err(|_| "descriptor chain expose")?;
            write32(self.base, 0x50, 0);
        }
        let deadline = virtio_ring::deadline_after(counter(), self.frequency)
            .map_err(|_| "request deadline")?;
        loop {
            let index = unsafe { u16::from_le(core::ptr::read_volatile((used + 2) as *const u16)) };
            if index == sequence {
                break;
            }
            if virtio_ring::deadline_expired(counter(), deadline) {
                return Err("GET timeout");
            }
            core::hint::spin_loop();
        }
        let (id, length, result) = unsafe {
            asm!("dmb oshld", options(nostack, preserves_flags));
            let entry = used + 4 + u64::from(sequence - 1) * 8;
            (
                u32::from_le(core::ptr::read_volatile(entry as *const u32)),
                u32::from_le(core::ptr::read_volatile((entry + 4) as *const u32)),
                core::ptr::read_volatile(status as *const u8),
            )
        };
        if id != 0
            || virtio_ring::validate_read_completion_length(length, 512) != Ok(())
            || virtio_ring::read_block_status(result) != Ok(())
        {
            return Err("GET identity/length/status");
        }
        self.descriptor_owners.complete(owner, id as u16, length)
            .map_err(|_| "descriptor chain completion")?;
        let completed = self.indices.observe_used(sequence, id)
            .map_err(|_| "used index ownership")?;
        if completed.head != 0 || completed.avail_index != sequence {
            return Err("used index ticket");
        }
        let completion = self.completions.complete(id, length, result)
            .map_err(|_| "completion owner match")?;
        if completion.handle != completion_owner
            || completion.handle.owner != data {
            return Err("wrong DMA completion owner");
        }
        for (offset, byte) in destination.iter_mut().enumerate() {
            *byte = unsafe { core::ptr::read_volatile((data + offset as u64) as *const u8) };
        }
        write_str(self.get_marker);
        write_str(": LBA=0x");
        write_hex(lba as usize);
        write_str(" USEDIDX=0x");
        write_hex(sequence as usize);
        write_str(" USEDID=0x");
        write_hex(id as usize);
        write_str(" USEDLEN=0x");
        write_hex(length as usize);
        write_str(" STATUS=0x");
        write_hex(result as usize);
        write_str("\r\n");
        self.descriptor_owners.reclaim_used(owner)
            .map_err(|_| "descriptor chain reclaim")?;
        self.next += 1;
        Ok(())
    }

    pub(super) fn close(
        mut self,
        allocator: &mut memory::FrameAllocator,
    ) -> Result<(), &'static str> {
        unsafe {
            write32(self.base, 0x70, 0);
            if read32(self.base, 0x70) != 0 || read32(self.base, 0x44) != 0 {
                return Err("final reset");
            }
        }
        if self.descriptor_owners.count() != 0 {
            return Err("descriptor owners live at close");
        }
        self.descriptor_owners.observe_reset(0, 0)
            .map_err(|_| "descriptor owner reset")?;
        let [_, avail, used, _] = self.frames;
        unsafe {
            core::ptr::write_volatile((avail + 2) as *mut u16, 0);
            core::ptr::write_volatile((used + 2) as *mut u16, 0);
            asm!("dmb oshst", options(nostack, preserves_flags));
        }
        self.indices.observe_reset(0, 0, 0, 0)
            .map_err(|_| "index owner reset")?;
        if self.completions.snapshot().outstanding != 0 {
            return Err("completion DMA owner live at close");
        }
        self.completions.observe_reset(0, 0)
            .map_err(|_| "completion owner reset")?;
        self.lease.mark_reset().map_err(|_| "lease reset")?;
        self.lease.release(allocator).map_err(|_| "lease release")?;
        if allocator.remaining_pages() != self.free_before {
            return Err("DMA leak");
        }
        Ok(())
    }
}