//! Bounded direct split-virtqueue descriptor chain ownership.

pub const NEXT: u16 = 1;
pub const WRITE: u16 = 2;
pub const INDIRECT: u16 = 4;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    InvalidQueue,
    InvalidChain,
    InvalidDescriptor,
    UnsupportedFlags,
    Overlap,
    Full,
    MissingChain,
    StaleHandle,
    WrongPhase,
    MismatchedCompletion,
    NotReset,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Descriptor {
    pub address: u64,
    pub length: u32,
    pub flags: u16,
    pub next: u16,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Handle { pub head: u16, pub generation: u64 }

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Phase { Prepared, Exposed, Used }

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Slot { owner: Handle, descriptor: Descriptor }

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Chain { handle: Handle, phase: Phase, writable_bytes: u64 }

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Snapshot<const N: usize, const C: usize> {
    slots: [Option<Slot>; N],
    chains: [Option<Chain>; C],
    pub count: usize,
    next_generation: u64,
}

pub struct Pool<const N: usize, const C: usize> {
    queue_size: u16,
    slots: [Option<Slot>; N],
    chains: [Option<Chain>; C],
    count: usize,
    next_generation: u64,
}

impl<const N: usize, const C: usize> Pool<N, C> {
    pub fn new(queue_size: u16) -> Result<Self, Error> {
        if queue_size == 0 || !queue_size.is_power_of_two()
            || queue_size > 1024 || usize::from(queue_size) > N {
            return Err(Error::InvalidQueue);
        }
        Ok(Self {queue_size, slots: [None; N], chains: [None; C],
            count: 0, next_generation: 1})
    }
    pub const fn count(&self) -> usize { self.count }
    pub fn snapshot(&self) -> Snapshot<N, C> {
        Snapshot {slots: self.slots, chains: self.chains,
            count: self.count, next_generation: self.next_generation}
    }
    fn chain_index(&self, handle: Handle) -> Result<usize, Error> {
        let index = self.chains.iter().position(|chain|
            chain.is_some_and(|chain| chain.handle.head == handle.head))
            .ok_or(Error::MissingChain)?;
        if self.chains[index].unwrap().handle != handle {
            return Err(Error::StaleHandle);
        }
        Ok(index)
    }
    /// `items` are in traversal order, beginning with `head`.
    pub fn claim(&mut self, head: u16,
        items: &[(u16, Descriptor)]) -> Result<Handle, Error> {
        if items.is_empty() || items.len() > usize::from(self.queue_size)
            || items[0].0 != head {
            return Err(Error::InvalidChain);
        }
        let mut seen = [false; N];
        let mut total = 0u64;
        let mut writable = 0u64;
        let mut writing = false;
        for (position, &(index, descriptor)) in items.iter().enumerate() {
            if index >= self.queue_size || seen[usize::from(index)] {
                return Err(Error::InvalidChain);
            }
            seen[usize::from(index)] = true;
            if self.slots[usize::from(index)].is_some() { return Err(Error::Overlap); }
            if descriptor.flags & INDIRECT != 0
                || descriptor.flags & !(NEXT | WRITE | INDIRECT) != 0 {
                return Err(Error::UnsupportedFlags);
            }
            if descriptor.address == 0
                || descriptor.address.checked_add(u64::from(descriptor.length)).is_none() {
                return Err(Error::InvalidDescriptor);
            }
            total = total.checked_add(u64::from(descriptor.length))
                .ok_or(Error::InvalidChain)?;
            if total > (1u64 << 32) { return Err(Error::InvalidChain); }
            if descriptor.flags & WRITE != 0 {
                writing = true;
                writable += u64::from(descriptor.length);
            } else if writing {
                return Err(Error::InvalidChain);
            }
            let next = items.get(position + 1);
            if next.is_some() != (descriptor.flags & NEXT != 0)
                || next.is_some_and(|next| descriptor.next != next.0) {
                return Err(Error::InvalidChain);
            }
        }
        let slot = self.chains.iter().position(Option::is_none).ok_or(Error::Full)?;
        let generation = self.next_generation;
        self.next_generation = generation.checked_add(1).ok_or(Error::Full)?;
        let handle = Handle {head, generation};
        for &(index, descriptor) in items {
            self.slots[usize::from(index)] = Some(Slot {owner: handle, descriptor});
        }
        self.chains[slot] = Some(Chain {handle, phase: Phase::Prepared,
            writable_bytes: writable});
        self.count += 1;
        Ok(handle)
    }
    pub fn phase(&self, handle: Handle) -> Result<Phase, Error> {
        Ok(self.chains[self.chain_index(handle)?].unwrap().phase)
    }
    /// Caller has finished descriptor/avail writes and a release barrier.
    pub fn expose(&mut self, handle: Handle) -> Result<(), Error> {
        let index = self.chain_index(handle)?;
        if self.chains[index].unwrap().phase != Phase::Prepared {
            return Err(Error::WrongPhase);
        }
        self.chains[index].as_mut().unwrap().phase = Phase::Exposed;
        Ok(())
    }
    pub fn complete(&mut self, handle: Handle, used_head: u16,
        used_length: u32) -> Result<(), Error> {
        let index = self.chain_index(handle)?;
        let chain = self.chains[index].unwrap();
        if chain.phase != Phase::Exposed { return Err(Error::WrongPhase); }
        if used_head != handle.head
            || u64::from(used_length) > chain.writable_bytes {
            return Err(Error::MismatchedCompletion);
        }
        self.chains[index].as_mut().unwrap().phase = Phase::Used;
        Ok(())
    }
    fn clear(&mut self, index: usize) {
        let handle = self.chains[index].unwrap().handle;
        for slot in &mut self.slots {
            if slot.is_some_and(|entry| entry.owner == handle) { *slot = None; }
        }
        self.chains[index] = None;
        self.count -= 1;
    }
    pub fn cancel_prepared(&mut self, handle: Handle) -> Result<(), Error> {
        let index = self.chain_index(handle)?;
        if self.chains[index].unwrap().phase != Phase::Prepared {
            return Err(Error::WrongPhase);
        }
        self.clear(index);
        Ok(())
    }
    pub fn reclaim_used(&mut self, handle: Handle) -> Result<(), Error> {
        let index = self.chain_index(handle)?;
        if self.chains[index].unwrap().phase != Phase::Used {
            return Err(Error::WrongPhase);
        }
        self.clear(index);
        Ok(())
    }
    pub fn observe_reset(&mut self, status: u8,
        queue_enable: u16) -> Result<(), Error> {
        if status != 0 || queue_enable != 0 { return Err(Error::NotReset); }
        self.slots = [None; N];
        self.chains = [None; C];
        self.count = 0;
        Ok(())
    }
}