brec 0.6.0

use std::{
    io::{BufRead, Cursor},
    ops::RangeInclusive,
};

use crate::*;

/// Iterator over a sequence of `Slot`s, yielding the byte ranges that contain actual packet data.
///
/// Each iteration returns a `RangeInclusive<u64>` representing the location of the used data region
/// within the slot. The offset includes the size of the slot header and accounts for cumulative offset.
///
/// The iterator skips over empty slots and automatically adjusts for the internal layout.
///
/// Useful for scanning files or buffers that store serialized packets in slot-based format.
pub struct PacketsLocatorIterator<'a, I: Iterator<Item = &'a Slot>> {
    offset: u64,
    slots: I,
}

impl<'a, I: Iterator<Item = &'a Slot>> PacketsLocatorIterator<'a, I> {
    /// Creates a new `PacketsLocatorIterator` over the provided slice of slots.
    pub fn new(slots: I) -> Self {
        Self { offset: 0, slots }
    }

    /// Seeks to the specified packet index across the slots.
    ///
    /// # Arguments
    /// * `packet` - The logical packet index to locate.
    ///
    /// # Returns
    /// Returns a `RangeInclusive<u64>` representing the byte range of the packet data within the source
    ///
    /// # Errors
    /// Returns `Error::OutOfBounds` if the packet index exceeds the total number of packets available in the slots.
    /// Returns `Error::EmptySource` if there are no slots to search through.
    pub fn from(&mut self, packet: usize) -> Result<RangeInclusive<u64>, Error> {
        let mut count = 0;
        let mut target = packet;
        for (idx, slot) in self.slots.by_ref().enumerate() {
            count += 1;
            if slot.count() <= target {
                target -= slot.count();
                self.offset += slot.size() + slot.width();
                continue;
            }
            if !slot.is_used(target) {
                return Err(Error::OutOfBounds(
                    slot.count() + idx * DEFAULT_SLOT_CAPACITY,
                    packet,
                ));
            }
            let Some(packet_offset) = slot.offset_of(target) else {
                return Err(Error::OutOfBounds(
                    slot.count() + idx * DEFAULT_SLOT_CAPACITY,
                    packet,
                ));
            };
            let slot_offset = self.offset;
            self.offset += slot.size() + slot.width();
            return Ok(RangeInclusive::new(
                slot_offset + slot.size() + packet_offset,
                slot_offset + slot.size() + slot.width(),
            ));
        }
        if count == 0 {
            Err(Error::EmptySource)
        } else {
            Err(Error::OutOfBounds(count * DEFAULT_SLOT_CAPACITY, packet))
        }
    }
}

impl<'a, I: Iterator<Item = &'a Slot>> Iterator for PacketsLocatorIterator<'a, I> {
    type Item = RangeInclusive<u64>;

    /// Returns the next occupied range of packet data, or `None` if finished.
    fn next(&mut self) -> Option<Self::Item> {
        let slot = self.slots.next()?;
        let slot_width = slot.width();
        if slot_width == 0 {
            return None;
        }
        let location = RangeInclusive::new(
            self.offset + slot.size(),
            self.offset + slot_width + slot.size(),
        );
        self.offset += slot_width + slot.size();
        Some(location)
    }
}

/// An iterator over stored packets distributed across multiple slots.
///
/// `ReaderIterator` reads packets from a `Read + Seek` source (e.g. file or memory stream)
/// using a list of `Slot`s and their recorded layout. It internally tracks the position and
/// reuses an internal buffer (`Cursor<Vec<u8>>`) to efficiently read packet data.
///
/// It yields deserialized `PacketDef<B, P, Inner>` values or parsing errors.
///
/// # Type Parameters
/// - `S`: Source implementing `Read + Seek`
/// - `B`: Block type (must implement `BlockDef`)
/// - `P`: Payload container type (must implement `PayloadDef`)
/// - `Inner`: Inner payload object (must implement `PayloadInnerDef`)
pub struct ReaderIterator<
    'a,
    I: Iterator<Item = &'a Slot>,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> {
    locator: PacketsLocatorIterator<'a, I>,
    source: &'a mut S,
    buffer: Cursor<Vec<u8>>,
    ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
    _block: std::marker::PhantomData<B>,
    _payload: std::marker::PhantomData<P>,
    _payload_inner: std::marker::PhantomData<Inner>,
}

