#[path = "../../../os-device-registry/src/lib.rs"]
mod device_registry;

use device_registry::{Device, DeviceKey, DeviceKind, Error, Registry, Transport};
use device_registry::{DeviceClass, DeviceIdentity, Driver, MatchError, MatchedRegistry,
    ResourceKind, ResourceRange, QEMU_VIRTIO_VENDOR, DetachError, DetachPhase, SafeRegistry};

const BLOCK: DeviceKey = DeviceKey { transport: Transport::Mmio, location: 0x0a00_0000 };
const NETWORK: DeviceKey = DeviceKey { transport: Transport::Pci, location: 0x100 };

#[test]
fn os_device_registry_identity_positive() {
    let mut registry = Registry::<2>::new();
    let block = registry.attach(BLOCK, DeviceKind::Block).unwrap();
    let network = registry.attach(NETWORK, DeviceKind::Network).unwrap();
    assert_eq!(registry.live_count(), 2);
    assert_eq!(registry.get(block), Ok(Device { key: BLOCK, kind: DeviceKind::Block }));
    assert_eq!(registry.get(network), Ok(Device { key: NETWORK, kind: DeviceKind::Network }));
    assert_eq!(registry.detach(block), Ok(Device { key: BLOCK, kind: DeviceKind::Block }));
    assert_eq!(registry.get(block), Err(Error::StaleHandle));
    let new_block = registry.attach(BLOCK, DeviceKind::Block).unwrap();
    assert_eq!(new_block.slot(), block.slot());
    assert_eq!(new_block.generation(), block.generation() + 1);
    assert_eq!(registry.live_count(), 2);
    assert_eq!(registry.get(new_block).unwrap().key, BLOCK);
    assert_eq!(registry.get(network).unwrap().key, NETWORK);
}

#[test]
fn os_device_registry_identity_negative() {
    let mut registry = Registry::<2>::new();
    let block = registry.attach(BLOCK, DeviceKind::Block).unwrap();
    let network = registry.attach(NETWORK, DeviceKind::Network).unwrap();
    let before_duplicate = registry.snapshot();
    assert_eq!(registry.attach(BLOCK, DeviceKind::Network), Err(Error::AlreadyAttached));
    assert_eq!(registry.snapshot(), before_duplicate);

    assert_eq!(registry.detach(block).unwrap().key, BLOCK);
    let replacement = registry.attach(BLOCK, DeviceKind::Block).unwrap();
    let before_stale = registry.snapshot();
    assert_eq!(registry.get(block), Err(Error::StaleHandle));
    assert_eq!(registry.detach(block), Err(Error::StaleHandle));
    assert_eq!(registry.snapshot(), before_stale);
    assert_eq!(registry.get(replacement).unwrap().key, BLOCK);
    assert_eq!(registry.get(network).unwrap().key, NETWORK);

    let before_full = registry.snapshot();
    assert_eq!(registry.attach(DeviceKey { transport: Transport::Mmio, location: 0x0a00_0200 }, DeviceKind::Block), Err(Error::Full));
    assert_eq!(registry.snapshot(), before_full);

    let mut exhausted = Registry::<1>::with_generation_seeds([u64::MAX]);
    let before_exhausted = exhausted.snapshot();
    assert_eq!(exhausted.attach(BLOCK, DeviceKind::Block), Err(Error::GenerationExhausted));
    assert_eq!(exhausted.snapshot(), before_exhausted);

    let mut empty = Registry::<0>::new();
    let before_empty = empty.snapshot();
    assert_eq!(empty.attach(BLOCK, DeviceKind::Block), Err(Error::Full));
    assert_eq!(empty.snapshot(), before_empty);
}

