// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

// UNSAFETY: test code implements a custom `GuestMemory` backing, which requires
// unsafe.
#![expect(unsafe_code)]
#![cfg(test)]

use crate::DeviceTraits;
use crate::PciInterruptModel;
use crate::Resources;
use crate::VirtioDevice;
use crate::VirtioQueue;
use crate::VirtioQueueCallbackWork;
use crate::queue::QueueParams;
use crate::spec::pci::*;
use crate::spec::queue::*;
use crate::spec::*;
use crate::transport::VirtioMmioDevice;
use crate::transport::VirtioPciDevice;
use chipset_device::mmio::ExternallyManagedMmioIntercepts;
use chipset_device::mmio::MmioIntercept;
use chipset_device::pci::PciConfigSpace;
use futures::StreamExt;
use guestmem::DoorbellRegistration;
use guestmem::GuestMemory;
use guestmem::GuestMemoryAccess;
use guestmem::GuestMemoryBackingError;
use inspect::InspectMut;
use pal_async::DefaultDriver;
use pal_async::async_test;
use pal_async::timer::PolledTimer;
use pal_async::wait::PolledWait;
use pal_event::Event;
use parking_lot::Mutex;
use pci_core::msi::MsiConnection;
use pci_core::spec::caps::CapabilityId;
use pci_core::spec::cfg_space;
use pci_core::test_helpers::TestPciInterruptController;
use std::collections::BTreeMap;
use std::future::poll_fn;
use std::io;
use std::ptr::NonNull;
use std::sync::Arc;
use std::time::Duration;
use task_control::AsyncRun;
use task_control::Cancelled;
use task_control::StopTask;
use task_control::TaskControl;
use test_with_tracing::test;
use vmcore::interrupt::Interrupt;
use vmcore::line_interrupt::LineInterrupt;
use vmcore::line_interrupt::test_helpers::TestLineInterruptTarget;
use vmcore::vm_task::SingleDriverBackend;
use vmcore::vm_task::VmTaskDriverSource;

// Device features - first bank
const VIRTIO_F_RING_INDIRECT_DESC: u32 = 0x10000000;
const VIRTIO_F_RING_EVENT_IDX: u32 = 0x20000000;
// Device features - second bank
const VIRTIO_F_VERSION_1: u32 = 1;
const VIRTIO_F_RING_PACKED: u32 = 4;

// Device status
const VIRTIO_ACKNOWLEDGE: u32 = 1;
const VIRTIO_DRIVER: u32 = 2;
const VIRTIO_DRIVER_OK: u32 = 4;
const VIRTIO_FEATURES_OK: u32 = 8;
const _VIRTIO_DEVICE_NEEDS_RESET: u32 = 0x40;
const _VIRTIO_FAILED: u32 = 0x80;

async fn must_recv_in_timeout<T: 'static + Send>(
    recv: &mut mesh::Receiver<T>,
    timeout: Duration,
) -> T {
    mesh::CancelContext::new()
        .with_timeout(timeout)
        .until_cancelled(recv.next())
        .await
        .unwrap()
        .unwrap()
}

async fn assert_no_recv_in_timeout<T: 'static + Send>(
    recv: &mut mesh::Receiver<T>,
    timeout: Duration,
) {
    if mesh::CancelContext::new()
        .with_timeout(timeout)
        .until_cancelled(recv.next())
        .await
        .is_ok()
    {
        panic!("Expected timeout, but received a value");
    }
}

#[derive(Default)]
struct VirtioTestMemoryAccess {
    memory_map: Mutex<MemoryMap>,
}

#[derive(Default)]
struct MemoryMap {
    map: BTreeMap<u64, (bool, Vec<u8>)>,
}

impl MemoryMap {
    fn get(&mut self, address: u64, len: usize) -> Option<(bool, &mut [u8])> {
        let (&base, &mut (writable, ref mut data)) = self.map.range_mut(..=address).last()?;
        let data = data
            .get_mut(usize::try_from(address - base).ok()?..)?
            .get_mut(..len)?;

        Some((writable, data))
    }

    fn insert(&mut self, address: u64, data: &[u8], writable: bool) {
        if let Some((is_writable, v)) = self.get(address, data.len()) {
            assert_eq!(writable, is_writable);
            v.copy_from_slice(data);
            return;
        }

        let end = address + data.len() as u64;
        let mut data = data.to_vec();
        if let Some((&next, &(next_writable, ref next_data))) = self.map.range(address..).next() {
            if end > next {
                let next_end = next + next_data.len() as u64;
                panic!(
                    "overlapping memory map: {address:#x}..{end:#x} > {next:#x}..={next_end:#x}"
                );
            }
            if end == next && next_writable == writable {
                data.extend(next_data.as_slice());
                self.map.remove(&next).unwrap();
            }
        }

        if let Some((&prev, &mut (prev_writable, ref mut prev_data))) =
            self.map.range_mut(..address).last()
        {
            let prev_end = prev + prev_data.len() as u64;
            if prev_end > address {
                panic!(
                    "overlapping memory map: {prev:#x}..{prev_end:#x} > {address:#x}..={end:#x}"
                );
            }
            if prev_end == address && prev_writable == writable {
                prev_data.extend_from_slice(&data);
                return;
            }
        }

        self.map.insert(address, (writable, data));
    }
}

impl VirtioTestMemoryAccess {
    fn new() -> Arc<Self> {
        Default::default()
    }

    fn modify_memory_map(&self, address: u64, data: &[u8], writeable: bool) {
        self.memory_map.lock().insert(address, data, writeable);
    }

    fn memory_map_get_u16(&self, address: u64) -> u16 {
        let mut map = self.memory_map.lock();
        let (_, data) = map.get(address, 2).unwrap();
        u16::from_le_bytes(data.try_into().unwrap())
    }

    fn memory_map_get_u32(&self, address: u64) -> u32 {
        let mut map = self.memory_map.lock();
        let (_, data) = map.get(address, 4).unwrap();
        u32::from_le_bytes(data.try_into().unwrap())
    }
}

// SAFETY: test code
unsafe impl GuestMemoryAccess for VirtioTestMemoryAccess {
    fn mapping(&self) -> Option<NonNull<u8>> {
        None
    }

    fn max_address(&self) -> u64 {
        // No real bound, so use the max physical address width on
        // AMD64/ARM64.
        1 << 52
    }

    unsafe fn read_fallback(
        &self,
        address: u64,
        dest: *mut u8,
        len: usize,
    ) -> Result<(), GuestMemoryBackingError> {
        match self.memory_map.lock().get(address, len) {
            Some((_, value)) => {
                // SAFETY: guaranteed by caller
                unsafe {
                    std::ptr::copy(value.as_ptr(), dest, len);
                }
            }
            None => panic!(
                "Unexpected read request for {} bytes at address {:x}",
                len, address
            ),
        }
        Ok(())
    }

    unsafe fn write_fallback(
        &self,
        address: u64,
        src: *const u8,
        len: usize,
    ) -> Result<(), GuestMemoryBackingError> {
        match self.memory_map.lock().get(address, len) {
            Some((true, value)) => {
                // SAFETY: guaranteed by caller
                unsafe {
                    std::ptr::copy(src, value.as_mut_ptr(), len);
                }
            }
            _ => panic!(
                "Unexpected write request for {} bytes at address {:x}",
                len, address
            ),
        }
        Ok(())
    }

    fn fill_fallback(
        &self,
        address: u64,
        val: u8,
        len: usize,
    ) -> Result<(), GuestMemoryBackingError> {
        match self.memory_map.lock().get(address, len) {
            Some((true, value)) => value.fill(val),
            _ => panic!("Unexpected write request at address {:x}", address),
        };
        Ok(())
    }
}

struct DoorbellEntry;
impl Drop for DoorbellEntry {
    fn drop(&mut self) {}
}

impl DoorbellRegistration for VirtioTestMemoryAccess {
    fn register_doorbell(
        &self,
        _: u64,
        _: Option<u64>,
        _: Option<u32>,
        _: &Event,
    ) -> io::Result<Box<dyn Send + Sync>> {
        Ok(Box::new(DoorbellEntry))
    }
}

type VirtioTestWorkCallback = Box<dyn Fn(VirtioQueueCallbackWork) + Sync + Send>;
struct CreateDirectQueueParams {
    process_work: VirtioTestWorkCallback,
    notify: Interrupt,
    event: Event,
}

struct SplitQueue {
    last_avail_index: Vec<u16>,
    last_used_index: Vec<u16>,
}

struct PackedQueue {
    next_ready_index: Vec<u16>,
    ready_wrapped_bit: Vec<bool>,
    next_completed_index: Vec<u16>,
    completed_wrapped_bit: Vec<bool>,
    buffer_and_count: Vec<BTreeMap<u16, u16>>,
}

enum VirtioQueueInfo {
    Split(SplitQueue),
    Packed(PackedQueue),
}

struct VirtioTestGuest {
    test_mem: Arc<VirtioTestMemoryAccess>,
    driver: DefaultDriver,
    num_queues: u16,
    queue_size: u16,
    allow_indirect_descriptors: bool,
    info: VirtioQueueInfo,
    use_ring_event_index: bool,
    avail_descriptors: Vec<Vec<bool>>,
    auto_arm_event: bool,
}

impl VirtioTestGuest {
    fn new_split(
        driver: &DefaultDriver,
        test_mem: &Arc<VirtioTestMemoryAccess>,
        num_queues: u16,
        queue_size: u16,
        use_ring_event_index: bool,
    ) -> Self {
        let last_avail_index: Vec<u16> = vec![0; num_queues as usize];
        let last_used_index: Vec<u16> = vec![0; num_queues as usize];
        let avail_descriptors: Vec<Vec<bool>> =
            vec![vec![true; queue_size as usize]; num_queues as usize];
        let test_guest = Self {
            test_mem: test_mem.clone(),
            driver: driver.clone(),
            num_queues,
            queue_size,
            allow_indirect_descriptors: true,
            info: VirtioQueueInfo::Split(SplitQueue {
                last_avail_index,
                last_used_index,
            }),
            use_ring_event_index,
            avail_descriptors,
            auto_arm_event: true,
        };
        for i in 0..num_queues {
            test_guest.add_queue_memory(i);
        }
        test_guest
    }

    fn new_packed(
        driver: &DefaultDriver,
        test_mem: &Arc<VirtioTestMemoryAccess>,
        num_queues: u16,
        queue_size: u16,
        use_ring_event_index: bool,
    ) -> Self {
        let avail_descriptors: Vec<Vec<bool>> =
            vec![vec![true; queue_size as usize]; num_queues as usize];
        let test_guest = Self {
            test_mem: test_mem.clone(),
            driver: driver.clone(),
            num_queues,
            queue_size,
            allow_indirect_descriptors: true,
            info: VirtioQueueInfo::Packed(PackedQueue {
                next_ready_index: vec![0; num_queues as usize],
                ready_wrapped_bit: vec![true; num_queues as usize],
                next_completed_index: vec![0; num_queues as usize],
                completed_wrapped_bit: vec![true; num_queues as usize],
                buffer_and_count: vec![BTreeMap::new(); num_queues as usize],
            }),
            use_ring_event_index,
            avail_descriptors,
            auto_arm_event: true,
        };
        for i in 0..num_queues {
            test_guest.add_queue_memory(i);
        }
        test_guest
    }

    fn driver(&self) -> DefaultDriver {
        self.driver.clone()
    }

    fn mem(&self) -> GuestMemory {
        GuestMemory::new("test", self.test_mem.clone())
    }

    fn create_direct_queues<F>(&self, f: F) -> Vec<TaskControl<TestQueueWorker, TestQueueState>>
    where
        F: Fn(u16) -> CreateDirectQueueParams,
    {
        (0..self.num_queues)
            .map(|i| {
                let params = f(i);
                let queue_event = PolledWait::new(&self.driver, params.event).unwrap();
                let queue = VirtioQueue::new(
                    self.queue_features(),
                    self.queue_params(i),
                    self.mem(),
                    params.notify,
                    queue_event,
                )
                .expect("failed to create virtio queue");
                let mut tc = TaskControl::new(TestQueueWorker {
                    callback: params.process_work,
                });
                tc.insert(
                    self.driver.clone(),
                    "virtio-test-queue",
                    TestQueueState { queue },
                );
                tc.start();
                tc
            })
            .collect::<Vec<_>>()
    }