impl<
    'a,
    I: Iterator<Item = &'a Slot>,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> ReaderIterator<'a, I, S, B, P, Inner>
{
    /// Constructs a new `ReaderIterator` from the given stream and slot layout.
    pub fn new(
        source: &'a mut S,
        slots: I,
        ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
    ) -> Self {
        Self {
            locator: PacketsLocatorIterator::new(slots),
            source,
            buffer: Cursor::new(Vec::new()),
            ctx,
            _block: std::marker::PhantomData,
            _payload: std::marker::PhantomData,
            _payload_inner: std::marker::PhantomData,
        }
    }
    /// Seeks to the specified packet index across the slots.
    ///
    /// # Arguments
    /// * `packet` - The logical packet index to locate.
    ///
    /// # Returns
    /// Returns a `RangeInclusive<u64>` representing the byte range of the packet data within the source
    ///
    /// # Errors
    /// Returns `Error::OutOfBounds` if the packet index exceeds the total number of packets available in the slots.
    /// Returns `Error::EmptySource` if there are no slots to search through.
    /// Returns other errors related to IO operation
    pub fn seek(mut self, packet: usize) -> Result<Self, Error> {
        let location = self.locator.from(packet)?;
        self.source
            .seek(std::io::SeekFrom::Start(*location.start()))?;
        let size = (location.end() - location.start()) as usize;
        let mut inner = vec![0u8; size];
        self.source.read_exact(&mut inner).unwrap();
        self.buffer = Cursor::new(inner);
        Ok(self)
    }
}

impl<
    'a,
    I: Iterator<Item = &'a Slot>,
    S: std::io::Read + std::io::Write + std::io::Seek,
    B: BlockDef,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> Iterator for ReaderIterator<'a, I, S, B, P, Inner>
{
    type Item = Result<PacketDef<B, P, Inner>, Error>;

    /// Reads and yields the next packet located in the slots.
    ///
    /// Loads the slot's region into an internal buffer and calls `PacketDef::read`.
    /// Returns `None` if all slots are exhausted.
    fn next(&mut self) -> Option<Self::Item> {
        // TODO: remove unwraps and handle errors properly
        if self.buffer.fill_buf().unwrap().is_empty() {
            let location = self.locator.next()?;
            if let Err(err) = self
                .source
                .seek(std::io::SeekFrom::Start(*location.start()))
            {
                return Some(Err(err.into()));
            }
            let size = (location.end() - location.start()) as usize;
            let mut inner = vec![0u8; size];
            self.source.read_exact(&mut inner).unwrap();
            self.buffer = Cursor::new(inner);
        }
        match <PacketDef<B, P, Inner> as ReadPacketFrom>::read(&mut self.buffer, self.ctx) {
            Err(err) => Some(Err(err)),
            Ok(pkg) => Some(Ok(pkg)),
        }
    }
}

/// An iterator over packets stored in slots with rule-based filtering.
///
/// This iterator functions like `ReaderIterator`, but applies `RulesDef`-based
/// filters to decide whether to yield, skip, or reject packets. The filtering is performed
/// during parsing using `PacketDef::filtered`, which allows:
/// - prefiltering by blocks
/// - filtering by payload
/// - filtering by fully parsed packet
///
/// # Type Parameters
/// - `S`: Input stream (`Read + Seek`)
/// - `B`: Block type
/// - `BR`: Referred block type for rule filtering
/// - `P`: Payload wrapper type
/// - `Inner`: Inner payload type
pub struct ReaderFilteredIterator<
    'a,
    I: Iterator<Item = &'a Slot>,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> {
    locator: PacketsLocatorIterator<'a, I>,
    source: &'a mut S,
    rules: &'a RulesDef<B, BR, P, Inner>,
    buffer: Cursor<Vec<u8>>,
    ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
}

