prns-core 0.3.4

use crate::engine::CommandId;
use crate::routing::links::request::RequestId;
use crate::routing::links::resources::ResourceHash;
use crate::routing::links::LinkId;

pub trait IncomingAssemblyTable {
    fn capacity(&self) -> usize;
    fn len(&self) -> usize;

    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn link_ids(&self) -> &[LinkId];
    fn original_hashes(&self) -> &[ResourceHash];
    fn total_segments(&self) -> &[u64];
    fn segments_received(&self) -> &[u64];
    fn received_totals(&self) -> &[u64];

    fn push(&mut self, link_id: LinkId, original_hash: ResourceHash, total_segments: u64);
    fn set_progress(&mut self, index: usize, segments_received: u64, received_total: u64);
    fn swap_remove(&mut self, index: usize);
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SegmentFit {
    Expected,
    Unexpected,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AssemblyProgress {
    Assembling,
    Complete { total_size_bytes: u64 },
}

#[derive(Debug, Default)]
pub struct IncomingAssemblies<C: IncomingAssemblyTable> {
    table: C,
}

impl<C: IncomingAssemblyTable> IncomingAssemblies<C> {
    /// Open a chain on `link_id`; any prior chain is replaced. A link reassembles one transfer at a time, the same one-resource-per-link invariant [`IncomingResources`](super::super::table::IncomingResources) keeps.
    pub fn begin(&mut self, link_id: LinkId, original_hash: ResourceHash, total_segments: u64) {
        if let Some(index) = self.index_of(&link_id) {
            self.table.swap_remove(index);
        }
        if self.table.len() < self.table.capacity() {
            self.table.push(link_id, original_hash, total_segments);
        }
    }

    pub fn fit(
        &self,
        link_id: &LinkId,
        original_hash: &ResourceHash,
        segment_index: u64,
    ) -> SegmentFit {
        let matches = self.index_of(link_id).is_some_and(|index| {
            self.table.original_hashes()[index] == *original_hash
                && segment_index == self.table.segments_received()[index] + 1
        });
        if matches {
            SegmentFit::Expected
        } else {
            SegmentFit::Unexpected
        }
    }

    pub fn advance(&mut self, link_id: &LinkId, segment_bytes: u64) -> Option<AssemblyProgress> {
        let index = self.index_of(link_id)?;
        let segments_received = self.table.segments_received()[index] + 1;
        let received_total = self.table.received_totals()[index].saturating_add(segment_bytes);
        self.table
            .set_progress(index, segments_received, received_total);
        if segments_received >= self.table.total_segments()[index] {
            Some(AssemblyProgress::Complete {
                total_size_bytes: received_total,
            })
        } else {
            Some(AssemblyProgress::Assembling)
        }
    }

    pub fn original_hash(&self, link_id: &LinkId) -> Option<ResourceHash> {
        self.index_of(link_id)
            .map(|index| self.table.original_hashes()[index])
    }

    pub fn clear(&mut self, link_id: &LinkId) {
        if let Some(index) = self.index_of(link_id) {
            self.table.swap_remove(index);
        }
    }

    fn index_of(&self, link_id: &LinkId) -> Option<usize> {
        self.table
            .link_ids()
            .iter()
            .position(|candidate| candidate == link_id)
    }
}

pub trait OutgoingAssemblyTable {
    fn capacity(&self) -> usize;
    fn len(&self) -> usize;

    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn link_ids(&self) -> &[LinkId];
    fn original_hashes(&self) -> &[ResourceHash];

    fn supports_static_continuations(&self) -> bool {
        false
    }

    fn static_continuation(&self, _index: usize) -> Option<StaticResponseContinuation> {
        None
    }

    fn set_static_continuation(
        &mut self,
        _index: usize,
        _continuation: StaticResponseContinuation,
    ) -> bool {
        false
    }

    fn push(&mut self, link_id: LinkId, original_hash: ResourceHash);
    fn swap_remove(&mut self, index: usize);
}

/// The flash-backed source for the next automatically segmented response window.
///
/// Only offsets and a static borrow survive between proofs. The live encrypted resource owns at
/// most one segment, so the complete file is never copied into the transfer store.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StaticResponseContinuation {
    pub command_id: CommandId,
    pub request_id: RequestId,
    pub bytes: &'static [u8],
    pub next_offset: usize,
    pub next_segment_index: u64,
    pub total_segments: u64,
    pub total_data_bytes: u64,
    pub metadata_packed_len: u32,
    pub segment_stream_bytes: usize,
}

#[derive(Debug, Default)]
pub struct OutgoingAssemblies<C: OutgoingAssemblyTable> {
    table: C,
}

impl<C: OutgoingAssemblyTable> OutgoingAssemblies<C> {
    /// Record the `original_hash` the chain's first segment minted so every later segment advertises the same one. Any prior chain on the link is replaced (one outgoing transfer per link).
    pub fn begin(&mut self, link_id: LinkId, original_hash: ResourceHash) {
        if let Some(index) = self.index_of(&link_id) {
            self.table.swap_remove(index);
        }
        if self.table.len() < self.table.capacity() {
            self.table.push(link_id, original_hash);
        }
    }