    fn queue_features(&self) -> VirtioDeviceFeatures {
        let flags0 = if self.use_ring_event_index {
            VIRTIO_F_RING_EVENT_IDX
        } else {
            0
        };
        let flags0 = if self.allow_indirect_descriptors {
            flags0 | VIRTIO_F_RING_INDIRECT_DESC
        } else {
            flags0
        };

        let flags1 = VIRTIO_F_VERSION_1;
        let flags1 = if matches!(self.info, VirtioQueueInfo::Packed(_)) {
            flags1 | VIRTIO_F_RING_PACKED
        } else {
            flags1
        };

        VirtioDeviceFeatures::new()
            .with_bank(0, flags0)
            .with_bank(1, flags1)
    }

    fn queue_params(&self, i: u16) -> QueueParams {
        QueueParams {
            size: self.queue_size,
            enable: true,
            desc_addr: self.get_queue_descriptor_base_address(i),
            avail_addr: self.get_queue_available_base_address(i),
            used_addr: self.get_queue_used_base_address(i),
        }
    }

    fn get_queue_base_address(&self, index: u16) -> u64 {
        0x10000000 * index as u64
    }

    fn get_queue_descriptor_base_address(&self, index: u16) -> u64 {
        self.get_queue_base_address(index) + 0x1000
    }

    fn get_queue_available_base_address(&self, index: u16) -> u64 {
        self.get_queue_base_address(index) + 0x2000
    }

    fn get_queue_used_base_address(&self, index: u16) -> u64 {
        self.get_queue_base_address(index) + 0x3000
    }

    fn get_queue_descriptor_backing_memory_address(&self, index: u16) -> u64 {
        self.get_queue_base_address(index) + 0x4000
    }

    fn setup_chipset_device(
        &self,
        dev: &mut VirtioMmioDevice,
        driver_features: VirtioDeviceFeatures,
    ) {
        dev.write_u32(112, VIRTIO_ACKNOWLEDGE);
        dev.write_u32(112, VIRTIO_DRIVER);
        dev.write_u32(36, 0);
        dev.write_u32(32, driver_features.bank(0));
        dev.write_u32(36, 1);
        dev.write_u32(32, driver_features.bank(1));
        dev.write_u32(112, VIRTIO_FEATURES_OK);
        for i in 0..self.num_queues {
            let queue_index = i;
            dev.write_u32(48, i as u32);
            dev.write_u32(56, self.queue_size as u32);
            let desc_addr = self.get_queue_descriptor_base_address(queue_index);
            dev.write_u32(128, desc_addr as u32);
            dev.write_u32(132, (desc_addr >> 32) as u32);
            let avail_addr = self.get_queue_available_base_address(queue_index);
            dev.write_u32(144, avail_addr as u32);
            dev.write_u32(148, (avail_addr >> 32) as u32);
            let used_addr = self.get_queue_used_base_address(queue_index);
            dev.write_u32(160, used_addr as u32);
            dev.write_u32(164, (used_addr >> 32) as u32);
            // enable the queue
            dev.write_u32(68, 1);
        }
        dev.write_u32(112, VIRTIO_DRIVER_OK);
        assert_eq!(dev.read_u32(0xfc), 2);
    }

    fn setup_pci_device(
        &self,
        dev: &mut VirtioPciTestDevice,
        driver_features: VirtioDeviceFeatures,
    ) {
        let bar_address1: u64 = 0x10000000000;
        dev.pci_device
            .pci_cfg_write(0x14, (bar_address1 >> 32) as u32)
            .unwrap();
        dev.pci_device
            .pci_cfg_write(0x10, bar_address1 as u32)
            .unwrap();

        let bar_address2: u64 = 0x20000000000;
        dev.pci_device
            .pci_cfg_write(0x1c, (bar_address2 >> 32) as u32)
            .unwrap();
        dev.pci_device
            .pci_cfg_write(0x18, bar_address2 as u32)
            .unwrap();

        dev.pci_device
            .pci_cfg_write(
                0x4,
                cfg_space::Command::new()
                    .with_mmio_enabled(true)
                    .into_bits() as u32,
            )
            .unwrap();

        let mut device_status = VIRTIO_ACKNOWLEDGE as u8;
        dev.pci_device
            .mmio_write(bar_address1 + 20, &device_status.to_le_bytes())
            .unwrap();
        device_status = VIRTIO_DRIVER as u8;
        dev.pci_device
            .mmio_write(bar_address1 + 20, &device_status.to_le_bytes())
            .unwrap();
        dev.write_u32(bar_address1 + 8, 0);
        dev.write_u32(bar_address1 + 12, driver_features.bank(0));
        dev.write_u32(bar_address1 + 8, 1);
        dev.write_u32(bar_address1 + 12, driver_features.bank(1));
        device_status = VIRTIO_FEATURES_OK as u8;
        dev.pci_device
            .mmio_write(bar_address1 + 20, &device_status.to_le_bytes())
            .unwrap();
        // setup config interrupt
        dev.pci_device
            .mmio_write(bar_address2, &0_u64.to_le_bytes())
            .unwrap(); // vector
        dev.pci_device
            .mmio_write(bar_address2 + 8, &0_u32.to_le_bytes())
            .unwrap(); // data
        dev.pci_device
            .mmio_write(bar_address2 + 12, &0_u32.to_le_bytes())
            .unwrap();
        for i in 0..self.num_queues {
            let queue_index = i;
            dev.pci_device
                .mmio_write(bar_address1 + 22, &queue_index.to_le_bytes())
                .unwrap();
            dev.pci_device
                .mmio_write(bar_address1 + 24, &self.queue_size.to_le_bytes())
                .unwrap();
            // setup MSI information for the queue
            let msix_vector = queue_index + 1;
            let address = bar_address2 + 0x10 * msix_vector as u64;
            dev.pci_device
                .mmio_write(address, &(msix_vector as u64).to_le_bytes())
                .unwrap();
            let address = bar_address2 + 0x10 * msix_vector as u64 + 8;
            dev.pci_device
                .mmio_write(address, &0_u32.to_le_bytes())
                .unwrap();
            let address = bar_address2 + 0x10 * msix_vector as u64 + 12;
            dev.pci_device
                .mmio_write(address, &0_u32.to_le_bytes())
                .unwrap();
            dev.pci_device
                .mmio_write(bar_address1 + 26, &msix_vector.to_le_bytes())
                .unwrap();
            // setup queue addresses
            let desc_addr = self.get_queue_descriptor_base_address(queue_index);
            dev.write_u32(bar_address1 + 32, desc_addr as u32);
            dev.write_u32(bar_address1 + 36, (desc_addr >> 32) as u32);
            let avail_addr = self.get_queue_available_base_address(queue_index);
            dev.write_u32(bar_address1 + 40, avail_addr as u32);
            dev.write_u32(bar_address1 + 44, (avail_addr >> 32) as u32);
            let used_addr = self.get_queue_used_base_address(queue_index);
            dev.write_u32(bar_address1 + 48, used_addr as u32);
            dev.write_u32(bar_address1 + 52, (used_addr >> 32) as u32);
            // enable the queue
            let enabled: u16 = 1;
            dev.pci_device
                .mmio_write(bar_address1 + 28, &enabled.to_le_bytes())
                .unwrap();
        }
        // enable all device MSI interrupts
        dev.pci_device.pci_cfg_write(0x40, 0x80000000).unwrap();
        // run device
        device_status = VIRTIO_DRIVER_OK as u8;
        dev.pci_device
            .mmio_write(bar_address1 + 20, &device_status.to_le_bytes())
            .unwrap();
        let mut config_generation: [u8; 1] = [0];
        dev.pci_device
            .mmio_read(bar_address1 + 21, &mut config_generation)
            .unwrap();
        assert_eq!(config_generation[0], 2);
    }

    fn get_queue_descriptor(&self, queue_index: u16, descriptor_index: u16) -> u64 {
        self.get_queue_descriptor_base_address(queue_index) + 0x10 * descriptor_index as u64
    }

    fn add_queue_memory(&self, queue_index: u16) {
        let is_packed_descriptors = matches!(self.info, VirtioQueueInfo::Packed(_));
        // descriptors
        for i in 0..self.queue_size {
            let base = self.get_queue_descriptor(queue_index, i);
            // physical address
            self.test_mem.modify_memory_map(
                base,
                &(self.get_queue_descriptor_backing_memory_address(queue_index)
                    + 0x1000 * i as u64)
                    .to_le_bytes(),
                is_packed_descriptors,
            );
            // length
            self.test_mem.modify_memory_map(
                base + 8,
                &0x1000u32.to_le_bytes(),
                is_packed_descriptors,
            );
            // split: flags, packed: buffer_id
            self.test_mem
                .modify_memory_map(base + 12, &0u16.to_le_bytes(), is_packed_descriptors);
            // split: next index, packed: flags
            self.test_mem
                .modify_memory_map(base + 14, &0u16.to_le_bytes(), is_packed_descriptors);
        }

        if is_packed_descriptors {
            let base = self.get_queue_descriptor(queue_index, self.queue_size);
            // Device and driver event fields
            self.test_mem
                .modify_memory_map(base, &0_u32.to_le_bytes(), true);
            self.test_mem
                .modify_memory_map(base + 4, &0_u32.to_le_bytes(), false);
            self.test_mem.modify_memory_map(
                self.get_queue_available_base_address(queue_index),
                &0_u32.to_le_bytes(),
                false,
            );
            self.test_mem.modify_memory_map(
                self.get_queue_used_base_address(queue_index),
                &0_u32.to_le_bytes(),
                true,
            );
            return;
        }

        // available queue (flags, index)
        let base = self.get_queue_available_base_address(queue_index);
        self.test_mem
            .modify_memory_map(base, &0u16.to_le_bytes(), false);
        self.test_mem
            .modify_memory_map(base + 2, &0u16.to_le_bytes(), false);
        // available queue ring buffer
        for i in 0..self.queue_size {
            let base = base + 4 + 2 * i as u64;
            self.test_mem
                .modify_memory_map(base, &0u16.to_le_bytes(), false);
        }
        // used event
        if self.use_ring_event_index {
            self.test_mem.modify_memory_map(
                base + 4 + 2 * self.queue_size as u64,
                &0u16.to_le_bytes(),
                false,
            );
        }

        // used queue (flags, index)
        let base = self.get_queue_used_base_address(queue_index);
        self.test_mem
            .modify_memory_map(base, &0u16.to_le_bytes(), true);
        self.test_mem
            .modify_memory_map(base + 2, &0u16.to_le_bytes(), true);
        for i in 0..self.queue_size {
            let base = base + 4 + 8 * i as u64;
            // index
            self.test_mem
                .modify_memory_map(base, &0u32.to_le_bytes(), true);
            // length
            self.test_mem
                .modify_memory_map(base + 4, &0u32.to_le_bytes(), true);
        }
        // available event
        if self.use_ring_event_index {
            self.test_mem.modify_memory_map(
                base + 4 + 8 * self.queue_size as u64,
                &0u16.to_le_bytes(),
                true,
            );
        }
    }

    fn reserve_split_descriptor(&mut self, queue_index: u16) -> u16 {
        if let Some((desc_index, desc)) = self.avail_descriptors[queue_index as usize]
            .iter_mut()
            .enumerate()
            .find(|(_, desc)| **desc)
        {
            *desc = false;
            return desc_index as u16;
        }
        panic!("No descriptors are available!");
    }

    fn reserve_packed_descriptors(&mut self, queue_index: u16, count: u16) -> u16 {
        let queue_size = self.queue_size;
        let starting_descriptor = if let VirtioQueueInfo::Packed(packed) = &self.info {
            packed.next_ready_index[queue_index as usize]
        } else {
            panic!("Not a packed queue");
        };
        let avail_descriptors = &mut self.avail_descriptors[queue_index as usize];
        for i in 0..count {
            let desc_index = (starting_descriptor + i) % queue_size;
            if !avail_descriptors[desc_index as usize] {
                panic!("Not enough available descriptors!");
            }
            avail_descriptors[desc_index as usize] = false;
        }
        // Reset default descriptor values.
        for i in 0..count {
            let desc_index = (starting_descriptor + i) % queue_size;
            let base = self.get_queue_descriptor(queue_index, desc_index);
            // physical address
            self.test_mem.modify_memory_map(
                base,
                &(self.get_queue_descriptor_backing_memory_address(queue_index)
                    + 0x1000 * desc_index as u64)
                    .to_le_bytes(),
                true,
            );
            // length
            self.test_mem
                .modify_memory_map(base + 8, &0x1000u32.to_le_bytes(), true);
            // buffer_id
            self.test_mem
                .modify_memory_map(base + 12, &0u16.to_le_bytes(), true);
            // flags
            self.test_mem
                .modify_memory_map(base + 14, &0u16.to_le_bytes(), true);
        }
        starting_descriptor
    }