impl<
    'a,
    I: Iterator<Item = &'a Slot>,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> ReaderFilteredIterator<'a, I, S, B, BR, P, Inner>
{
    /// Constructs a new filtered packet iterator from the given stream, slots and rules.
    pub fn new(
        source: &'a mut S,
        slots: I,
        rules: &'a RulesDef<B, BR, P, Inner>,
        ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
    ) -> Self {
        Self {
            locator: PacketsLocatorIterator::new(slots),
            source,
            rules,
            buffer: Cursor::new(Vec::new()),
            ctx,
        }
    }
}

impl<
    'a,
    I: Iterator<Item = &'a Slot>,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> Iterator for ReaderFilteredIterator<'a, I, S, B, BR, P, Inner>
{
    type Item = Result<PacketDef<B, P, Inner>, Error>;

    /// Attempts to read and yield the next packet that passes all configured rules.
    ///
    /// Internally uses `PacketDef::filtered` to apply:
    /// - block-level filtering
    /// - payload-level filtering
    /// - full-packet filtering
    ///
    /// Returns `Some(Ok(...))` if accepted, skips if denied, and `Err(...)` on error.
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.buffer.fill_buf().unwrap().is_empty() {
                let location = self.locator.next()?;
                if let Err(err) = self
                    .source
                    .seek(std::io::SeekFrom::Start(*location.start()))
                {
                    return Some(Err(err.into()));
                }
                let size = (location.end() - location.start()) as usize;
                let mut inner = vec![0u8; size];
                self.source.read_exact(&mut inner).unwrap();
                self.buffer = Cursor::new(inner);
            }
            match PacketDef::filtered(&mut self.buffer, self.rules, self.ctx) {
                Ok(LookInStatus::Accepted(_, packet)) => return Some(Ok(packet)),
                Ok(LookInStatus::Denied(_)) => {
                    continue;
                }
                Ok(LookInStatus::NotEnoughData(needed)) => {
                    return Some(Err(Error::NotEnoughData(needed)));
                }
                Err(err) => return Some(Err(err)),
            }
        }
    }
}

/// An iterator over a specified range of packets within a `ReaderDef`.
///
/// Unlike `ReaderIterator`, this variant yields a bounded number of packets starting from a specific index.
/// It uses the internal `storage.nth(n)` method to fetch each packet by logical index.
///
/// # Type Parameters
/// - `S`: Underlying storage stream (`Read + Write + Seek`)
/// - `B`: Block type
/// - `BR`: Referred block type
/// - `P`: Payload type wrapper
/// - `Inner`: Inner payload object
pub struct ReaderRangeIterator<
    'a,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> {
    storage: &'a mut ReaderDef<S, B, BR, P, Inner>,
    len: usize,
    from: usize,
    ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
    _block: std::marker::PhantomData<B>,
    _payload: std::marker::PhantomData<P>,
    _payload_inner: std::marker::PhantomData<Inner>,
}

impl<
    'a,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> ReaderRangeIterator<'a, S, B, BR, P, Inner>
{
    pub fn new(
        storage: &'a mut ReaderDef<S, B, BR, P, Inner>,
        from: usize,
        len: usize,
        ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
    ) -> Self {
        Self {
            storage,
            len,
            from,
            ctx,
            _block: std::marker::PhantomData,
            _payload: std::marker::PhantomData,
            _payload_inner: std::marker::PhantomData,
        }
    }
}

impl<
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> Iterator for ReaderRangeIterator<'_, S, B, BR, P, Inner>
{
    type Item = Result<PacketDef<B, P, Inner>, Error>;

    /// Returns the next packet from the range by calling `storage.nth(current_index)`.
    ///
    /// Ends after yielding `len` elements or if an error occurs.
    fn next(&mut self) -> Option<Self::Item> {
        if self.len == 0 {
            return None;
        }
        let item = self.storage.nth(self.from, self.ctx);
        self.from += 1;
        self.len -= 1;
        match item {
            Ok(None) => None,
            Ok(Some(packet)) => Some(Ok(packet)),
            Err(err) => Some(Err(err)),
        }
    }
}