const BLOCK_IDENTITY: DeviceIdentity = DeviceIdentity {
    transport: Transport::Mmio, vendor_id: QEMU_VIRTIO_VENDOR,
    device_id: 2, class: DeviceClass::Block,
};
const NETWORK_IDENTITY: DeviceIdentity = DeviceIdentity {
    transport: Transport::Mmio, vendor_id: QEMU_VIRTIO_VENDOR,
    device_id: 1, class: DeviceClass::Network,
};
const MMIO_NETWORK: DeviceKey = DeviceKey { transport: Transport::Mmio, location: 0x0a00_0200 };
const BLOCK_RANGE: [ResourceRange; 1] = [ResourceRange {
    kind: ResourceKind::Mmio, start: 0x0a00_0000, length: 0x200,
}];
const NETWORK_RANGE: [ResourceRange; 1] = [ResourceRange {
    kind: ResourceKind::Mmio, start: 0x0a00_0200, length: 0x200,
}];

#[test]
fn os_device_registry_driver_match_positive() {
    let mut registry = MatchedRegistry::<2, 2>::new();
    let block = registry.attach(BLOCK, BLOCK_IDENTITY, &BLOCK_RANGE).unwrap();
    let network = registry.attach(MMIO_NETWORK, NETWORK_IDENTITY, &NETWORK_RANGE).unwrap();
    assert_eq!((registry.binding_count(), registry.resource_count()), (2, 2));
    let (block_device, block_binding) = registry.get(block).unwrap();
    let (network_device, network_binding) = registry.get(network).unwrap();
    assert_eq!((block_device.key, block_device.kind, block_binding.driver),
        (BLOCK, DeviceKind::Block, Driver::VirtioBlock));
    assert_eq!((network_device.key, network_device.kind, network_binding.driver),
        (MMIO_NETWORK, DeviceKind::Network, Driver::VirtioNetwork));
    let snapshot = registry.snapshot();
    assert_eq!(snapshot.bindings.iter().flatten().count(), 2);
    assert_eq!(snapshot.resources.claims.iter().flatten().count(), 2);
    assert_eq!(registry.detach(block).unwrap().key, BLOCK);
    assert_eq!((registry.binding_count(), registry.resource_count()), (1, 1));
    assert_eq!(registry.get(network).unwrap().1.driver, Driver::VirtioNetwork);
    assert_eq!(registry.detach(network).unwrap().key, MMIO_NETWORK);
    assert_eq!((registry.binding_count(), registry.resource_count()), (0, 0));
}

#[test]
fn os_device_registry_driver_match_negative() {
    let mut registry = MatchedRegistry::<2, 2>::new();
    let block = registry.attach(BLOCK, BLOCK_IDENTITY, &BLOCK_RANGE).unwrap();
    let before = registry.snapshot();
    let unsupported = DeviceIdentity {class: DeviceClass::Other(0xff), ..BLOCK_IDENTITY};
    assert_eq!(registry.attach(MMIO_NETWORK, unsupported, &NETWORK_RANGE),
        Err(MatchError::UnsupportedClass));
    assert_eq!(registry.snapshot(), before);
    let mismatched = DeviceIdentity {class: DeviceClass::Network, ..BLOCK_IDENTITY};
    assert_eq!(registry.attach(MMIO_NETWORK, mismatched, &NETWORK_RANGE),
        Err(MatchError::UnsupportedDevice));
    assert_eq!(registry.snapshot(), before);
    let wrong_vendor = DeviceIdentity {vendor_id: 0, ..NETWORK_IDENTITY};
    assert_eq!(registry.attach(MMIO_NETWORK, wrong_vendor, &NETWORK_RANGE),
        Err(MatchError::UnsupportedVendor));
    assert_eq!(registry.snapshot(), before);
    let wrong_transport = DeviceIdentity {transport: Transport::Pci, ..NETWORK_IDENTITY};
    assert_eq!(registry.attach(MMIO_NETWORK, wrong_transport, &NETWORK_RANGE),
        Err(MatchError::KeyTransportMismatch));
    assert_eq!(registry.snapshot(), before);
    assert_eq!(registry.get(block).unwrap().1.driver, Driver::VirtioBlock);
    let network = registry.attach(MMIO_NETWORK, NETWORK_IDENTITY, &NETWORK_RANGE).unwrap();
    assert_eq!(registry.get(network).unwrap().1.driver, Driver::VirtioNetwork);
    assert_eq!((registry.binding_count(), registry.resource_count()), (2, 2));
}