    fn free_descriptor(&mut self, queue_index: u16, desc_index: u16) {
        assert!(desc_index < self.queue_size);
        let desc_addr = self.get_queue_descriptor(queue_index, desc_index);
        if matches!(self.info, VirtioQueueInfo::Split(_)) {
            let flags: DescriptorFlags = self.test_mem.memory_map_get_u16(desc_addr + 12).into();
            if flags.next() {
                let next = self.test_mem.memory_map_get_u16(desc_addr + 14);
                self.free_descriptor(queue_index, next);
            }
        } else {
            // Linked packed descriptors are handled in get_next_packed_completed()
        }
        let avail_descriptors = &mut self.avail_descriptors[queue_index as usize];
        assert_eq!(avail_descriptors[desc_index as usize], false);
        avail_descriptors[desc_index as usize] = true;
    }

    fn queue_available_desc(&mut self, queue_index: u16, desc_index: u16) {
        let avail_base_addr = self.get_queue_available_base_address(queue_index);
        let (last_avail_index, next_index) = if let VirtioQueueInfo::Split(split) = &mut self.info {
            let last_avail_index = split.last_avail_index[queue_index as usize];
            let next_index = last_avail_index % self.queue_size;
            let last_avail_index = last_avail_index.wrapping_add(1);
            split.last_avail_index[queue_index as usize] = last_avail_index;
            (last_avail_index, next_index)
        } else {
            panic!("Not a split queue");
        };
        self.test_mem.modify_memory_map(
            avail_base_addr + 4 + 2 * next_index as u64,
            &desc_index.to_le_bytes(),
            false,
        );
        self.test_mem.modify_memory_map(
            avail_base_addr + 2,
            &last_avail_index.to_le_bytes(),
            false,
        );
    }

    fn add_to_avail_queue(&mut self, queue_index: u16) {
        if matches!(self.info, VirtioQueueInfo::Packed(_)) {
            self.make_packed_descriptors_available(queue_index, vec![DescriptorFlags::new()], None);
        } else {
            let next_descriptor = self.reserve_split_descriptor(queue_index);
            // flags
            self.test_mem.modify_memory_map(
                self.get_queue_descriptor(queue_index, next_descriptor) + 12,
                &0u16.to_le_bytes(),
                false,
            );
            self.queue_available_desc(queue_index, next_descriptor);
        }
    }

    fn make_packed_descriptors_available(
        &mut self,
        queue_index: u16,
        flags: Vec<DescriptorFlags>,
        indirect_count: Option<u32>,
    ) {
        let count = flags.len() as u16;
        let first_descriptor = self.reserve_packed_descriptors(queue_index, count);
        let wrapped_bit = if let VirtioQueueInfo::Packed(packed) = &self.info {
            packed.ready_wrapped_bit[queue_index as usize]
        } else {
            unreachable!("Not a packed queue");
        };
        let buffer_id = queue_index << 8 | first_descriptor;
        // The first descriptor needs to be the last marked available, as once
        // it is valid, all the descriptors can be examined at any time.
        for desc_index in (0..count).rev() {
            let descriptor_index = (first_descriptor + desc_index) % self.queue_size;
            let wrapped_bit = if descriptor_index >= first_descriptor {
                wrapped_bit
            } else {
                !wrapped_bit
            };
            let desc_addr = self.get_queue_descriptor(queue_index, descriptor_index);
            let flags = flags[desc_index as usize]
                .with_available(wrapped_bit)
                .with_used(!wrapped_bit)
                .with_next(desc_index < count - 1);
            self.test_mem
                .modify_memory_map(desc_addr + 14, &flags.into_bits().to_le_bytes(), true);
            if desc_index == count - 1 {
                // The buffer ID is assigned to the last descriptor.
                self.test_mem
                    .modify_memory_map(desc_addr + 12, &buffer_id.to_le_bytes(), true);
                // Indirect packed descriptors consume all of the indirect
                // buffer, so set the size accordingly.
                if let Some(indirect_count) = indirect_count {
                    let length = indirect_count * 0x10;
                    self.test_mem
                        .modify_memory_map(desc_addr + 8, &length.to_le_bytes(), true);
                }
            }
        }
        if let VirtioQueueInfo::Packed(packed) = &mut self.info {
            packed.buffer_and_count[queue_index as usize].insert(buffer_id, count);
            let next_ready_index = (first_descriptor + count) % self.queue_size;
            if next_ready_index < packed.next_ready_index[queue_index as usize] {
                packed.ready_wrapped_bit[queue_index as usize] =
                    !packed.ready_wrapped_bit[queue_index as usize];
            }
            packed.next_ready_index[queue_index as usize] = next_ready_index;
        } else {
            unreachable!("Not a packed queue");
        };
    }

    fn add_indirect_to_avail_queue(&mut self, queue_index: u16) {
        // create another (indirect) descriptor in the buffer
        let buffer_addr = self.get_queue_descriptor_backing_memory_address(queue_index);
        // physical address
        self.test_mem
            .modify_memory_map(buffer_addr, &0xffffffff00000000u64.to_le_bytes(), false);
        // length
        self.test_mem
            .modify_memory_map(buffer_addr + 8, &0x1000u32.to_le_bytes(), false);
        // split: flags, packed: buffer_id (ignored for indirect descriptors)
        self.test_mem
            .modify_memory_map(buffer_addr + 12, &0u16.to_le_bytes(), false);
        // split: next, packed: flags (wait bit ignored)
        self.test_mem
            .modify_memory_map(buffer_addr + 14, &0u16.to_le_bytes(), false);
        if matches!(self.info, VirtioQueueInfo::Packed(_)) {
            self.make_packed_descriptors_available(
                queue_index,
                vec![DescriptorFlags::new().with_indirect(true)],
                Some(1),
            );
        } else {
            let next_descriptor = self.reserve_split_descriptor(queue_index);
            // flags on primary descriptor
            self.test_mem.modify_memory_map(
                self.get_queue_descriptor(queue_index, next_descriptor) + 12,
                &u16::from(DescriptorFlags::new().with_indirect(true)).to_le_bytes(),
                false,
            );
            self.queue_available_desc(queue_index, next_descriptor);
        }
    }

    fn add_linked_to_avail_queue(&mut self, queue_index: u16, desc_count: u16) {
        if matches!(self.info, VirtioQueueInfo::Packed(_)) {
            let flags = vec![DescriptorFlags::new(); desc_count as usize];
            return self.make_packed_descriptors_available(queue_index, flags, None);
        }

        let mut descriptors = Vec::with_capacity(desc_count as usize);
        for _ in 0..desc_count {
            descriptors.push(self.reserve_split_descriptor(queue_index));
        }

        for i in 0..descriptors.len() {
            let base = self.get_queue_descriptor(queue_index, descriptors[i]);
            let flags = if i < descriptors.len() - 1 {
                u16::from(DescriptorFlags::new().with_next(true))
            } else {
                0
            };
            self.test_mem
                .modify_memory_map(base + 12, &flags.to_le_bytes(), false);
            let next = if i < descriptors.len() - 1 {
                descriptors[i + 1]
            } else {
                0
            };
            self.test_mem
                .modify_memory_map(base + 14, &next.to_le_bytes(), false);
        }
        self.queue_available_desc(queue_index, descriptors[0]);
    }

    fn add_indirect_linked_to_avail_queue(&mut self, queue_index: u16, desc_count: u16) {
        // create indirect descriptors in the buffer
        let buffer_addr = self.get_queue_descriptor_backing_memory_address(queue_index);
        for i in 0..desc_count {
            let base = buffer_addr + 0x10 * i as u64;
            let indirect_buffer_addr = 0xffffffff00000000u64 + 0x1000 * i as u64;
            // physical address
            self.test_mem
                .modify_memory_map(base, &indirect_buffer_addr.to_le_bytes(), false);
            // length
            self.test_mem
                .modify_memory_map(base + 8, &0x1000u32.to_le_bytes(), false);
            // The next field is ignored for packed indirect descriptors. It
            // will consume all descriptors in the buffer, so is dependent on
            // the size of the buffer.
            let flags = if matches!(self.info, VirtioQueueInfo::Split(_)) && i < desc_count - 1 {
                DescriptorFlags::new().with_next(true).into_bits()
            } else {
                0
            };
            if matches!(self.info, VirtioQueueInfo::Packed(_)) {
                // buffer id (ignored for indirect descriptors)
                self.test_mem
                    .modify_memory_map(base + 12, &0_u16.to_le_bytes(), false);
                // flags
                self.test_mem
                    .modify_memory_map(base + 14, &flags.to_le_bytes(), false);
            } else {
                // flags
                self.test_mem
                    .modify_memory_map(base + 12, &flags.to_le_bytes(), false);
                // next index
                let next = if i < desc_count - 1 { i + 1 } else { 0 };
                self.test_mem
                    .modify_memory_map(base + 14, &next.to_le_bytes(), false);
            }
        }

        if matches!(self.info, VirtioQueueInfo::Packed(_)) {
            self.make_packed_descriptors_available(
                queue_index,
                vec![DescriptorFlags::new().with_indirect(true)],
                Some(desc_count as u32),
            );
        } else {
            let next_descriptor = self.reserve_split_descriptor(queue_index);
            // flags on primary descriptor
            self.test_mem.modify_memory_map(
                self.get_queue_descriptor(queue_index, next_descriptor) + 12,
                &DescriptorFlags::new()
                    .with_indirect(true)
                    .into_bits()
                    .to_le_bytes(),
                false,
            );
            self.queue_available_desc(queue_index, next_descriptor);
        }
    }

    fn get_next_completed(&mut self, queue_index: u16) -> Option<(u16, u32)> {
        if matches!(self.info, VirtioQueueInfo::Split(_)) {
            self.get_next_split_completed(queue_index)
        } else {
            self.get_next_packed_completed(queue_index)
        }
    }

    fn get_next_split_completed(&mut self, queue_index: u16) -> Option<(u16, u32)> {
        let avail_base_addr = self.get_queue_available_base_address(queue_index);
        let used_base_addr = self.get_queue_used_base_address(queue_index);
        let cur_used_index = self.test_mem.memory_map_get_u16(used_base_addr + 2);
        let next_index = if let VirtioQueueInfo::Split(split) = &mut self.info {
            let last_used_index = split.last_used_index[queue_index as usize];
            if last_used_index == cur_used_index {
                return None;
            }
            split.last_used_index[queue_index as usize] = last_used_index.wrapping_add(1);
            last_used_index % self.queue_size
        } else {
            panic!("Not a split queue");
        };

        if self.auto_arm_event && self.use_ring_event_index {
            self.test_mem.modify_memory_map(
                avail_base_addr + 4 + 2 * self.queue_size as u64,
                &cur_used_index.to_le_bytes(),
                false,
            );
        }

        let desc_index = self
            .test_mem
            .memory_map_get_u32(used_base_addr + 4 + 8 * next_index as u64);
        let desc_index = desc_index as u16;
        let bytes_written = self
            .test_mem
            .memory_map_get_u32(used_base_addr + 8 + 8 * next_index as u64);
        self.free_descriptor(queue_index, desc_index);
        Some((desc_index, bytes_written))
    }

    fn get_next_packed_completed(&mut self, queue_index: u16) -> Option<(u16, u32)> {
        let queue_size = self.queue_size;
        let (descriptor_index, wrapped_bit) =
            if let VirtioQueueInfo::Packed(packed) = &mut self.info {
                (
                    packed.next_completed_index[queue_index as usize],
                    packed.completed_wrapped_bit[queue_index as usize],
                )
            } else {
                panic!("Not a packed queue");
            };
        let desc_addr = self.get_queue_descriptor(queue_index, descriptor_index);
        let flags: DescriptorFlags = self.test_mem.memory_map_get_u16(desc_addr + 14).into();
        if flags.available() != wrapped_bit || flags.used() != wrapped_bit {
            return None;
        }

        let bytes_written = self.test_mem.memory_map_get_u32(desc_addr + 8);
        let buffer_id = self.test_mem.memory_map_get_u16(desc_addr + 12);
        let count = if let VirtioQueueInfo::Packed(packed) = &mut self.info {
            let err_msg = format!("Buffer ID {} not found in queue {}", buffer_id, queue_index);
            let count = packed
                .buffer_and_count
                .get_mut(queue_index as usize)
                .expect("Invalid queue index")
                .remove(&buffer_id)
                .expect(err_msg.as_str());
            let next_completed_index = (descriptor_index + count) % queue_size;
            if next_completed_index < packed.next_completed_index[queue_index as usize] {
                packed.completed_wrapped_bit[queue_index as usize] =
                    !packed.completed_wrapped_bit[queue_index as usize];
            }
            packed.next_completed_index[queue_index as usize] = next_completed_index;
            count
        } else {
            unreachable!("Not a packed queue");
        };

        for desc_index in 0..count {
            let descriptor_index = (descriptor_index + desc_index) % queue_size;
            self.free_descriptor(queue_index, descriptor_index);
        }
        Some((descriptor_index, bytes_written))
    }