/// An iterator over a specific range of packets from storage with rule-based filtering.
///
/// Similar to `ReaderRangeIterator`, but each packet is passed through the configured filtering rules.
/// Internally uses `storage.nth_filtered(index)` to fetch and filter packets.
///
/// Only packets that match all rule conditions are yielded (`LookInStatus::Accepted`).
///
/// # Type Parameters
/// - `S`: Source stream with `Read + Write + Seek`
/// - `B`: Block type
/// - `BR`: Block reference type for filtering
/// - `P`: Payload container type
/// - `Inner`: Inner payload object
pub struct ReaderRangeFilteredIterator<
    'a,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> {
    storage: &'a mut ReaderDef<S, B, BR, P, Inner>,
    len: usize,
    from: usize,
    ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
}

impl<
    'a,
    S: std::io::Read + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> ReaderRangeFilteredIterator<'a, S, B, BR, P, Inner>
{
    pub fn new(
        storage: &'a mut ReaderDef<S, B, BR, P, Inner>,
        from: usize,
        len: usize,
        ctx: &'a mut <Inner as ProtocolSchema>::Context<'a>,
    ) -> Self {
        Self {
            storage,
            len,
            from,
            ctx,
        }
    }
}

impl<
    S: std::io::Read + std::io::Write + std::io::Seek,
    B: BlockDef,
    BR: BlockReferredDef<B>,
    P: PayloadDef<Inner>,
    Inner: PayloadInnerDef,
> Iterator for ReaderRangeFilteredIterator<'_, S, B, BR, P, Inner>
{
    type Item = Result<PacketDef<B, P, Inner>, Error>;

    /// Attempts to read and yield the next packet in range that passes all filtering rules.
    ///
    /// Skips over packets that are denied, stops when `len` is exhausted or `None` is returned from storage.
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.len == 0 {
                return None;
            }
            let item = self.storage.nth_filtered(self.from, self.ctx);
            self.from += 1;
            match item {
                Ok(None) => return None,
                Ok(Some(LookInStatus::Accepted(_, packet))) => {
                    self.len -= 1;
                    return Some(Ok(packet));
                }
                Ok(Some(LookInStatus::Denied(_))) => {
                    continue;
                }
                Ok(Some(LookInStatus::NotEnoughData(needed))) => {
                    return Some(Err(Error::NotEnoughData(needed)));
                }
                Err(err) => return Some(Err(err)),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn slot_with_lengths(lengths: &[u64]) -> Slot {
        let mut slot = Slot::new(lengths.to_vec(), lengths.len() as u64, [0; 4]);
        slot.overwrite_crc();
        slot
    }

    #[test]
    fn packets_locator_next_tracks_offsets_across_slots() {
        let slot_a = slot_with_lengths(&[10, 20, 0]);
        let slot_b = slot_with_lengths(&[7, 0]);
        let slots = [slot_a, slot_b];

        let mut it = PacketsLocatorIterator::new(slots.iter());

        let first = it.next().expect("first slot range");
        let second = it.next().expect("second slot range");
        assert!(it.next().is_none());

        let first_size = slots[0].size();
        assert_eq!(*first.start(), first_size);
        assert_eq!(*first.end(), first_size + slots[0].width());

        let second_base = slots[0].size() + slots[0].width();
        assert_eq!(*second.start(), second_base + slots[1].size());
        assert_eq!(
            *second.end(),
            second_base + slots[1].size() + slots[1].width()
        );
    }

    #[test]
    fn packets_locator_from_handles_valid_empty_and_oob() {
        let empty = slot_with_lengths(&[0, 0]);
        let used = slot_with_lengths(&[5, 6, 0]);
        let slots = [empty, used];
        let mut it = PacketsLocatorIterator::new(slots.iter());

        let range = it
            .from(0)
            .expect("first packet should resolve in second slot");
        let base = slots[0].size() + slots[0].width();
        assert_eq!(*range.start(), base + slots[1].size());
        assert_eq!(*range.end(), base + slots[1].size() + slots[1].width());

        let mut oob_it = PacketsLocatorIterator::new(slots.iter());
        assert!(matches!(oob_it.from(10), Err(Error::OutOfBounds(_, 10))));

        let mut no_slots = PacketsLocatorIterator::new([].iter());
        assert!(matches!(no_slots.from(0), Err(Error::EmptySource)));
    }
}