#[test]
fn os_device_registry_detach_positive() {
    let mut registry = SafeRegistry::<2, 2>::new();
    let empty = registry.attach(BLOCK, BLOCK_IDENTITY, &BLOCK_RANGE, 2).unwrap();
    let full = registry.attach(MMIO_NETWORK, NETWORK_IDENTITY, &NETWORK_RANGE, 2).unwrap();
    registry.post(full, 2).unwrap();
    assert_eq!(registry.get(empty).unwrap().2.outstanding, 0);
    assert_eq!(registry.get(full).unwrap().2.outstanding, 2);
    assert_eq!((registry.binding_count(), registry.resource_count()), (2, 2));
    registry.stop_admissions(empty).unwrap();
    registry.mark_reset(empty).unwrap();
    assert_eq!(registry.drain_after_reset(empty), Ok(0));
    registry.mark_irq_detached(empty).unwrap();
    assert_eq!(registry.finish_detach(empty).unwrap().key, BLOCK);
    assert_eq!((registry.binding_count(), registry.resource_count()), (1, 1));
    registry.stop_admissions(full).unwrap();
    registry.mark_reset(full).unwrap();
    assert_eq!(registry.drain_after_reset(full), Ok(2));
    registry.mark_irq_detached(full).unwrap();
    assert_eq!(registry.finish_detach(full).unwrap().key, MMIO_NETWORK);
    assert_eq!((registry.binding_count(), registry.resource_count()), (0, 0));
    assert!(registry.snapshot().queues.iter().all(Option::is_none));
}

#[test]
fn os_device_registry_detach_negative() {
    let mut registry = SafeRegistry::<1, 1>::new();
    let handle = registry.attach(BLOCK, BLOCK_IDENTITY, &BLOCK_RANGE, 2).unwrap();
    registry.post(handle, 2).unwrap();
    let full = registry.snapshot();
    assert_eq!(registry.post(handle, 1), Err(DetachError::QueueFull));
    assert_eq!(registry.finish_detach(handle), Err(DetachError::ResetRequired));
    assert_eq!(registry.snapshot(), full);
    registry.stop_admissions(handle).unwrap();
    let stopped = registry.snapshot();
    assert_eq!(registry.post(handle, 1), Err(DetachError::AdmissionStopped));
    assert_eq!(registry.drain_after_reset(handle), Err(DetachError::ResetRequired));
    assert_eq!(registry.finish_detach(handle), Err(DetachError::ResetRequired));
    assert_eq!(registry.snapshot(), stopped);
    registry.mark_reset(handle).unwrap();
    let reset = registry.snapshot();
    assert_eq!(registry.finish_detach(handle), Err(DetachError::DrainRequired));
    assert_eq!(registry.mark_irq_detached(handle), Err(DetachError::DrainRequired));
    assert_eq!(registry.snapshot(), reset);
    assert_eq!(registry.drain_after_reset(handle), Ok(2));
    let drained = registry.snapshot();
    assert_eq!(registry.finish_detach(handle), Err(DetachError::IrqDetachRequired));
    assert_eq!(registry.snapshot(), drained);
    assert_eq!(registry.get(handle).unwrap().2.phase, DetachPhase::Drained);
    registry.mark_irq_detached(handle).unwrap();
    assert_eq!(registry.finish_detach(handle).unwrap().key, BLOCK);
    assert_eq!((registry.binding_count(), registry.resource_count()), (0, 0));
    assert!(registry.snapshot().queues.iter().all(Option::is_none));
}