    fn enable_interrupt(&mut self, queue_index: u16, desc_index: Option<u16>) {
        assert!(desc_index.is_none() || self.use_ring_event_index);
        let base = self.get_queue_available_base_address(queue_index);
        if let VirtioQueueInfo::Packed(packed) = &mut self.info {
            if let Some(desc_index) = desc_index {
                let wrapped_bit = if packed.next_ready_index[queue_index as usize] > desc_index {
                    if packed.ready_wrapped_bit[queue_index as usize] {
                        0
                    } else {
                        1
                    }
                } else {
                    if packed.ready_wrapped_bit[queue_index as usize] {
                        1
                    } else {
                        0
                    }
                };
                let packed_event = desc_index as u32 | (wrapped_bit << 15) | (2_u32 << 16);
                self.test_mem
                    .modify_memory_map(base, &packed_event.to_le_bytes(), false);
                self.auto_arm_event = false;
            } else {
                self.test_mem
                    .modify_memory_map(base, &0_u32.to_le_bytes(), false);
                self.auto_arm_event = true;
            }
        } else {
            if let Some(desc_index) = desc_index {
                self.test_mem.modify_memory_map(
                    base + 4 + 2 * self.queue_size as u64,
                    &desc_index.to_le_bytes(),
                    false,
                );
                self.auto_arm_event = false;
            } else {
                if self.use_ring_event_index {
                    let next_index = if let VirtioQueueInfo::Split(split) = &self.info {
                        split.last_used_index[queue_index as usize]
                    } else {
                        panic!("Not a split queue");
                    };
                    self.test_mem.modify_memory_map(
                        base + 4 + 2 * self.queue_size as u64,
                        &next_index.to_le_bytes(),
                        false,
                    );
                } else {
                    self.test_mem
                        .modify_memory_map(base, &0_u16.to_le_bytes(), false);
                }
                self.auto_arm_event = true;
            }
        }
    }

    fn disable_interrupt(&mut self, queue_index: u16) {
        let base = self.get_queue_available_base_address(queue_index);
        if matches!(self.info, VirtioQueueInfo::Packed(_)) {
            self.test_mem
                .modify_memory_map(base, &(1_u32 << 16).to_le_bytes(), false);
        } else {
            if self.use_ring_event_index {
                // Can't really disable, but can set the next event to be far away.
                let last_index = if let VirtioQueueInfo::Split(split) = &self.info {
                    split.last_used_index[queue_index as usize] - 1
                } else {
                    panic!("Not a split queue");
                };
                self.test_mem.modify_memory_map(
                    base + 4 + 2 * self.queue_size as u64,
                    &last_index.to_le_bytes(),
                    false,
                );
            } else {
                self.test_mem
                    .modify_memory_map(base, &1_u16.to_le_bytes(), false);
            }
        }
        self.auto_arm_event = false;
    }
}

struct TestQueueWorker {
    callback: VirtioTestWorkCallback,
}

struct TestQueueState {
    queue: VirtioQueue,
}

impl AsyncRun<TestQueueState> for TestQueueWorker {
    async fn run(
        &mut self,
        stop: &mut StopTask<'_>,
        state: &mut TestQueueState,
    ) -> Result<(), Cancelled> {
        loop {
            let work = stop.until_stopped(state.queue.next()).await?;
            let Some(work) = work else { break };
            match work {
                Ok(work) => (self.callback)(work),
                Err(err) => panic!("queue error: {}", err),
            }
        }
        Ok(())
    }
}
struct VirtioPciTestDevice {
    pci_device: VirtioPciDevice,
    test_intc: Arc<TestPciInterruptController>,
}

type TestDeviceQueueWorkFn = Arc<dyn Fn(u16, VirtioQueueCallbackWork) + Send + Sync>;

/// A minimal VirtioDevice whose enable() always returns an error.
/// Used to test that transports correctly handle enable failures.
#[derive(InspectMut)]
#[inspect(skip)]
struct FailingTestDevice {
    traits: DeviceTraits,
}

impl VirtioDevice for FailingTestDevice {
    fn traits(&self) -> DeviceTraits {
        self.traits.clone()
    }

    fn read_registers_u32(&mut self, _offset: u16) -> u32 {
        0
    }

    fn write_registers_u32(&mut self, _offset: u16, _val: u32) {}

    fn enable(&mut self, _resources: Resources) -> anyhow::Result<()> {
        anyhow::bail!("intentional enable failure for testing")
    }

    fn poll_disable(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<()> {
        std::task::Poll::Ready(())
    }
}

#[derive(InspectMut)]
#[inspect(skip)]
struct TestDevice {
    traits: DeviceTraits,
    queue_work: Option<TestDeviceQueueWorkFn>,
    driver: vmcore::vm_task::VmTaskDriver,
    mem: GuestMemory,
    workers: Vec<TaskControl<TestDeviceTask, TestDeviceQueue>>,
}

impl TestDevice {
    fn new(
        driver_source: &VmTaskDriverSource,
        traits: DeviceTraits,
        queue_work: Option<TestDeviceQueueWorkFn>,
        mem: &GuestMemory,
    ) -> Self {
        Self {
            traits,
            queue_work,
            driver: driver_source.simple(),
            mem: mem.clone(),
            workers: Vec::new(),
        }
    }
}

impl VirtioDevice for TestDevice {
    fn traits(&self) -> DeviceTraits {
        self.traits.clone()
    }

    fn read_registers_u32(&mut self, _offset: u16) -> u32 {
        0
    }

    fn write_registers_u32(&mut self, _offset: u16, _val: u32) {}

    fn enable(&mut self, resources: Resources) -> anyhow::Result<()> {
        self.workers = resources
            .queues
            .into_iter()
            .enumerate()
            .filter_map(|(i, queue_resources)| {
                if !queue_resources.params.enable {
                    return None;
                }

                let mut tc = TaskControl::new(TestDeviceTask {
                    index: i as u16,
                    queue_work: self.queue_work.clone(),
                });

                let queue_event = PolledWait::new(&self.driver, queue_resources.event).unwrap();
                let queue = VirtioQueue::new(
                    resources.features.clone(),
                    queue_resources.params,
                    self.mem.clone(),
                    queue_resources.notify,
                    queue_event,
                )
                .expect("failed to create virtio queue");

                tc.insert(
                    self.driver.clone(),
                    "virtio-test-queue",
                    TestDeviceQueue { queue },
                );
                tc.start();
                Some(tc)
            })
            .collect();
        Ok(())
    }

    fn poll_disable(&mut self, cx: &mut std::task::Context<'_>) -> std::task::Poll<()> {
        for worker in &mut self.workers {
            std::task::ready!(worker.poll_stop(cx));
        }
        self.workers.clear();
        std::task::Poll::Ready(())
    }
}

struct TestDeviceTask {
    index: u16,
    queue_work: Option<TestDeviceQueueWorkFn>,
}

struct TestDeviceQueue {
    queue: VirtioQueue,
}

impl AsyncRun<TestDeviceQueue> for TestDeviceTask {
    async fn run(
        &mut self,
        stop: &mut StopTask<'_>,
        state: &mut TestDeviceQueue,
    ) -> Result<(), Cancelled> {
        loop {
            let work = stop.until_stopped(state.queue.next()).await?;
            let Some(work) = work else { break };
            match work {
                Ok(work) => {
                    if let Some(ref func) = self.queue_work {
                        (func)(self.index, work);
                    }
                }
                Err(err) => {
                    panic!(
                        "Invalid virtio queue state index {} error {}",
                        self.index, err
                    );
                }
            }
        }
        Ok(())
    }
}

impl VirtioPciTestDevice {
    fn new(
        driver: &DefaultDriver,
        num_queues: u16,
        test_mem: &Arc<VirtioTestMemoryAccess>,
        queue_work: Option<TestDeviceQueueWorkFn>,
    ) -> Self {
        let doorbell_registration: Arc<dyn DoorbellRegistration> = test_mem.clone();
        let mem = GuestMemory::new("test", test_mem.clone());
        let msi_conn = MsiConnection::new();
        let driver_source = VmTaskDriverSource::new(SingleDriverBackend::new(driver.clone()));

        let dev = VirtioPciDevice::new(
            Box::new(TestDevice::new(
                &driver_source,
                DeviceTraits {
                    device_id: 3,
                    device_features: VirtioDeviceFeatures::new().with_bank(0, 2),
                    max_queues: num_queues,
                    device_register_length: 12,
                    ..Default::default()
                },
                queue_work,
                &mem,
            )),
            PciInterruptModel::Msix(msi_conn.target()),
            Some(doorbell_registration),
            &mut ExternallyManagedMmioIntercepts,
            None,
        )
        .unwrap();

        let test_intc = Arc::new(TestPciInterruptController::new());
        msi_conn.connect(test_intc.signal_msi());

        Self {
            pci_device: dev,
            test_intc,
        }
    }

    fn read_u32(&mut self, address: u64) -> u32 {
        let mut value = [0; 4];
        self.pci_device.mmio_read(address, &mut value).unwrap();
        u32::from_ne_bytes(value)
    }

