haematite 0.6.1

use super::{CarrierRefusal, OpaqueBytes};

/// Maximum opaque payload size in one carrier envelope (16 MiB).
pub const MAX_CARRIER_FRAME_BYTES: usize = 16_777_216;
const LENGTH_PREFIX_BYTES: usize = 4;

/// Encodes one opaque payload as a big-endian length followed by its bytes.
pub fn encode_frame(payload: &[u8]) -> Result<Vec<u8>, CarrierRefusal> {
    if payload.len() > MAX_CARRIER_FRAME_BYTES {
        let announced = u32::try_from(payload.len()).unwrap_or(u32::MAX);
        return Err(CarrierRefusal::FrameTooLarge {
            announced,
            maximum: MAX_CARRIER_FRAME_BYTES,
        });
    }

    let announced = u32::try_from(payload.len()).unwrap_or(u32::MAX);
    let encoded_len = payload.len() + LENGTH_PREFIX_BYTES;
    let mut encoded = Vec::with_capacity(encoded_len);
    encoded.extend_from_slice(&announced.to_be_bytes());
    encoded.extend_from_slice(payload);
    Ok(encoded)
}

/// Decodes exactly one complete envelope.
///
/// All length and shape checks happen before payload allocation.
pub fn decode_envelope(envelope: &[u8]) -> Result<OpaqueBytes, CarrierRefusal> {
    if envelope.len() < LENGTH_PREFIX_BYTES {
        return Err(CarrierRefusal::TruncatedLengthPrefix {
            received: envelope.len(),
        });
    }

    let prefix = [envelope[0], envelope[1], envelope[2], envelope[3]];
    let announced = u32::from_be_bytes(prefix);
    if u64::from(announced) > MAX_CARRIER_FRAME_BYTES as u64 {
        return Err(CarrierRefusal::FrameTooLarge {
            announced,
            maximum: MAX_CARRIER_FRAME_BYTES,
        });
    }

    let received = envelope.len() - LENGTH_PREFIX_BYTES;
    let announced_len = announced as usize;
    if received < announced_len {
        return Err(CarrierRefusal::TruncatedFrame {
            announced,
            received,
        });
    }
    if received > announced_len {
        return Err(CarrierRefusal::TrailingFrameBytes {
            announced,
            received,
        });
    }

    Ok(OpaqueBytes::from(envelope[LENGTH_PREFIX_BYTES..].to_vec()))
}

/// Fatal incremental decoder for a stream of carrier envelopes.
#[derive(Debug)]
pub struct FrameDecoder {
    prefix: [u8; LENGTH_PREFIX_BYTES],
    prefix_received: usize,
    announced: Option<u32>,
    payload: Vec<u8>,
    frames_decoded: usize,
    payload_allocations: usize,
    refusal: Option<CarrierRefusal>,
}

impl Default for FrameDecoder {
    fn default() -> Self {
        Self::new()
    }
}

impl FrameDecoder {
    /// Creates an empty incremental decoder.
    pub const fn new() -> Self {
        Self {
            prefix: [0; LENGTH_PREFIX_BYTES],
            prefix_received: 0,
            announced: None,
            payload: Vec::new(),
            frames_decoded: 0,
            payload_allocations: 0,
            refusal: None,
        }
    }

    /// Feeds a stream fragment and returns every newly completed payload.
    ///
    /// Once an error is returned, every later call returns the same error.
    pub fn push(&mut self, mut input: &[u8]) -> Result<Vec<OpaqueBytes>, CarrierRefusal> {
        if let Some(refusal) = self.refusal {
            return Err(refusal);
        }

        let mut frames = Vec::new();
        while !input.is_empty() {
            if self.announced.is_none() {
                let needed = LENGTH_PREFIX_BYTES - self.prefix_received;
                let taken = needed.min(input.len());
                let end = self.prefix_received + taken;
                self.prefix[self.prefix_received..end].copy_from_slice(&input[..taken]);
                self.prefix_received = end;
                input = &input[taken..];

                if self.prefix_received < LENGTH_PREFIX_BYTES {
                    continue;
                }

                let announced = u32::from_be_bytes(self.prefix);
                if u64::from(announced) > MAX_CARRIER_FRAME_BYTES as u64 {
                    return self.fail(CarrierRefusal::FrameTooLarge {
                        announced,
                        maximum: MAX_CARRIER_FRAME_BYTES,
                    });
                }

                self.announced = Some(announced);
                let announced_len = announced as usize;
                if announced_len > 0 {
                    self.payload.reserve_exact(announced_len);
                    self.payload_allocations += 1;
                }
                if announced_len == 0 {
                    frames.push(OpaqueBytes::default());
                    self.finish_frame();
                }
                continue;
            }

            let announced = self.announced.unwrap_or_default() as usize;
            let remaining = announced - self.payload.len();
            let taken = remaining.min(input.len());
            self.payload.extend_from_slice(&input[..taken]);
            input = &input[taken..];

            if self.payload.len() == announced {
                let payload = std::mem::take(&mut self.payload);
                frames.push(OpaqueBytes::from(payload));
                self.finish_frame();
            }
        }

        Ok(frames)
    }

    /// Marks the input stream complete and reports a partial envelope fatally.
    pub const fn finish(&mut self) -> Result<(), CarrierRefusal> {
        if let Some(refusal) = self.refusal {
            return Err(refusal);
        }

        if let Some(announced) = self.announced {
            return self.fail(CarrierRefusal::TruncatedFrame {
                announced,
                received: self.payload.len(),
            });
        }
        if self.prefix_received > 0 || self.frames_decoded == 0 {
            return self.fail(CarrierRefusal::TruncatedLengthPrefix {
                received: self.prefix_received,
            });
        }
        Ok(())
    }

    /// Reports payload allocation operations, for bound verification.
    pub const fn payload_allocation_count(&self) -> usize {
        self.payload_allocations
    }

    /// Reports payload capacity currently owned by the decoder.
    pub const fn buffered_payload_capacity(&self) -> usize {
        self.payload.capacity()
    }

    /// Returns the fatal refusal, if one has occurred.
    pub const fn refusal(&self) -> Option<CarrierRefusal> {
        self.refusal
    }

    const fn finish_frame(&mut self) {
        self.prefix = [0; LENGTH_PREFIX_BYTES];
        self.prefix_received = 0;
        self.announced = None;
        self.frames_decoded += 1;
    }

    const fn fail<T>(&mut self, refusal: CarrierRefusal) -> Result<T, CarrierRefusal> {
        self.refusal = Some(refusal);
        Err(refusal)
    }
}