//! 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};
use os_virtio_backpressure::{Budget, Error as BudgetError};
use os_block_geometry::{Error as GeometryError, Geometry};
use os_block_io::{Error as RequestError, Requests, Token as RequestToken};

pub(super) struct RoDisk {
    base: u64,
    capacity: u64,
    geometry: Geometry,
    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>,
    budget: Budget<8>,
    requests: Requests<8>,
    request_ids: [u64; 2],
    request_count: usize,
    foreign_completion_rejected: bool,
    backpressure_rejected: bool,
    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 geometry = Geometry::new(capacity).map_err(|_| "block geometry")?;
        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,
            4,
            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) < 4 || read32(base, 0x44) != 0 {
                return Err("queue capacity");
            }
            write32(base, 0x38, 4);
            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,
            geometry,
            frames,
            frequency,
            next: 1,
            free_before,
            lease,
            descriptor_owners: Pool::new(4).map_err(|_| "descriptor pool")?,
            indices: IndexQueue::new(4).map_err(|_| "index queue")?,
            completions: CompletionTable::new(4).map_err(|_| "completion table")?,
            budget: Budget::new(4).map_err(|_| "descriptor budget")?,
            requests: Requests::new(4, 1).map_err(|_| "block requests")?,
            request_ids: [0; 2],
            request_count: 0,
            foreign_completion_rejected: false,
            backpressure_rejected: false,
            get_marker,
        })
    }

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

    pub(super) const fn backpressure_rejected(&self) -> bool {
        self.backpressure_rejected
    }

    pub(super) fn request_ids(&self) -> [u64; 2] { self.request_ids }

    pub(super) fn foreign_completion_rejected(&self) -> bool {
        self.foreign_completion_rejected
    }

    pub(super) fn reject_stale_request_without_publish(&mut self)
        -> Result<(), &'static str> {
        let before = self.requests.snapshot();
        let descriptors = self.descriptor_owners.snapshot();
        let indices = self.indices.snapshot();
        let completions = self.completions.snapshot();
        let budget = self.budget.snapshot();
        let data = self.lease.request_subregion(512, 512)
            .map_err(|_| "stale probe buffer")?;
        let now = counter();
        let deadline = virtio_ring::deadline_after(now, self.frequency)
            .map_err(|_| "stale probe deadline")?;
        if self.requests.prepare(0, data, 0, deadline, now, 0)
                != Err(RequestError::StaleDevice)
            || self.requests.snapshot() != before
            || self.descriptor_owners.snapshot() != descriptors
            || self.indices.snapshot() != indices
            || self.completions.snapshot() != completions
            || self.budget.snapshot() != budget {
            return Err("stale device generation changed ownership");
        }
        Ok(())
    }

    pub(super) fn reject_outside_without_publish(&mut self)
        -> Result<(), &'static str> {
        let descriptors = self.descriptor_owners.snapshot();
        let indices = self.indices.snapshot();
        let completions = self.completions.snapshot();
        let budget = self.budget.snapshot();
        if self.geometry.read_span(self.capacity, 1)
                != Err(GeometryError::OutsideCapacity)
            || self.geometry.read_span(u64::MAX, 2)
                != Err(GeometryError::Overflow)
            || self.descriptor_owners.snapshot() != descriptors
            || self.indices.snapshot() != indices
            || self.completions.snapshot() != completions
            || self.budget.snapshot() != budget {
            return Err("outside LBA changed queue ownership");
        }
        Ok(())
    }

    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(())
    }

    fn abort_prepared_read(&mut self, request: RequestToken,
        owner: os_virtio_descriptors::Handle,
        completion: os_virtio_completions::Handle,
        budget_claimed: bool) -> Result<(), &'static str> {
        if budget_claimed {
            self.budget.complete(0).map_err(|_| "budget rollback")?;
        }
        self.completions.cancel_pending(completion)
            .map_err(|_| "completion rollback")?;
        self.descriptor_owners.cancel_prepared(owner)
            .map_err(|_| "descriptor rollback")?;
        self.requests.abort_prepared(request)
            .map_err(|_| "request rollback")?;
        Ok(())
    }

    pub(super) fn read(
        &mut self,
        lba: u64,
        destination: &mut [u8; 512],
    ) -> Result<(), &'static str> {
        if self.geometry.read_span(lba, 1).is_err() {
            return Err("LBA outside device");
        }
        if self.next > 4 {
            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 now = counter();
        let deadline = virtio_ring::deadline_after(now, self.frequency)
            .map_err(|_| "request deadline")?;
        let request = self.requests.prepare(0, data, lba, deadline, now, 1)
            .map_err(|_| "owned request prepare")?;
        let owner = match self.descriptor_owners.claim(0,
            &[(0, owned[0]), (1, owned[1]), (2, owned[2])]) {
            Ok(owner) => owner,
            Err(_) => {
                self.requests.abort_prepared(request).map_err(|_| "request rollback")?;
                return Err("descriptor chain claim");
            }
        };
        let completion_owner = match self.completions.submit(0, data, 513) {
            Ok(owner) => owner,
            Err(_) => {
                self.descriptor_owners.cancel_prepared(owner)
                    .map_err(|_| "descriptor rollback")?;
                self.requests.abort_prepared(request).map_err(|_| "request rollback")?;
                return Err("completion owner claim");
            }
        };
        let budget_claim = match self.budget.claim(0, 3) {
            Ok(claim) => claim,
            Err(_) => {
                self.abort_prepared_read(request, owner, completion_owner, false)?;
                return Err("descriptor budget claim");
            }
        };
        let sequence = self.next;
        if sequence == 1 {
            let before = self.budget.snapshot();
            if self.budget.claim(3, 3) != Err(BudgetError::Again)
                || self.budget.snapshot() != before {
                self.abort_prepared_read(request, owner, completion_owner, true)?;
                return Err("full descriptor budget changed first owner");
            }
            self.backpressure_rejected = true;
        }
        if let Err(_) = self.requests.publish(request, counter()) {
            self.abort_prepared_read(request, owner, completion_owner, true)?;
            return Err("owned request publish deadline");
        }
        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 foreign = RequestToken {id: request.id + 1, ..request};
        let before = self.requests.snapshot();
        if self.requests.complete(foreign, 0, 513, 0)
                != Err(RequestError::WrongCompletion)
            || self.requests.snapshot() != before {
            return Err("foreign request consumed DMA buffer");
        }
        self.foreign_completion_rejected = true;
        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");
        }
        if self.budget.complete(0) != Ok(budget_claim) {
            return Err("descriptor budget completion");
        }
        let matched = self.requests.complete(request, id, length, result)
            .map_err(|_| "owned request completion mismatch")?;
        if matched.buffer != data || matched.lba != lba
            || matched.deadline != deadline || matched.token != request {
            return Err("owned request fields changed");
        }
        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")?;
        if self.request_count < self.request_ids.len() {
            self.request_ids[self.request_count] = request.id;
        }
        self.request_count += 1;
        self.next += 1;
        Ok(())
    }

    /// Prepare one real read but hold its physical avail index until after the
    /// deadline. This remains conservative: all DMA owners are marked live.
    /// The late device completion is observed after cancellation and before
    /// reset; none of the live DMA owners are released on that late response.
    pub(super) fn unanswered_then_late(&mut self) -> Result<(), &'static str> {
        let [desc, avail, used, _] = self.frames;
        let header = self.lease.request_subregion(0, 16)
            .map_err(|_| "timeout header")?;
        let status = self.lease.request_subregion(16, 1)
            .map_err(|_| "timeout status")?;
        let data = self.lease.request_subregion(512, 512)
            .map_err(|_| "timeout data")?;
        self.geometry.read_span(0, 1).map_err(|_| "timeout geometry")?;
        let chain = virtio_ring::read_descriptor_chain(header, data, status);
        self.lease.validate_read_chain(&chain).map_err(|_| "timeout DMA chain")?;
        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 now = counter();
        let interval = self.frequency / 10;
        let deadline = virtio_ring::deadline_after(now, interval)
            .map_err(|_| "timeout deadline")?;
        let request = self.requests.prepare(0, data, 0, deadline, now, 1)
            .map_err(|_| "timeout request prepare")?;
        let owner = self.descriptor_owners.claim(0,
            &[(0, owned[0]), (1, owned[1]), (2, owned[2])])
            .map_err(|_| "timeout descriptor claim")?;
        self.completions.submit(0, data, 513)
            .map_err(|_| "timeout completion claim")?;
        self.budget.claim(0, 3).map_err(|_| "timeout budget claim")?;
        self.requests.publish(request, counter())
            .map_err(|_| "timeout request publish")?;
        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(0));
            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) as *mut u16, 0u16.to_le());
            asm!("dmb oshst", options(nostack, preserves_flags));
            if self.indices.publish(0).map_err(|_| "timeout index publish")?
                    .avail_index != 1 { return Err("timeout avail index"); }
            self.descriptor_owners.expose(owner)
                .map_err(|_| "timeout expose")?;
            // QueueReady is live, but avail.idx stays zero until cancellation.
        }
        while counter() < deadline {
            let observed = unsafe {
                u16::from_le(core::ptr::read_volatile((used + 2) as *const u16))
            };
            if observed != 0 { return Err("timeout read completed before avail"); }
            core::hint::spin_loop();
        }
        self.requests.timeout(request, counter())
            .map_err(|_| "timeout transition")?;
        if self.requests.snapshot().live != 1
            || self.descriptor_owners.count() != 1
            || self.completions.snapshot().outstanding != 1
            || self.budget.snapshot().used_descriptors != 3 {
            return Err("timeout released DMA early");
        }
        self.requests.cancel(request).map_err(|_| "timeout cancel")?;
        let cancelled = self.requests.snapshot();
        unsafe {
            core::ptr::write_volatile((avail + 2) as *mut u16, 1u16.to_le());
            asm!("dmb oshst", options(nostack, preserves_flags));
            write32(self.base, 0x50, 0);
        }
        let late_deadline = virtio_ring::deadline_after(counter(), self.frequency)
            .map_err(|_| "late deadline")?;
        loop {
            let observed = unsafe {
                u16::from_le(core::ptr::read_volatile((used + 2) as *const u16))
            };
            if observed == 1 { break; }
            if virtio_ring::deadline_expired(counter(), late_deadline) {
                return Err("late device completion absent");
            }
            core::hint::spin_loop();
        }
        let (id, length, result) = unsafe {
            asm!("dmb oshld", options(nostack, preserves_flags));
            let entry = used + 4;
            (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 || length != 513 || result != 0
            || self.requests.complete(request, id, length, result)
                != Err(RequestError::LateCompletion)
            || self.requests.snapshot() != cancelled
            || self.descriptor_owners.count() != 1
            || self.completions.snapshot().outstanding != 1 {
            return Err("late completion consumed cancelled DMA");
        }
        Ok(())
    }

    pub(super) fn close_after_timeout(
        self, allocator: &mut memory::FrameAllocator) -> Result<(), &'static str> {
        self.close_reset(allocator, true)
    }

    pub(super) fn close(
        self, allocator: &mut memory::FrameAllocator) -> Result<(), &'static str> {
        self.close_reset(allocator, false)
    }

    fn close_reset(
        mut self,
        allocator: &mut memory::FrameAllocator,
        allow_live: bool,
    ) -> 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 !allow_live && 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 !allow_live && self.completions.snapshot().outstanding != 0 {
            return Err("completion DMA owner live at close");
        }
        self.completions.observe_reset(0, 0)
            .map_err(|_| "completion owner reset")?;
        if !allow_live && self.requests.snapshot().live != 0 {
            return Err("block request buffer live at close");
        }
        self.requests.observe_reset(0, 0, 2)
            .map_err(|_| "block requests reset")?;
        if !allow_live && self.budget.snapshot().used_descriptors != 0 {
            return Err("descriptor budget live at close");
        }
        self.budget.observe_reset(0, 0)
            .map_err(|_| "descriptor budget 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(())
    }
}