    fn write_u32(&mut self, address: u64, value: u32) {
        self.pci_device
            .mmio_write(address, &value.to_ne_bytes())
            .unwrap();
    }
}

#[async_test]
async fn verify_chipset_config(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    let doorbell_registration: Arc<dyn DoorbellRegistration> = test_mem.clone();
    let mem = GuestMemory::new("test", test_mem);
    let interrupt = LineInterrupt::detached();
    let driver_source = VmTaskDriverSource::new(SingleDriverBackend::new(driver));

    let mut dev = VirtioMmioDevice::new(
        Box::new(TestDevice::new(
            &driver_source,
            DeviceTraits {
                device_id: 3,
                device_features: VirtioDeviceFeatures::new().with_bank(0, 2),
                max_queues: 1,
                device_register_length: 0,
                ..Default::default()
            },
            None,
            &mem,
        )),
        interrupt,
        Some(doorbell_registration),
        0,
        1,
    );
    // magic value
    assert_eq!(dev.read_u32(0), u32::from_le_bytes(*b"virt"));
    // version
    assert_eq!(dev.read_u32(4), 2);
    // device ID
    assert_eq!(dev.read_u32(8), 3);
    // vendor ID
    assert_eq!(dev.read_u32(12), 0x1af4);
    // device feature (bank 0)
    assert_eq!(
        dev.read_u32(16),
        VIRTIO_F_RING_INDIRECT_DESC | VIRTIO_F_RING_EVENT_IDX | 2
    );
    // device feature bank index
    assert_eq!(dev.read_u32(20), 0);
    // device feature (bank 1)
    dev.write_u32(20, 1);
    assert_eq!(dev.read_u32(20), 1);
    assert_eq!(dev.read_u32(16), VIRTIO_F_VERSION_1 | VIRTIO_F_RING_PACKED);
    // device feature (bank 2)
    dev.write_u32(20, 2);
    assert_eq!(dev.read_u32(16), 0);
    // driver feature (bank 0)
    assert_eq!(dev.read_u32(32), 0);
    dev.write_u32(32, 2);
    assert_eq!(dev.read_u32(32), 2);
    dev.write_u32(32, 0xffffffff);
    assert_eq!(
        dev.read_u32(32),
        VIRTIO_F_RING_INDIRECT_DESC | VIRTIO_F_RING_EVENT_IDX | 2
    );
    // driver feature bank index
    assert_eq!(dev.read_u32(36), 0);
    dev.write_u32(36, 1);
    assert_eq!(dev.read_u32(36), 1);
    // driver feature (bank 1)
    assert_eq!(dev.read_u32(32), 0);
    dev.write_u32(32, 0xffffffff);
    assert_eq!(dev.read_u32(32), VIRTIO_F_VERSION_1 | VIRTIO_F_RING_PACKED);
    // driver feature (bank 2)
    dev.write_u32(36, 2);
    assert_eq!(dev.read_u32(32), 0);
    dev.write_u32(32, 0xffffffff);
    assert_eq!(dev.read_u32(32), 0);
    // host notify
    assert_eq!(dev.read_u32(80), 0);
    // interrupt status
    assert_eq!(dev.read_u32(96), 0);
    // interrupt ACK (queue 0)
    assert_eq!(dev.read_u32(100), 0);
    // device status
    assert_eq!(dev.read_u32(112), 0);
    // config generation
    assert_eq!(dev.read_u32(0xfc), 0);

    // queue index
    assert_eq!(dev.read_u32(48), 0);
    // queue max size (queue 0)
    assert_eq!(dev.read_u32(52), 0x40);
    // queue size (queue 0)
    assert_eq!(dev.read_u32(56), 0x40);
    dev.write_u32(56, 0x20);
    assert_eq!(dev.read_u32(56), 0x20);
    // queue enable (queue 0)
    assert_eq!(dev.read_u32(68), 0);
    dev.write_u32(68, 1);
    assert_eq!(dev.read_u32(68), 1);
    dev.write_u32(68, 0xffffffff);
    assert_eq!(dev.read_u32(68), 1);
    dev.write_u32(68, 0);
    assert_eq!(dev.read_u32(68), 0);
    // queue descriptor address low (queue 0)
    assert_eq!(dev.read_u32(128), 0);
    dev.write_u32(128, 0xffff);
    assert_eq!(dev.read_u32(128), 0xffff);
    // queue descriptor address high (queue 0)
    assert_eq!(dev.read_u32(132), 0);
    dev.write_u32(132, 1);
    assert_eq!(dev.read_u32(132), 1);
    // queue available address low (queue 0)
    assert_eq!(dev.read_u32(144), 0);
    dev.write_u32(144, 0xeeee);
    assert_eq!(dev.read_u32(144), 0xeeee);
    // queue available address high (queue 0)
    assert_eq!(dev.read_u32(148), 0);
    dev.write_u32(148, 2);
    assert_eq!(dev.read_u32(148), 2);
    // queue used address low (queue 0)
    assert_eq!(dev.read_u32(160), 0);
    dev.write_u32(160, 0xdddd);
    assert_eq!(dev.read_u32(160), 0xdddd);
    // queue used address high (queue 0)
    assert_eq!(dev.read_u32(164), 0);
    dev.write_u32(164, 3);
    assert_eq!(dev.read_u32(164), 3);

    // switch to queue #1
    dev.write_u32(48, 1);
    assert_eq!(dev.read_u32(48), 1);
    // queue max size (queue 1)
    assert_eq!(dev.read_u32(52), 0);
    // queue size (queue 1)
    assert_eq!(dev.read_u32(56), 0);
    dev.write_u32(56, 2);
    assert_eq!(dev.read_u32(56), 0);
    // queue enable (queue 1)
    assert_eq!(dev.read_u32(68), 0);
    dev.write_u32(68, 1);
    assert_eq!(dev.read_u32(68), 0);
    // queue descriptor address low (queue 1)
    assert_eq!(dev.read_u32(128), 0);
    dev.write_u32(128, 1);
    assert_eq!(dev.read_u32(128), 0);
    // queue descriptor address high (queue 1)
    assert_eq!(dev.read_u32(132), 0);
    dev.write_u32(132, 1);
    assert_eq!(dev.read_u32(132), 0);
    // queue available address low (queue 1)
    assert_eq!(dev.read_u32(144), 0);
    dev.write_u32(144, 1);
    assert_eq!(dev.read_u32(144), 0);
    // queue available address high (queue 1)
    assert_eq!(dev.read_u32(148), 0);
    dev.write_u32(148, 1);
    assert_eq!(dev.read_u32(148), 0);
    // queue used address low (queue 1)
    assert_eq!(dev.read_u32(160), 0);
    dev.write_u32(160, 1);
    assert_eq!(dev.read_u32(160), 0);
    // queue used address high (queue 1)
    assert_eq!(dev.read_u32(164), 0);
    dev.write_u32(164, 1);
    assert_eq!(dev.read_u32(164), 0);
}

#[async_test]
async fn verify_pci_config(driver: DefaultDriver) {
    let mut pci_test_device =
        VirtioPciTestDevice::new(&driver, 1, &VirtioTestMemoryAccess::new(), None);
    let mut capabilities = 0;
    pci_test_device
        .pci_device
        .pci_cfg_read(4, &mut capabilities)
        .unwrap();
    assert_eq!(
        capabilities,
        (cfg_space::Status::new()
            .with_capabilities_list(true)
            .into_bits() as u32)
            << 16
    );
    let mut next_cap_offset = 0;
    pci_test_device
        .pci_device
        .pci_cfg_read(0x34, &mut next_cap_offset)
        .unwrap();
    assert_ne!(next_cap_offset, 0);

    let mut header = 0;
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16, &mut header)
        .unwrap();
    let header = header.to_le_bytes();
    assert_eq!(header[0], CapabilityId::MSIX.0);
    next_cap_offset = header[1] as u32;
    assert_ne!(next_cap_offset, 0);

    let mut header = 0;
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16, &mut header)
        .unwrap();
    let header = header.to_le_bytes();
    assert_eq!(header[0], CapabilityId::VENDOR_SPECIFIC.0);
    assert_eq!(header[3], VirtioPciCapType::COMMON_CFG.0);
    assert_eq!(header[2], 16);
    let mut buf = 0;

    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 4, &mut buf)
        .unwrap();
    assert_eq!(buf, 0);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 8, &mut buf)
        .unwrap();
    assert_eq!(buf, 0);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 12, &mut buf)
        .unwrap();
    assert_eq!(buf, 0x38);
    next_cap_offset = header[1] as u32;
    assert_ne!(next_cap_offset, 0);

    let mut header = 0;
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16, &mut header)
        .unwrap();
    let header = header.to_le_bytes();
    assert_eq!(header[0], CapabilityId::VENDOR_SPECIFIC.0);
    assert_eq!(header[3], VirtioPciCapType::NOTIFY_CFG.0);
    assert_eq!(header[2], 20);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 4, &mut buf)
        .unwrap();
    assert_eq!(buf, 0);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 8, &mut buf)
        .unwrap();
    assert_eq!(buf, 0x38);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 12, &mut buf)
        .unwrap();
    assert_eq!(buf, 4);
    next_cap_offset = header[1] as u32;
    assert_ne!(next_cap_offset, 0);

    let mut header = 0;
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16, &mut header)
        .unwrap();
    let header = header.to_le_bytes();
    assert_eq!(header[0], CapabilityId::VENDOR_SPECIFIC.0);
    assert_eq!(header[3], VirtioPciCapType::ISR_CFG.0);
    assert_eq!(header[2], 16);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 4, &mut buf)
        .unwrap();
    assert_eq!(buf, 0);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 8, &mut buf)
        .unwrap();
    assert_eq!(buf, 0x3c);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 12, &mut buf)
        .unwrap();
    assert_eq!(buf, 4);
    next_cap_offset = header[1] as u32;
    assert_ne!(next_cap_offset, 0);

    let mut header = 0;
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16, &mut header)
        .unwrap();
    let header = header.to_le_bytes();
    assert_eq!(header[0], CapabilityId::VENDOR_SPECIFIC.0);
    assert_eq!(header[3], VirtioPciCapType::DEVICE_CFG.0);
    assert_eq!(header[2], 16);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 4, &mut buf)
        .unwrap();
    assert_eq!(buf, 0);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 8, &mut buf)
        .unwrap();
    assert_eq!(buf, 0x40);
    pci_test_device
        .pci_device
        .pci_cfg_read(next_cap_offset as u16 + 12, &mut buf)
        .unwrap();
    assert_eq!(buf, 12);
    next_cap_offset = header[1] as u32;
    assert_eq!(next_cap_offset, 0);
}

#[async_test]
async fn verify_pci_registers(driver: DefaultDriver) {
    let mut pci_test_device =
        VirtioPciTestDevice::new(&driver, 1, &VirtioTestMemoryAccess::new(), None);
    let bar_address1: u64 = 0x2000000000;
    pci_test_device
        .pci_device
        .pci_cfg_write(0x14, (bar_address1 >> 32) as u32)
        .unwrap();
    pci_test_device
        .pci_device
        .pci_cfg_write(0x10, bar_address1 as u32)
        .unwrap();

    let bar_address2: u64 = 0x4000;
    pci_test_device
        .pci_device
        .pci_cfg_write(0x1c, (bar_address2 >> 32) as u32)
        .unwrap();
    pci_test_device
        .pci_device
        .pci_cfg_write(0x18, bar_address2 as u32)
        .unwrap();

    pci_test_device
        .pci_device
        .pci_cfg_write(
            0x4,
            cfg_space::Command::new()
                .with_mmio_enabled(true)
                .into_bits() as u32,
        )
        .unwrap();

    // device feature bank index
    assert_eq!(pci_test_device.read_u32(bar_address1), 0);
    // device feature (bank 0)
    assert_eq!(
        pci_test_device.read_u32(bar_address1 + 4),
        VIRTIO_F_RING_INDIRECT_DESC | VIRTIO_F_RING_EVENT_IDX | 2
    );
    // device feature (bank 1)
    pci_test_device.write_u32(bar_address1, 1);
    assert_eq!(pci_test_device.read_u32(bar_address1), 1);
    assert_eq!(
        pci_test_device.read_u32(bar_address1 + 4),
        VIRTIO_F_VERSION_1 | VIRTIO_F_RING_PACKED
    );
    // device feature (bank 2)
    pci_test_device.write_u32(bar_address1, 2);
    assert_eq!(pci_test_device.read_u32(bar_address1), 2);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 4), 0);
    // driver feature bank index
    assert_eq!(pci_test_device.read_u32(bar_address1 + 8), 0);
    // driver feature (bank 0)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 12), 0);
    pci_test_device.write_u32(bar_address1 + 12, 2);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 12), 2);
    pci_test_device.write_u32(bar_address1 + 12, 0xffffffff);
    assert_eq!(
        pci_test_device.read_u32(bar_address1 + 12),
        VIRTIO_F_RING_INDIRECT_DESC | VIRTIO_F_RING_EVENT_IDX | 2
    );
    // driver feature (bank 1)
    pci_test_device.write_u32(bar_address1 + 8, 1);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 8), 1);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 12), 0);
    pci_test_device.write_u32(bar_address1 + 12, 0xffffffff);
    assert_eq!(
        pci_test_device.read_u32(bar_address1 + 12),
        VIRTIO_F_VERSION_1 | VIRTIO_F_RING_PACKED
    );
    // driver feature (bank 2)
    pci_test_device.write_u32(bar_address1 + 8, 2);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 8), 2);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 12), 0);
    pci_test_device.write_u32(bar_address1 + 12, 0xffffffff);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 12), 0);
    // max queues and the msix vector for config changes
    assert_eq!(pci_test_device.read_u32(bar_address1 + 16), 1 << 16);
    // queue index, config generation and device status
    assert_eq!(pci_test_device.read_u32(bar_address1 + 20), 0);
    // current queue size and msix vector
    assert_eq!(pci_test_device.read_u32(bar_address1 + 24), 0x40);
    pci_test_device.write_u32(bar_address1 + 24, 0x20);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 24), 0x20);
    // current queue enabled and notify offset
    assert_eq!(pci_test_device.read_u32(bar_address1 + 28), 0);
    pci_test_device.write_u32(bar_address1 + 28, 1);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 28), 1);
    pci_test_device.write_u32(bar_address1 + 28, 0xffff);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 28), 1);
    pci_test_device.write_u32(bar_address1 + 28, 0);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 28), 0);
    // current queue descriptor table address (low)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 32), 0);
    pci_test_device.write_u32(bar_address1 + 32, 0xffff);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 32), 0xffff);
    // current queue descriptor table address (high)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 36), 0);
    pci_test_device.write_u32(bar_address1 + 36, 1);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 36), 1);
    // current queue available ring address (low)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 40), 0);
    pci_test_device.write_u32(bar_address1 + 40, 0xeeee);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 40), 0xeeee);
    // current queue available ring address (high)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 44), 0);
    pci_test_device.write_u32(bar_address1 + 44, 2);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 44), 2);
    // current queue used ring address (low)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 48), 0);
    pci_test_device.write_u32(bar_address1 + 48, 0xdddd);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 48), 0xdddd);
    // current queue used ring address (high)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 52), 0);
    pci_test_device.write_u32(bar_address1 + 52, 3);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 52), 3);
    // VIRTIO_PCI_CAP_NOTIFY_CFG notification register
    assert_eq!(pci_test_device.read_u32(bar_address1 + 56), 0);
    // VIRTIO_PCI_CAP_ISR_CFG register
    assert_eq!(pci_test_device.read_u32(bar_address1 + 60), 0);

    // switch to queue #1 (disabled, only one queue on this device)
    let queue_index: u16 = 1;
    pci_test_device
        .pci_device
        .mmio_write(bar_address1 + 22, &queue_index.to_le_bytes())
        .unwrap();
    assert_eq!(pci_test_device.read_u32(bar_address1 + 20), 1 << 24);
    // current queue size and msix vector
    assert_eq!(pci_test_device.read_u32(bar_address1 + 24), 0);
    pci_test_device.write_u32(bar_address1 + 24, 2);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 24), 0);
    // current queue enabled and notify offset
    assert_eq!(pci_test_device.read_u32(bar_address1 + 28), 0);
    pci_test_device.write_u32(bar_address1 + 28, 1);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 28), 0);
    // current queue descriptor table address (low)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 32), 0);
    pci_test_device.write_u32(bar_address1 + 32, 0x10);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 32), 0);
    // current queue descriptor table address (high)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 36), 0);
    pci_test_device.write_u32(bar_address1 + 36, 0x10);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 36), 0);
    // current queue available ring address (low)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 40), 0);
    pci_test_device.write_u32(bar_address1 + 40, 0x10);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 40), 0);
    // current queue available ring address (high)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 44), 0);
    pci_test_device.write_u32(bar_address1 + 44, 0x10);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 44), 0);
    // current queue used ring address (low)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 48), 0);
    pci_test_device.write_u32(bar_address1 + 48, 0x10);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 48), 0);
    // current queue used ring address (high)
    assert_eq!(pci_test_device.read_u32(bar_address1 + 52), 0);
    pci_test_device.write_u32(bar_address1 + 52, 0x10);
    assert_eq!(pci_test_device.read_u32(bar_address1 + 52), 0);
}

