//! Bounded logical device identities. Transport resource ownership is outside
//! this table; callers release mappings, IRQs, and DMA before detaching.

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Transport {
    Mmio,
    Pci,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DeviceKey {
    pub transport: Transport,
    pub location: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeviceKind {
    Block,
    Network,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Device {
    pub key: DeviceKey,
    pub kind: DeviceKind,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DeviceHandle {
    slot: usize,
    generation: u64,
}

impl DeviceHandle {
    pub const fn slot(self) -> usize { self.slot }
    pub const fn generation(self) -> u64 { self.generation }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    AlreadyAttached,
    Full,
    GenerationExhausted,
    StaleHandle,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Slot {
    generation: u64,
    device: Option<Device>,
}

impl Slot {
    const fn empty() -> Self { Self { generation: 0, device: None } }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Snapshot<const N: usize> {
    pub generations: [u64; N],
    pub devices: [Option<Device>; N],
    pub live_count: usize,
}

pub struct Registry<const N: usize> {
    slots: [Slot; N],
    live_count: usize,
}

impl<const N: usize> Registry<N> {
    pub const fn new() -> Self {
        Self { slots: [Slot::empty(); N], live_count: 0 }
    }

    pub const fn live_count(&self) -> usize { self.live_count }

    #[cfg(test)]
    pub fn with_generation_seeds(generations: [u64; N]) -> Self {
        let mut registry = Self::new();
        for (slot, generation) in registry.slots.iter_mut().zip(generations) {
            slot.generation = generation;
        }
        registry
    }

    pub fn snapshot(&self) -> Snapshot<N> {
        Snapshot {
            generations: self.slots.map(|slot| slot.generation),
            devices: self.slots.map(|slot| slot.device),
            live_count: self.live_count,
        }
    }

    pub fn attach(&mut self, key: DeviceKey, kind: DeviceKind) -> Result<DeviceHandle, Error> {
        if self.slots.iter().any(|slot| slot.device.is_some_and(|d| d.key == key)) {
            return Err(Error::AlreadyAttached);
        }
        let mut exhausted = false;
        for (index, slot) in self.slots.iter_mut().enumerate() {
            if slot.device.is_some() { continue; }
            if let Some(next) = slot.generation.checked_add(1) {
                slot.generation = next;
                slot.device = Some(Device { key, kind });
                self.live_count += 1;
                return Ok(DeviceHandle { slot: index, generation: next });
            }
            exhausted = true;
        }
        Err(if exhausted { Error::GenerationExhausted } else { Error::Full })
    }

    pub fn get(&self, handle: DeviceHandle) -> Result<Device, Error> {
        let slot = self.slots.get(handle.slot).ok_or(Error::StaleHandle)?;
        if slot.generation != handle.generation { return Err(Error::StaleHandle); }
        slot.device.ok_or(Error::StaleHandle)
    }

    pub fn detach(&mut self, handle: DeviceHandle) -> Result<Device, Error> {
        let slot = self.slots.get_mut(handle.slot).ok_or(Error::StaleHandle)?;
        if slot.generation != handle.generation { return Err(Error::StaleHandle); }
        let device = slot.device.take().ok_or(Error::StaleHandle)?;
        self.live_count -= 1;
        Ok(device)
    }
}

impl<const N: usize> Default for Registry<N> {
    fn default() -> Self { Self::new() }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ResourceKind {
    Mmio,
    Irq,
    Dma,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ResourceRange {
    pub kind: ResourceKind,
    pub start: u64,
    pub length: u64,
}

impl ResourceRange {
    fn end(self) -> Option<u64> {
        if self.length == 0 { None } else { self.start.checked_add(self.length) }
    }

    fn overlaps(self, other: Self) -> bool {
        self.kind == other.kind && self.start < other.start + other.length
            && other.start < self.start + self.length
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ResourceClaim {
    pub range: ResourceRange,
    pub owner: DeviceHandle,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ResourceError {
    EmptyResources,
    InvalidRange,
    Conflict,
    Capacity,
    Device(Error),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ResourceSnapshot<const D: usize, const R: usize> {
    pub devices: Snapshot<D>,
    pub claims: [Option<ResourceClaim>; R],
    pub claim_count: usize,
}

pub struct ResourceRegistry<const D: usize, const R: usize> {
    devices: Registry<D>,
    claims: [Option<ResourceClaim>; R],
    claim_count: usize,
}

impl<const D: usize, const R: usize> ResourceRegistry<D, R> {
    pub const fn new() -> Self {
        Self { devices: Registry::new(), claims: [None; R], claim_count: 0 }
    }

    pub const fn device_count(&self) -> usize { self.devices.live_count() }
    pub const fn claim_count(&self) -> usize { self.claim_count }

    pub fn snapshot(&self) -> ResourceSnapshot<D, R> {
        ResourceSnapshot { devices: self.devices.snapshot(), claims: self.claims,
            claim_count: self.claim_count }
    }

    pub fn get(&self, handle: DeviceHandle) -> Result<Device, ResourceError> {
        self.devices.get(handle).map_err(ResourceError::Device)
    }

    pub fn attach(&mut self, key: DeviceKey, kind: DeviceKind,
        ranges: &[ResourceRange]) -> Result<DeviceHandle, ResourceError> {
        if ranges.is_empty() { return Err(ResourceError::EmptyResources); }
        if ranges.len() > R - self.claim_count { return Err(ResourceError::Capacity); }
        for (index, &range) in ranges.iter().enumerate() {
            if range.end().is_none() { return Err(ResourceError::InvalidRange); }
            for &earlier in &ranges[..index] {
                if earlier.overlaps(range) { return Err(ResourceError::Conflict); }
            }
            for existing in self.claims.iter().flatten() {
                if existing.range.overlaps(range) { return Err(ResourceError::Conflict); }
            }
        }
        let handle = self.devices.attach(key, kind).map_err(ResourceError::Device)?;
        let mut incoming = ranges.iter();
        for slot in &mut self.claims {
            if slot.is_none() {
                if let Some(&range) = incoming.next() {
                    *slot = Some(ResourceClaim { range, owner: handle });
                }
            }
        }
        self.claim_count += ranges.len();
        Ok(handle)
    }

    pub fn detach(&mut self, handle: DeviceHandle) -> Result<Device, ResourceError> {
        self.devices.get(handle).map_err(ResourceError::Device)?;
        let device = self.devices.detach(handle).map_err(ResourceError::Device)?;
        for slot in &mut self.claims {
            if slot.is_some_and(|claim| claim.owner == handle) {
                *slot = None;
                self.claim_count -= 1;
            }
        }
        Ok(device)
    }
}

impl<const D: usize, const R: usize> Default for ResourceRegistry<D, R> {
    fn default() -> Self { Self::new() }
}

pub const QEMU_VIRTIO_VENDOR: u32 = 0x554d_4551;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeviceClass {
    Block,
    Network,
    Other(u32),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DeviceIdentity {
    pub transport: Transport,
    pub vendor_id: u32,
    pub device_id: u32,
    pub class: DeviceClass,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Driver {
    VirtioBlock,
    VirtioNetwork,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DriverMatch {
    pub driver: Driver,
    pub kind: DeviceKind,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MatchError {
    UnsupportedTransport,
    UnsupportedVendor,
    UnsupportedClass,
    UnsupportedDevice,
    KeyTransportMismatch,
    Resource(ResourceError),
}

struct MatchRule {
    device_id: u32,
    class: DeviceClass,
    driver: Driver,
    kind: DeviceKind,
}

const MATCH_RULES: [MatchRule; 2] = [
    MatchRule { device_id: 2, class: DeviceClass::Block,
        driver: Driver::VirtioBlock, kind: DeviceKind::Block },
    MatchRule { device_id: 1, class: DeviceClass::Network,
        driver: Driver::VirtioNetwork, kind: DeviceKind::Network },
];

pub fn match_driver(identity: DeviceIdentity) -> Result<DriverMatch, MatchError> {
    if identity.transport != Transport::Mmio { return Err(MatchError::UnsupportedTransport); }
    if identity.vendor_id != QEMU_VIRTIO_VENDOR { return Err(MatchError::UnsupportedVendor); }
    if matches!(identity.class, DeviceClass::Other(_)) { return Err(MatchError::UnsupportedClass); }
    for rule in MATCH_RULES {
        if rule.device_id == identity.device_id && rule.class == identity.class {
            return Ok(DriverMatch { driver: rule.driver, kind: rule.kind });
        }
    }
    Err(MatchError::UnsupportedDevice)
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DriverBinding {
    pub handle: DeviceHandle,
    pub identity: DeviceIdentity,
    pub driver: Driver,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MatchedSnapshot<const D: usize, const R: usize> {
    pub resources: ResourceSnapshot<D, R>,
    pub bindings: [Option<DriverBinding>; D],
    pub binding_count: usize,
}

pub struct MatchedRegistry<const D: usize, const R: usize> {
    resources: ResourceRegistry<D, R>,
    bindings: [Option<DriverBinding>; D],
    binding_count: usize,
}

impl<const D: usize, const R: usize> MatchedRegistry<D, R> {
    pub const fn new() -> Self {
        Self { resources: ResourceRegistry::new(), bindings: [None; D], binding_count: 0 }
    }

    pub const fn binding_count(&self) -> usize { self.binding_count }
    pub const fn resource_count(&self) -> usize { self.resources.claim_count() }

    pub fn snapshot(&self) -> MatchedSnapshot<D, R> {
        MatchedSnapshot { resources: self.resources.snapshot(), bindings: self.bindings,
            binding_count: self.binding_count }
    }

    pub fn attach(&mut self, key: DeviceKey, identity: DeviceIdentity,
        ranges: &[ResourceRange]) -> Result<DeviceHandle, MatchError> {
        if key.transport != identity.transport { return Err(MatchError::KeyTransportMismatch); }
        let found = match_driver(identity)?;
        let handle = self.resources.attach(key, found.kind, ranges)
            .map_err(MatchError::Resource)?;
        self.bindings[handle.slot] = Some(DriverBinding {
            handle, identity, driver: found.driver });
        self.binding_count += 1;
        Ok(handle)
    }

    pub fn get(&self, handle: DeviceHandle) -> Result<(Device, DriverBinding), MatchError> {
        let device = self.resources.get(handle).map_err(MatchError::Resource)?;
        let binding = self.bindings.get(handle.slot).and_then(|slot| *slot)
            .filter(|binding| binding.handle == handle)
            .ok_or(MatchError::Resource(ResourceError::Device(Error::StaleHandle)))?;
        Ok((device, binding))
    }

    pub fn detach(&mut self, handle: DeviceHandle) -> Result<Device, MatchError> {
        self.get(handle)?;
        let device = self.resources.detach(handle).map_err(MatchError::Resource)?;
        self.bindings[handle.slot] = None;
        self.binding_count -= 1;
        Ok(device)
    }
}

impl<const D: usize, const R: usize> Default for MatchedRegistry<D, R> {
    fn default() -> Self { Self::new() }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DetachPhase {
    Running,
    Stopped,
    Reset,
    Drained,
    IrqDetached,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct QueueOwnership {
    pub handle: DeviceHandle,
    pub capacity: u16,
    pub outstanding: u16,
    pub phase: DetachPhase,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DetachError {
    InvalidCapacity,
    QueueFull,
    AdmissionStopped,
    ResetRequired,
    DrainRequired,
    IrqDetachRequired,
    WrongPhase,
    Device(MatchError),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SafeSnapshot<const D: usize, const R: usize> {
    pub matched: MatchedSnapshot<D, R>,
    pub queues: [Option<QueueOwnership>; D],
}

pub struct SafeRegistry<const D: usize, const R: usize> {
    matched: MatchedRegistry<D, R>,
    queues: [Option<QueueOwnership>; D],
}

impl<const D: usize, const R: usize> SafeRegistry<D, R> {
    pub const fn new() -> Self {
        Self { matched: MatchedRegistry::new(), queues: [None; D] }
    }

    pub const fn binding_count(&self) -> usize { self.matched.binding_count() }
    pub const fn resource_count(&self) -> usize { self.matched.resource_count() }

    pub fn snapshot(&self) -> SafeSnapshot<D, R> {
        SafeSnapshot { matched: self.matched.snapshot(), queues: self.queues }
    }

    pub fn attach(&mut self, key: DeviceKey, identity: DeviceIdentity,
        ranges: &[ResourceRange], capacity: u16) -> Result<DeviceHandle, DetachError> {
        if capacity == 0 { return Err(DetachError::InvalidCapacity); }
        let handle = self.matched.attach(key, identity, ranges).map_err(DetachError::Device)?;
        self.queues[handle.slot] = Some(QueueOwnership {
            handle, capacity, outstanding: 0, phase: DetachPhase::Running });
        Ok(handle)
    }

    pub fn get(&self, handle: DeviceHandle) -> Result<(Device, DriverBinding, QueueOwnership), DetachError> {
        let (device, binding) = self.matched.get(handle).map_err(DetachError::Device)?;
        let queue = self.queues.get(handle.slot).and_then(|slot| *slot)
            .filter(|queue| queue.handle == handle)
            .ok_or(DetachError::Device(MatchError::Resource(
                ResourceError::Device(Error::StaleHandle))))?;
        Ok((device, binding, queue))
    }

    fn queue_mut(&mut self, handle: DeviceHandle) -> Result<&mut QueueOwnership, DetachError> {
        self.matched.get(handle).map_err(DetachError::Device)?;
        self.queues.get_mut(handle.slot).and_then(Option::as_mut)
            .filter(|queue| queue.handle == handle)
            .ok_or(DetachError::Device(MatchError::Resource(
                ResourceError::Device(Error::StaleHandle))))
    }

    pub fn post(&mut self, handle: DeviceHandle, descriptors: u16) -> Result<(), DetachError> {
        let queue = self.queue_mut(handle)?;
        if queue.phase != DetachPhase::Running { return Err(DetachError::AdmissionStopped); }
        let next = queue.outstanding.checked_add(descriptors).ok_or(DetachError::QueueFull)?;
        if next > queue.capacity { return Err(DetachError::QueueFull); }
        queue.outstanding = next;
        Ok(())
    }

    pub fn stop_admissions(&mut self, handle: DeviceHandle) -> Result<(), DetachError> {
        let queue = self.queue_mut(handle)?;
        if queue.phase != DetachPhase::Running { return Err(DetachError::WrongPhase); }
        queue.phase = DetachPhase::Stopped;
        Ok(())
    }

    pub fn mark_reset(&mut self, handle: DeviceHandle) -> Result<(), DetachError> {
        let queue = self.queue_mut(handle)?;
        if queue.phase != DetachPhase::Stopped { return Err(DetachError::WrongPhase); }
        queue.phase = DetachPhase::Reset;
        Ok(())
    }

    pub fn drain_after_reset(&mut self, handle: DeviceHandle) -> Result<u16, DetachError> {
        let queue = self.queue_mut(handle)?;
        if queue.phase != DetachPhase::Reset { return Err(DetachError::ResetRequired); }
        let cancelled = queue.outstanding;
        queue.outstanding = 0;
        queue.phase = DetachPhase::Drained;
        Ok(cancelled)
    }

    pub fn mark_irq_detached(&mut self, handle: DeviceHandle) -> Result<(), DetachError> {
        let queue = self.queue_mut(handle)?;
        if queue.phase != DetachPhase::Drained { return Err(DetachError::DrainRequired); }
        queue.phase = DetachPhase::IrqDetached;
        Ok(())
    }

    pub fn finish_detach(&mut self, handle: DeviceHandle) -> Result<Device, DetachError> {
        let (_, _, queue) = self.get(handle)?;
        let _ = queue;
        let device = self.matched.detach(handle).map_err(DetachError::Device)?;
        self.queues[handle.slot] = None;
        Ok(device)
    }
}

impl<const D: usize, const R: usize> Default for SafeRegistry<D, R> {
    fn default() -> Self { Self::new() }
}