    pub fn supports_static_continuations(&self) -> bool {
        self.table.supports_static_continuations()
    }

    pub fn set_static_continuation(
        &mut self,
        link_id: &LinkId,
        continuation: StaticResponseContinuation,
    ) -> bool {
        let Some(index) = self.index_of(link_id) else {
            return false;
        };
        self.table.set_static_continuation(index, continuation)
    }

    pub fn static_continuation(&self, link_id: &LinkId) -> Option<StaticResponseContinuation> {
        self.index_of(link_id)
            .and_then(|index| self.table.static_continuation(index))
    }

    pub fn original_hash(&self, link_id: &LinkId) -> Option<ResourceHash> {
        self.index_of(link_id)
            .map(|index| self.table.original_hashes()[index])
    }

    pub fn clear(&mut self, link_id: &LinkId) {
        if let Some(index) = self.index_of(link_id) {
            self.table.swap_remove(index);
        }
    }

    fn index_of(&self, link_id: &LinkId) -> Option<usize> {
        self.table
            .link_ids()
            .iter()
            .position(|candidate| candidate == link_id)
    }
}

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

    fn link(byte: u8) -> LinkId {
        LinkId::new([byte; 16])
    }

    fn hash(byte: u8) -> ResourceHash {
        ResourceHash::new([byte; 32])
    }

    fn table() -> IncomingAssemblies<FixedIncomingAssemblyTable<4>> {
        IncomingAssemblies::default()
    }

    #[test]
    fn advance_assembles_until_the_last_segment_completes() {
        let mut assemblies = table();
        assemblies.begin(link(1), hash(0xA), 3);
        assert_eq!(
            assemblies.advance(&link(1), 100),
            Some(AssemblyProgress::Assembling)
        );
        assert_eq!(
            assemblies.advance(&link(1), 100),
            Some(AssemblyProgress::Assembling)
        );
        assert_eq!(
            assemblies.advance(&link(1), 50),
            Some(AssemblyProgress::Complete {
                total_size_bytes: 250
            })
        );
    }

    #[test]
    fn fit_expects_the_next_segment_of_the_right_chain() {
        let mut assemblies = table();
        assemblies.begin(link(1), hash(0xA), 3);
        assemblies.advance(&link(1), 100);
        assert_eq!(
            assemblies.fit(&link(1), &hash(0xA), 2),
            SegmentFit::Expected
        );
        assert_eq!(
            assemblies.fit(&link(1), &hash(0xA), 3),
            SegmentFit::Unexpected
        );
        assert_eq!(
            assemblies.fit(&link(1), &hash(0xB), 2),
            SegmentFit::Unexpected
        );
        assert_eq!(
            assemblies.fit(&link(2), &hash(0xA), 2),
            SegmentFit::Unexpected
        );
    }

    #[test]
    fn clear_retires_the_chain() {
        let mut assemblies = table();
        assemblies.begin(link(1), hash(0xA), 3);
        assemblies.clear(&link(1));
        assert_eq!(assemblies.advance(&link(1), 100), None);
        assert_eq!(assemblies.original_hash(&link(1)), None);
    }

    #[test]
    fn begin_replaces_a_prior_chain_on_the_same_link() {
        let mut assemblies = table();
        assemblies.begin(link(1), hash(0xA), 2);
        assemblies.begin(link(1), hash(0xB), 3);
        assert_eq!(assemblies.original_hash(&link(1)), Some(hash(0xB)));
    }
}