async fn verify_queue_simple_inner(mut guest: VirtioTestGuest) {
    let base_addr = guest.get_queue_descriptor_backing_memory_address(0);
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                assert_eq!(work.payload.len(), 1);
                assert_eq!(work.payload[0].length, 0x1000);
                match work.payload[0].address {
                    addr if addr == base_addr => work.complete(123),
                    addr if addr == base_addr + 0x1000 => work.complete(456),
                    _ => panic!("Unexpected address {}", work.payload[0].address),
                }
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    guest.add_to_avail_queue(0);
    guest.add_to_avail_queue(0);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 123);
    let (desc, len) = match guest.get_next_completed(0) {
        Some(v) => v,
        None => {
            must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
            guest.get_next_completed(0).unwrap()
        }
    };
    assert_eq!(desc, 1u16);
    assert_eq!(len, 456);
    assert_eq!(guest.get_next_completed(0).is_none(), true);
    queues[0].stop().await;
}

#[async_test]
async fn verify_split_queue_simple(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_simple_inner(VirtioTestGuest::new_split(&driver, &test_mem, 1, 2, true)).await;
}
#[async_test]
async fn verify_packed_queue_simple(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_simple_inner(VirtioTestGuest::new_packed(&driver, &test_mem, 1, 2, true)).await;
}

async fn verify_queue_simple_interrupt_control_inner(mut guest: VirtioTestGuest, with_index: bool) {
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                assert_eq!(work.payload.len(), 1);
                assert_eq!(work.payload[0].length, 0x1000);
                work.complete(123);
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    // interrupt on a specific descriptor
    if with_index {
        guest.enable_interrupt(0, Some(1));
        guest.add_to_avail_queue(0);
        event.signal();
        assert_no_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
        guest.add_to_avail_queue(0);
        event.signal();
        must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;

        let (_, len) = guest.get_next_completed(0).unwrap();
        assert_eq!(len, 123);
        let (_, len) = guest.get_next_completed(0).unwrap();
        assert_eq!(len, 123);
        assert_eq!(guest.get_next_completed(0).is_none(), true);
    }
    // always interrupt
    guest.enable_interrupt(0, None);
    guest.add_to_avail_queue(0);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (_, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(len, 123);
    assert_eq!(guest.get_next_completed(0).is_none(), true);

    // never interrupt
    guest.disable_interrupt(0);
    guest.add_to_avail_queue(0);
    guest.add_to_avail_queue(0);
    event.signal();
    assert_no_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (_, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(len, 123);
    let (_, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(len, 123);

    queues[0].stop().await;
}

#[async_test]
async fn verify_split_queue_simple_interrupt_control(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_simple_interrupt_control_inner(
        VirtioTestGuest::new_split(&driver, &test_mem, 1, 4, false),
        false,
    )
    .await;
}
#[async_test]
async fn verify_split_queue_simple_interrupt_control_with_index(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_simple_interrupt_control_inner(
        VirtioTestGuest::new_split(&driver, &test_mem, 1, 4, true),
        true,
    )
    .await;
}
#[async_test]
async fn verify_packed_queue_simple_interrupt_control(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_simple_interrupt_control_inner(
        VirtioTestGuest::new_packed(&driver, &test_mem, 1, 4, false),
        false,
    )
    .await;
}
#[async_test]
async fn verify_packed_queue_simple_interrupt_control_with_index(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_simple_interrupt_control_inner(
        VirtioTestGuest::new_packed(&driver, &test_mem, 1, 4, true),
        true,
    )
    .await;
}

async fn verify_queue_indirect_inner(mut guest: VirtioTestGuest) {
    let base_addr = guest.get_queue_descriptor_backing_memory_address(0);
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                assert_eq!(work.payload.len(), 1);
                assert_eq!(work.payload[0].length, 0x1000);
                match work.payload[0].address {
                    0xffffffff00000000u64 => work.complete(123),
                    addr if addr == base_addr + 0x1000 => work.complete(456),
                    _ => panic!("Unexpected address {}", work.payload[0].address),
                }
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    guest.add_indirect_to_avail_queue(0);
    guest.add_to_avail_queue(0);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 123);
    let (desc, len) = match guest.get_next_completed(0) {
        Some(v) => v,
        None => {
            must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
            guest.get_next_completed(0).unwrap()
        }
    };
    assert_eq!(desc, 1u16);
    assert_eq!(len, 456);
    assert_eq!(guest.get_next_completed(0).is_none(), true);
    queues[0].stop().await;
}

#[async_test]
async fn verify_split_queue_indirect(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_indirect_inner(VirtioTestGuest::new_split(&driver, &test_mem, 1, 2, true)).await;
}
#[async_test]
async fn verify_packed_queue_indirect(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_indirect_inner(VirtioTestGuest::new_packed(&driver, &test_mem, 1, 2, true)).await;
}

async fn verify_queue_linked_inner(mut guest: VirtioTestGuest) {
    let (tx, mut rx) = mesh::mpsc_channel();
    let base_address = guest.get_queue_descriptor_backing_memory_address(0);
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                if work.payload.len() == 3 {
                    for i in 0..work.payload.len() {
                        assert_eq!(work.payload[i].address, base_address + 0x1000 * i as u64);
                        assert_eq!(work.payload[i].length, 0x1000);
                    }
                    work.complete(123);
                } else {
                    assert_eq!(work.payload.len(), 1);
                    work.complete(456);
                }
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    guest.add_linked_to_avail_queue(0, 3);
    guest.add_to_avail_queue(0);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 123);
    let (desc, len) = match guest.get_next_completed(0) {
        Some(v) => v,
        None => {
            must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
            guest.get_next_completed(0).unwrap()
        }
    };
    assert_eq!(desc, 3u16);
    assert_eq!(len, 456);
    assert_eq!(guest.get_next_completed(0).is_none(), true);
    queues[0].stop().await;
}

#[async_test]
async fn verify_split_queue_linked(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_linked_inner(VirtioTestGuest::new_split(&driver, &test_mem, 1, 8, true)).await;
}
#[async_test]
async fn verify_packed_queue_linked(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_linked_inner(VirtioTestGuest::new_packed(&driver, &test_mem, 1, 8, true)).await;
}

async fn verify_queue_indirect_linked_inner(mut guest: VirtioTestGuest) {
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                if work.payload.len() == 3 {
                    for i in 0..work.payload.len() {
                        assert_eq!(
                            work.payload[i].address,
                            0xffffffff00000000u64 + 0x1000 * i as u64
                        );
                        assert_eq!(work.payload[i].length, 0x1000);
                    }
                    work.complete(123);
                } else {
                    assert_eq!(work.payload.len(), 1);
                    work.complete(456);
                }
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    guest.add_indirect_linked_to_avail_queue(0, 3);
    guest.add_to_avail_queue(0);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 123);
    let (desc, len) = match guest.get_next_completed(0) {
        Some(v) => v,
        None => {
            must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
            guest.get_next_completed(0).unwrap()
        }
    };
    assert_eq!(desc, 1u16);
    assert_eq!(len, 456);
    assert_eq!(guest.get_next_completed(0).is_none(), true);
    queues[0].stop().await;
}

#[async_test]
async fn verify_split_queue_indirect_linked(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_indirect_linked_inner(VirtioTestGuest::new_split(&driver, &test_mem, 1, 8, true))
        .await;
}
#[async_test]
async fn verify_packed_queue_indirect_linked(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_indirect_linked_inner(VirtioTestGuest::new_packed(&driver, &test_mem, 1, 8, true))
        .await;
}

async fn verify_queue_avail_rollover_inner(mut guest: VirtioTestGuest) {
    let base_addr = guest.get_queue_descriptor_backing_memory_address(0);
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                assert_eq!(work.payload.len(), 1);
                assert_eq!(work.payload[0].length, 0x1000);
                if work.payload[0].address == base_addr {
                    work.complete(123);
                } else if work.payload[0].address == base_addr + 0x1000 {
                    work.complete(456);
                } else {
                    panic!(
                        "Unexpected descriptor address {:x}",
                        work.payload[0].address
                    );
                }
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    for _ in 0..3 {
        guest.add_to_avail_queue(0);
        event.signal();
        must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
        let (desc, len) = guest.get_next_completed(0).unwrap();
        if desc == 0 {
            assert_eq!(len, 123);
        } else if desc == 1 {
            assert_eq!(len, 456);
        } else {
            panic!("Unexpected descriptor index");
        }
        assert_eq!(guest.get_next_completed(0).is_none(), true);
    }

    queues[0].stop().await;
}

#[async_test]
async fn verify_split_queue_avail_rollover(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_avail_rollover_inner(VirtioTestGuest::new_split(&driver, &test_mem, 1, 2, true))
        .await;
}
#[async_test]
async fn verify_packed_queue_avail_rollover(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    verify_queue_avail_rollover_inner(VirtioTestGuest::new_packed(&driver, &test_mem, 1, 2, true))
        .await;
}

async fn verify_multi_queue_inner(mut guest: VirtioTestGuest) {
    let (tx, mut rx) = mesh::mpsc_channel();
    let events = (0..guest.num_queues)
        .map(|_| Event::new())
        .collect::<Vec<_>>();
    let mut queues = guest.create_direct_queues(|queue_index| {
        let tx = tx.clone();
        let base_addr = guest.get_queue_descriptor_backing_memory_address(queue_index);
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                assert_eq!(work.payload.len(), 1);
                assert_eq!(work.payload[0].address, base_addr);
                assert_eq!(work.payload[0].length, 0x1000);
                work.complete(123 * queue_index as u32);
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(queue_index as usize);
            }),
            event: events[queue_index as usize].clone(),
        }
    });

    for (i, event) in events.iter().enumerate() {
        let queue_index = i as u16;
        guest.add_to_avail_queue(queue_index);
        event.signal();
    }
    // wait for all queue processing to finish
    for _ in 0..guest.num_queues {
        must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    }
    // check results
    for queue_index in 0..guest.num_queues {
        let (desc, len) = guest.get_next_completed(queue_index).unwrap();
        assert_eq!(desc, 0u16);
        assert_eq!(len, 123 * queue_index as u32);
    }
    // verify no extraneous completions
    for (i, queue) in queues.iter_mut().enumerate() {
        let queue_index = i as u16;
        assert_eq!(guest.get_next_completed(queue_index).is_none(), true);
        queue.stop().await;
    }
}

#[async_test]
async fn verify_split_multi_queue(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_split(&driver, &test_mem, 5, 2, true);
    verify_multi_queue_inner(guest).await;
}
#[async_test]
async fn verify_packed_multi_queue(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_packed(&driver, &test_mem, 5, 2, true);
    verify_multi_queue_inner(guest).await;
}

fn take_mmio_interrupt_status(dev: &mut VirtioMmioDevice, mask: u32) -> u32 {
    let mut v = [0; 4];
    dev.mmio_read(96, &mut v).unwrap();
    dev.mmio_write(100, &mask.to_ne_bytes()).unwrap();
    u32::from_ne_bytes(v)
}

async fn expect_mmio_interrupt(
    dev: &mut VirtioMmioDevice,
    target: &TestLineInterruptTarget,
    mask: u32,
    multiple_expected: bool,
) {
    poll_fn(|cx| target.poll_high(cx, 0)).await;
    let v = take_mmio_interrupt_status(dev, mask);
    assert_eq!(v & mask, mask);
    assert!(multiple_expected || !target.is_high(0));
}

async fn verify_device_queue_simple_inner(
    test_mem: Arc<VirtioTestMemoryAccess>,
    mut guest: VirtioTestGuest,
    features: VirtioDeviceFeatures,
) {
    let doorbell_registration: Arc<dyn DoorbellRegistration> = test_mem.clone();
    let mem = guest.mem();
    let target = TestLineInterruptTarget::new_arc();
    let interrupt = LineInterrupt::new_with_target("test", target.clone(), 0);
    let base_addr = guest.get_queue_descriptor_backing_memory_address(0);
    let queue_work = Arc::new(move |_: u16, mut work: VirtioQueueCallbackWork| {
        assert_eq!(work.payload.len(), 1);
        assert_eq!(work.payload[0].address, base_addr);
        assert_eq!(work.payload[0].length, 0x1000);
        work.complete(123);
    });
    let driver_source = VmTaskDriverSource::new(SingleDriverBackend::new(guest.driver()));
    let mut dev = VirtioMmioDevice::new(
        Box::new(TestDevice::new(
            &driver_source,
            DeviceTraits {
                device_id: 3,
                device_features: features.clone(),
                max_queues: 1,
                device_register_length: 0,
                ..Default::default()
            },
            Some(queue_work),
            &mem,
        )),
        interrupt,
        Some(doorbell_registration),
        0,
        1,
    );

    guest.setup_chipset_device(&mut dev, features);
    expect_mmio_interrupt(
        &mut dev,
        &target,
        VIRTIO_MMIO_INTERRUPT_STATUS_CONFIG_CHANGE,
        false,
    )
    .await;
    guest.add_to_avail_queue(0);
    // notify device
    dev.write_u32(80, 0);
    expect_mmio_interrupt(
        &mut dev,
        &target,
        VIRTIO_MMIO_INTERRUPT_STATUS_USED_BUFFER,
        false,
    )
    .await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 123);
    assert_eq!(guest.get_next_completed(0).is_none(), true);
    // reset the device
    dev.write_u32(112, 0);
    drop(dev);
}

#[async_test]
async fn verify_device_split_queue_simple(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_split(&driver, &test_mem, 1, 2, true);
    let features = VirtioDeviceFeatures::new()
        .with_bank(0, VIRTIO_F_RING_EVENT_IDX | 2)
        .with_bank(1, VIRTIO_F_VERSION_1);
    verify_device_queue_simple_inner(test_mem, guest, features).await;
}
#[async_test]
async fn verify_device_packed_queue_simple(driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_packed(&driver, &test_mem, 1, 2, true);
    let features = VirtioDeviceFeatures::new()
        .with_bank(0, VIRTIO_F_RING_EVENT_IDX | 2)
        .with_bank(1, VIRTIO_F_VERSION_1 | VIRTIO_F_RING_PACKED);
    verify_device_queue_simple_inner(test_mem, guest, features).await;
}

async fn verify_device_multi_queue_inner(
    test_mem: Arc<VirtioTestMemoryAccess>,
    mut guest: VirtioTestGuest,
    num_queues: u16,
    features: VirtioDeviceFeatures,
) {
    let doorbell_registration: Arc<dyn DoorbellRegistration> = test_mem.clone();
    let mem = guest.mem();
    let target = TestLineInterruptTarget::new_arc();
    let interrupt = LineInterrupt::new_with_target("test", target.clone(), 0);
    let base_addr: Vec<_> = (0..num_queues)
        .map(|i| guest.get_queue_descriptor_backing_memory_address(i))
        .collect();
    let queue_work = Arc::new(move |i: u16, mut work: VirtioQueueCallbackWork| {
        assert_eq!(work.payload.len(), 1);
        assert_eq!(work.payload[0].address, base_addr[i as usize]);
        assert_eq!(work.payload[0].length, 0x1000);
        work.complete(123 * i as u32);
    });
    let driver_source = VmTaskDriverSource::new(SingleDriverBackend::new(guest.driver()));
    let mut dev = VirtioMmioDevice::new(
        Box::new(TestDevice::new(
            &driver_source,
            DeviceTraits {
                device_id: 3,
                device_features: features.clone(),
                max_queues: num_queues + 1,
                device_register_length: 0,
                ..Default::default()
            },
            Some(queue_work),
            &mem,
        )),
        interrupt,
        Some(doorbell_registration),
        0,
        1,
    );
    guest.setup_chipset_device(&mut dev, features);
    expect_mmio_interrupt(
        &mut dev,
        &target,
        VIRTIO_MMIO_INTERRUPT_STATUS_CONFIG_CHANGE,
        false,
    )
    .await;
    for i in 0..num_queues {
        guest.add_to_avail_queue(i);
        // notify device
        dev.write_u32(80, i as u32);
    }
    // check results
    for i in 0..num_queues {
        let (desc, len) = loop {
            if let Some(x) = guest.get_next_completed(i) {
                break x;
            }
            expect_mmio_interrupt(
                &mut dev,
                &target,
                VIRTIO_MMIO_INTERRUPT_STATUS_USED_BUFFER,
                i < (num_queues - 1),
            )
            .await;
        };
        assert_eq!(desc, 0u16);
        assert_eq!(len, 123 * i as u32);
    }
    // verify no extraneous completions
    for i in 0..num_queues {
        assert_eq!(guest.get_next_completed(i).is_none(), true);
    }
    // reset the device
    dev.write_u32(112, 0);
    drop(dev);
}

#[async_test]
async fn verify_device_split_multi_queue(driver: DefaultDriver) {
    let num_queues = 5;
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_split(&driver, &test_mem, num_queues, 2, true);
    let features = VirtioDeviceFeatures::new()
        .with_bank(0, VIRTIO_F_RING_EVENT_IDX | 2)
        .with_bank(1, VIRTIO_F_VERSION_1);
    verify_device_multi_queue_inner(test_mem, guest, num_queues, features).await;
}
#[async_test]
async fn verify_device_packed_multi_queue(driver: DefaultDriver) {
    let num_queues = 5;
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_packed(&driver, &test_mem, num_queues, 2, true);
    let features = VirtioDeviceFeatures::new()
        .with_bank(0, VIRTIO_F_RING_EVENT_IDX | 2)
        .with_bank(1, VIRTIO_F_VERSION_1 | VIRTIO_F_RING_PACKED);
    verify_device_multi_queue_inner(test_mem, guest, num_queues, features).await;
}

async fn verify_device_multi_queue_pci_inner(
    test_mem: Arc<VirtioTestMemoryAccess>,
    mut guest: VirtioTestGuest,
    num_queues: u16,
    features: VirtioDeviceFeatures,
) {
    let driver = guest.driver();
    let base_addr: Vec<_> = (0..num_queues)
        .map(|i| guest.get_queue_descriptor_backing_memory_address(i))
        .collect();
    let mut dev = VirtioPciTestDevice::new(
        &driver,
        num_queues + 1,
        &test_mem,
        Some(Arc::new(move |i, mut work| {
            assert_eq!(work.payload.len(), 1);
            assert_eq!(work.payload[0].address, base_addr[i as usize]);
            assert_eq!(work.payload[0].length, 0x1000);
            work.complete(123 * i as u32);
        })),
    );

    guest.setup_pci_device(&mut dev, features);

    let mut timer = PolledTimer::new(&driver);

    // expect a config generation interrupt
    timer.sleep(Duration::from_millis(100)).await;
    let delivered = dev.test_intc.get_next_interrupt().unwrap();
    assert_eq!(delivered.0, 0);
    assert!(dev.test_intc.get_next_interrupt().is_none());

    for i in 0..num_queues {
        guest.add_to_avail_queue(i);
        // notify device
        dev.write_u32(0x10000000000 + 0x38, i as u32);
    }
    // verify all queue processing finished
    timer.sleep(Duration::from_millis(100)).await;
    for _ in 0..num_queues {
        let delivered = dev.test_intc.get_next_interrupt();
        assert!(delivered.is_some());
    }
    // check results
    for i in 0..num_queues {
        let (desc, len) = guest.get_next_completed(i).unwrap();
        assert_eq!(desc, 0u16);
        assert_eq!(len, 123 * i as u32);
    }
    // verify no extraneous completions
    for i in 0..num_queues {
        assert_eq!(guest.get_next_completed(i).is_none(), true);
    }
    // reset the device
    let device_status: u8 = 0;
    dev.pci_device
        .mmio_write(0x10000000000 + 20, &device_status.to_le_bytes())
        .unwrap();
    drop(dev);
}

#[async_test]
async fn verify_device_split_multi_queue_pci(driver: DefaultDriver) {
    let num_queues = 5;
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_split(&driver, &test_mem, num_queues, 2, true);
    let features = VirtioDeviceFeatures::new()
        .with_bank(0, VIRTIO_F_RING_EVENT_IDX | 2)
        .with_bank(1, VIRTIO_F_VERSION_1);
    verify_device_multi_queue_pci_inner(test_mem, guest, num_queues, features).await;
}
#[async_test]
async fn verify_device_packed_multi_queue_pci(driver: DefaultDriver) {
    let num_queues = 5;
    let test_mem = VirtioTestMemoryAccess::new();
    let guest = VirtioTestGuest::new_packed(&driver, &test_mem, num_queues, 2, true);
    let features = VirtioDeviceFeatures::new()
        .with_bank(0, VIRTIO_F_RING_EVENT_IDX | 2)
        .with_bank(1, VIRTIO_F_VERSION_1 | VIRTIO_F_RING_PACKED);
    verify_device_multi_queue_pci_inner(test_mem, guest, num_queues, features).await;
}

#[async_test]
async fn verify_enable_failure_mmio_does_not_set_driver_ok(_driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    let doorbell_registration: Arc<dyn DoorbellRegistration> = test_mem.clone();
    let interrupt = LineInterrupt::detached();

    let mut dev = VirtioMmioDevice::new(
        Box::new(FailingTestDevice {
            traits: DeviceTraits {
                device_id: 3,
                device_features: VirtioDeviceFeatures::new().with_bank(0, 2),
                max_queues: 1,
                device_register_length: 0,
                ..Default::default()
            },
        }),
        interrupt,
        Some(doorbell_registration),
        0,
        1,
    );

    // Drive through ACKNOWLEDGE -> DRIVER -> FEATURES_OK -> DRIVER_OK
    dev.write_u32(112, VIRTIO_ACKNOWLEDGE);
    dev.write_u32(112, VIRTIO_DRIVER);
    dev.write_u32(36, 0);
    dev.write_u32(32, 2); // select matching features
    dev.write_u32(36, 1);
    dev.write_u32(32, VIRTIO_F_VERSION_1);
    dev.write_u32(112, VIRTIO_FEATURES_OK);

    // Set up one queue
    dev.write_u32(48, 0); // queue select
    dev.write_u32(56, 16); // queue size
    dev.write_u32(68, 1); // queue enable

    // Attempt DRIVER_OK — enable() will fail
    dev.write_u32(112, VIRTIO_DRIVER_OK);

    // Device status should NOT have DRIVER_OK set
    let status = dev.read_u32(112);
    assert_eq!(
        status & VIRTIO_DRIVER_OK,
        0,
        "DRIVER_OK must not be set when enable() fails"
    );
}

#[async_test]
async fn verify_enable_failure_pci_does_not_set_driver_ok(_driver: DefaultDriver) {
    let test_mem = VirtioTestMemoryAccess::new();
    let doorbell_registration: Arc<dyn DoorbellRegistration> = test_mem.clone();
    let msi_conn = MsiConnection::new();

    let mut dev = VirtioPciDevice::new(
        Box::new(FailingTestDevice {
            traits: DeviceTraits {
                device_id: 3,
                device_features: VirtioDeviceFeatures::new().with_bank(0, 2),
                max_queues: 1,
                device_register_length: 12,
                ..Default::default()
            },
        }),
        PciInterruptModel::Msix(msi_conn.target()),
        Some(doorbell_registration),
        &mut ExternallyManagedMmioIntercepts,
        None,
    )
    .unwrap();

    let bar_address1: u64 = 0x10000000000;
    dev.pci_cfg_write(0x14, (bar_address1 >> 32) as u32)
        .unwrap();
    dev.pci_cfg_write(0x10, bar_address1 as u32).unwrap();

    let bar_address2: u64 = 0x20000000000;
    dev.pci_cfg_write(0x1c, (bar_address2 >> 32) as u32)
        .unwrap();
    dev.pci_cfg_write(0x18, bar_address2 as u32).unwrap();

    dev.pci_cfg_write(
        0x4,
        cfg_space::Command::new()
            .with_mmio_enabled(true)
            .into_bits() as u32,
    )
    .unwrap();

    // Drive through ACKNOWLEDGE -> DRIVER -> FEATURES_OK
    let mut buf = [0u8; 1];
    buf[0] = VIRTIO_ACKNOWLEDGE as u8;
    dev.mmio_write(bar_address1 + 20, &buf).unwrap();
    buf[0] = VIRTIO_DRIVER as u8;
    dev.mmio_write(bar_address1 + 20, &buf).unwrap();
    // Select features
    let mut val;
    val = 0u32.to_le_bytes();
    dev.mmio_write(bar_address1 + 8, &val).unwrap();
    val = 2u32.to_le_bytes();
    dev.mmio_write(bar_address1 + 12, &val).unwrap();
    val = 1u32.to_le_bytes();
    dev.mmio_write(bar_address1 + 8, &val).unwrap();
    val = VIRTIO_F_VERSION_1.to_le_bytes();
    dev.mmio_write(bar_address1 + 12, &val).unwrap();
    buf[0] = VIRTIO_FEATURES_OK as u8;
    dev.mmio_write(bar_address1 + 20, &buf).unwrap();

    // Set up queue 0
    dev.mmio_write(bar_address1 + 22, &0u16.to_le_bytes())
        .unwrap(); // queue select
    dev.mmio_write(bar_address1 + 24, &16u16.to_le_bytes())
        .unwrap(); // queue size
    // Set up MSI for the queue
    dev.mmio_write(bar_address2, &0u64.to_le_bytes()).unwrap();
    dev.mmio_write(bar_address2 + 8, &0u32.to_le_bytes())
        .unwrap();
    dev.mmio_write(bar_address2 + 12, &0u32.to_le_bytes())
        .unwrap();
    let msix_vector: u16 = 1;
    let msix_addr = bar_address2 + 0x10 * msix_vector as u64;
    dev.mmio_write(msix_addr, &(msix_vector as u64).to_le_bytes())
        .unwrap();
    dev.mmio_write(msix_addr + 8, &0u32.to_le_bytes()).unwrap();
    dev.mmio_write(msix_addr + 12, &0u32.to_le_bytes()).unwrap();
    dev.mmio_write(bar_address1 + 26, &msix_vector.to_le_bytes())
        .unwrap();
    // Enable queue
    dev.mmio_write(bar_address1 + 28, &1u16.to_le_bytes())
        .unwrap();
    // Enable all MSI interrupts
    dev.pci_cfg_write(0x40, 0x80000000).unwrap();

    // Attempt DRIVER_OK — enable() will fail
    buf[0] = VIRTIO_DRIVER_OK as u8;
    dev.mmio_write(bar_address1 + 20, &buf).unwrap();

    // Read back device status
    let mut status_buf = [0u8; 1];
    dev.mmio_read(bar_address1 + 20, &mut status_buf).unwrap();
    let status = status_buf[0] as u32;
    assert_eq!(
        status & VIRTIO_DRIVER_OK,
        0,
        "DRIVER_OK must not be set when enable() fails"
    );
}

/// A linked chain using all queue_size descriptors (last has no NEXT flag)
/// must succeed — this is the maximum valid chain length per spec §2.7.5.3.1.
#[async_test]
async fn verify_chain_at_queue_size_succeeds(driver: DefaultDriver) {
    let queue_size: u16 = 4;
    let test_mem = VirtioTestMemoryAccess::new();
    let mut guest = VirtioTestGuest::new_split(&driver, &test_mem, 1, queue_size, true);
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                assert_eq!(work.payload.len(), queue_size as usize);
                work.complete(42);
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    guest.add_linked_to_avail_queue(0, queue_size);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 42);
    queues[0].stop().await;
}

/// A descriptor chain that forms a cycle (all queue_size descriptors point to
/// the next, with the last wrapping back to the first) must be rejected.
#[async_test]
async fn verify_chain_cycle_rejected(driver: DefaultDriver) {
    let queue_size: u16 = 4;
    let test_mem = VirtioTestMemoryAccess::new();
    let mut guest = VirtioTestGuest::new_split(&driver, &test_mem, 1, queue_size, true);
    let event = Event::new();
    let queue_event = PolledWait::new(&driver, event.clone()).unwrap();
    let mut queue = VirtioQueue::new(
        guest.queue_features(),
        guest.queue_params(0),
        guest.mem(),
        Interrupt::from_fn(|| {}),
        queue_event,
    )
    .unwrap();

    // Manually build a cycle: desc 0 → 1 → 2 → 3 → 0 (all with NEXT flag).
    for i in 0..queue_size {
        let base = guest.get_queue_descriptor(0, i);
        let flags = u16::from(DescriptorFlags::new().with_next(true));
        test_mem.modify_memory_map(base + 12, &flags.to_le_bytes(), false);
        let next = (i + 1) % queue_size;
        test_mem.modify_memory_map(base + 14, &next.to_le_bytes(), false);
    }
    guest.queue_available_desc(0, 0);

    let result = queue.try_next();
    assert!(result.is_err(), "Cyclic chain must produce an error");
}

/// An indirect descriptor table with more entries than queue_size must be
/// clamped to queue_size (spec §2.7.5.3.1: "A driver MUST NOT create a
/// descriptor chain longer than the Queue Size of the device").
#[async_test]
async fn verify_indirect_chain_clamped_to_queue_size(driver: DefaultDriver) {
    let queue_size: u16 = 4;
    let test_mem = VirtioTestMemoryAccess::new();
    let mut guest = VirtioTestGuest::new_split(&driver, &test_mem, 1, queue_size, true);
    let event = Event::new();
    let queue_event = PolledWait::new(&driver, event.clone()).unwrap();
    let mut queue = VirtioQueue::new(
        guest.queue_features(),
        guest.queue_params(0),
        guest.mem(),
        Interrupt::from_fn(|| {}),
        queue_event,
    )
    .unwrap();

    // Build an indirect descriptor pointing to an 8-entry table (> queue_size).
    let indirect_entries: u16 = 8;
    let next_descriptor = 0u16;
    let desc_base = guest.get_queue_descriptor(0, next_descriptor);
    // Set INDIRECT flag on the ring descriptor.
    test_mem.modify_memory_map(
        desc_base + 12,
        &u16::from(DescriptorFlags::new().with_indirect(true)).to_le_bytes(),
        false,
    );
    // Point addr/len to our indirect table in the backing buffer area.
    let buffer_addr = guest.get_queue_descriptor_backing_memory_address(0);
    test_mem.modify_memory_map(desc_base, &buffer_addr.to_le_bytes(), false);
    test_mem.modify_memory_map(
        desc_base + 8,
        &(indirect_entries as u32 * 16).to_le_bytes(),
        false,
    );
    // Fill out the indirect table: 8 linked entries.
    for i in 0..indirect_entries {
        let base = buffer_addr + 0x10 * i as u64;
        let indirect_buffer_addr = 0xffffffff00000000u64 + 0x1000 * i as u64;
        test_mem.modify_memory_map(base, &indirect_buffer_addr.to_le_bytes(), false);
        test_mem.modify_memory_map(base + 8, &0x1000u32.to_le_bytes(), false);
        let flags = if i < indirect_entries - 1 {
            u16::from(DescriptorFlags::new().with_next(true))
        } else {
            0
        };
        test_mem.modify_memory_map(base + 12, &flags.to_le_bytes(), false);
        let next = if i < indirect_entries - 1 { i + 1 } else { 0 };
        test_mem.modify_memory_map(base + 14, &next.to_le_bytes(), false);
    }
    guest.queue_available_desc(0, next_descriptor);

    let result = queue.try_next();
    assert!(
        result.is_err(),
        "Indirect chain exceeding queue_size must be rejected"
    );
}

/// An indirect table with exactly queue_size entries (last has no NEXT)
/// must succeed — this is the maximum valid indirect chain.
#[async_test]
async fn verify_indirect_chain_at_queue_size_succeeds(driver: DefaultDriver) {
    let queue_size: u16 = 4;
    let test_mem = VirtioTestMemoryAccess::new();
    let mut guest = VirtioTestGuest::new_split(&driver, &test_mem, 1, queue_size, true);
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                assert_eq!(work.payload.len(), queue_size as usize);
                work.complete(99);
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    guest.add_indirect_linked_to_avail_queue(0, queue_size);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 99);
    queues[0].stop().await;
}

/// A chain of normal descriptors followed by an indirect descriptor must work
/// (spec §2.7.5.3.2: "The device MUST handle the case of zero or more normal
/// chained descriptors followed by a single descriptor with
/// flags&VIRTQ_DESC_F_INDIRECT").
#[async_test]
async fn verify_normal_then_indirect_succeeds(driver: DefaultDriver) {
    let queue_size: u16 = 8;
    let test_mem = VirtioTestMemoryAccess::new();
    let mut guest = VirtioTestGuest::new_split(&driver, &test_mem, 1, queue_size, true);
    let (tx, mut rx) = mesh::mpsc_channel();
    let event = Event::new();
    let mut queues = guest.create_direct_queues(|i| {
        let tx = tx.clone();
        CreateDirectQueueParams {
            process_work: Box::new(move |mut work: VirtioQueueCallbackWork| {
                // 1 normal descriptor + 3 indirect entries = 4 payload entries.
                // (The indirect head descriptor is not itself a payload entry.)
                assert_eq!(work.payload.len(), 4);
                work.complete(77);
            }),
            notify: Interrupt::from_fn(move || {
                tx.send(i as usize);
            }),
            event: event.clone(),
        }
    });

    // Build: desc 0 (normal, NEXT→1) → desc 1 (INDIRECT, points to 3-entry table).
    // Reserve the two ring descriptors we're using manually.
    let d0 = guest.reserve_split_descriptor(0);
    let d1 = guest.reserve_split_descriptor(0);
    assert_eq!(d0, 0);
    assert_eq!(d1, 1);

    // Desc 0: normal descriptor with NEXT flag.
    let desc0_base = guest.get_queue_descriptor(0, 0);
    test_mem.modify_memory_map(
        desc0_base + 12,
        &u16::from(DescriptorFlags::new().with_next(true)).to_le_bytes(),
        false,
    );
    test_mem.modify_memory_map(desc0_base + 14, &1u16.to_le_bytes(), false);

    // Desc 1: INDIRECT descriptor pointing to a 3-entry indirect table.
    let desc1_base = guest.get_queue_descriptor(0, 1);
    let indirect_table_addr = 0xAABB00000000u64;
    let indirect_entries: u16 = 3;
    test_mem.modify_memory_map(desc1_base, &indirect_table_addr.to_le_bytes(), false);
    test_mem.modify_memory_map(
        desc1_base + 8,
        &(indirect_entries as u32 * 16).to_le_bytes(),
        false,
    );
    test_mem.modify_memory_map(
        desc1_base + 12,
        &u16::from(DescriptorFlags::new().with_indirect(true)).to_le_bytes(),
        false,
    );

    // Create the indirect table in memory at indirect_table_addr.
    for i in 0..indirect_entries {
        let entry_base = indirect_table_addr + 0x10 * i as u64;
        let entry_buf_addr = 0xDDEE00000000u64 + 0x1000 * i as u64;
        test_mem.modify_memory_map(entry_base, &entry_buf_addr.to_le_bytes(), false);
        test_mem.modify_memory_map(entry_base + 8, &0x1000u32.to_le_bytes(), false);
        let flags = if i < indirect_entries - 1 {
            u16::from(DescriptorFlags::new().with_next(true))
        } else {
            0
        };
        test_mem.modify_memory_map(entry_base + 12, &flags.to_le_bytes(), false);
        let next = if i < indirect_entries - 1 { i + 1 } else { 0 };
        test_mem.modify_memory_map(entry_base + 14, &next.to_le_bytes(), false);
    }

    guest.queue_available_desc(0, 0);
    event.signal();
    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
    let (desc, len) = guest.get_next_completed(0).unwrap();
    assert_eq!(desc, 0u16);
    assert_eq!(len, 77);
    queues[0].stop().await;
}
