/// The [`CobsDecoder`] type is used to decode a stream of bytes to a

/// given mutable output slice. This is often useful when heap data

/// structures are not available, or when not all message bytes are

/// received at a single point in time.

#[derive(Debug)]

pub struct CobsDecoder<'a> {

    /// Destination slice for decoded message

    dest: &'a mut [u8],

    /// Index of next byte to write in `dest`

    dest_idx: usize,

    /// Decoder state as an enum

    state: DecoderState,

}

/// The [`DecoderState`] is used to track the current state of a

/// streaming decoder. This struct does not contain the output buffer

/// (or a reference to one), and can be used when streaming the decoded

/// output to a custom data type.

#[derive(Debug)]

pub enum DecoderState {

    /// State machine has not received any non-zero bytes

    Idle,

    /// 1-254 bytes, can be header or 00

    Grab(u8),

    /// 255 bytes, will be a header next

    GrabChain(u8),

}

fn add(to: &mut [u8], idx: usize, data: u8) -> Result<(), DecodeError> {

    *to.get_mut(idx).ok_or(DecodeError::TargetBufTooSmall)? = data;

    Ok(())

}

/// [`DecodeResult`] represents the possible non-error outcomes of

/// pushing an encoded data byte into the [`DecoderState`] state machine

pub enum DecodeResult {

    /// The given input byte did not prompt an output byte, either because the

    /// state machine is still idle, or we have just processed a header byte.

    /// More data is needed to complete the message.

    NoData,

    /// We have received a complete and well-encoded COBS message. The

    /// contents of the associated output buffer may now be used

    DataComplete,

    /// The following byte should be appended to the current end of the decoded

    /// output buffer.

    /// More data is needed to complete the message.

    DataContinue(u8),

}

#[derive(Debug, PartialEq, Eq, thiserror::Error)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

#[cfg_attr(feature = "defmt", derive(defmt::Format))]

pub enum DecodeError {

    #[error("empty input frame")]

    EmptyFrame,

    #[error("frame with invalid format, written {decoded_bytes:?} to decoded buffer")]

    InvalidFrame {

        /// Number of bytes written to the decoded buffer.

        decoded_bytes: usize,

    },

    #[error("target buffer too small")]

    TargetBufTooSmall,

}

impl DecoderState {

    /// Push a single encoded byte into the state machine. If the input was

    /// unexpected, such as an early end of a framed message segment, an Error will

    /// be returned, and the current associated output buffer contents should be discarded.

    ///

    /// If a complete message is indicated, the decoding state machine will automatically

    /// reset itself to the Idle state, and may be used to begin decoding another message.

    ///

    /// NOTE: Sentinel value must be included in the input to this function for the

    /// decoding to complete

    pub fn feed(&mut self, data: u8) -> Result<DecodeResult, DecodeError> {

        use DecodeResult::*;

        use DecoderState::*;

        let (ret, state) = match (&self, data) {

            // Currently Idle, received a terminator, ignore, stay idle

            (Idle, 0x00) => (Ok(NoData), Idle),

            // Currently Idle, received a byte indicating the

            // next 255 bytes have no zeroes, so we will have 254 unmodified

            // data bytes, then an overhead byte

            (Idle, 0xFF) => (Ok(NoData), GrabChain(0xFE)),

            // Currently Idle, received a byte indicating there will be a

            // zero that must be modified in the next 1..=254 bytes

            (Idle, n) => (Ok(NoData), Grab(n - 1)),

            // We have reached the end of a data run indicated by an overhead

            // byte, AND we have recieved the message terminator. This was a

            // well framed message!

            (Grab(0), 0x00) => (Ok(DataComplete), Idle),

            // We have reached the end of a data run indicated by an overhead

            // byte, and the next segment of 254 bytes will have no modified

            // sentinel bytes

            (Grab(0), 0xFF) => (Ok(DataContinue(0)), GrabChain(0xFE)),

            // We have reached the end of a data run indicated by an overhead

            // byte, and we will treat this byte as a modified sentinel byte.

            // place the sentinel byte in the output, and begin processing the

            // next non-sentinel sequence

            (Grab(0), n) => (Ok(DataContinue(0)), Grab(n - 1)),

            // We were not expecting the sequence to terminate, but here we are.

            // Report an error due to early terminated message

            (Grab(_), 0) => (Err(DecodeError::InvalidFrame { decoded_bytes: 0 }), Idle),

            // We have not yet reached the end of a data run, decrement the run

            // counter, and place the byte into the decoded output

            (Grab(i), n) => (Ok(DataContinue(n)), Grab(*i - 1)),

            // We have reached the end of a data run indicated by an overhead

            // byte, AND we have recieved the message terminator. This was a

            // well framed message!

            (GrabChain(0), 0x00) => (Ok(DataComplete), Idle),

            // We have reached the end of a data run, and we will begin another

            // data run with an overhead byte expected at the end

            (GrabChain(0), 0xFF) => (Ok(NoData), GrabChain(0xFE)),

            // We have reached the end of a data run, and we will expect `n` data

            // bytes unmodified, followed by a sentinel byte that must be modified

            (GrabChain(0), n) => (Ok(NoData), Grab(n - 1)),

            // We were not expecting the sequence to terminate, but here we are.

            // Report an error due to early terminated message

            (GrabChain(_), 0) => (Err(DecodeError::InvalidFrame { decoded_bytes: 0 }), Idle),

            // We have not yet reached the end of a data run, decrement the run

            // counter, and place the byte into the decoded output

            (GrabChain(i), n) => (Ok(DataContinue(n)), GrabChain(*i - 1)),

        };

        *self = state;

        ret

    }

}

impl<'a> CobsDecoder<'a> {

    /// Create a new streaming Cobs Decoder. Provide the output buffer

    /// for the decoded message to be placed in

    pub fn new(dest: &'a mut [u8]) -> CobsDecoder<'a> {

        CobsDecoder {

            dest,

            dest_idx: 0,

            state: DecoderState::Idle,

        }

    }

    /// Push a single byte into the streaming CobsDecoder. Return values mean:

    ///

    /// * Ok(None) - State machine okay, more data needed

    /// * Ok(Some(N)) - A message of N bytes was successfully decoded

    /// * Err([DecodeError]) - Message decoding failed

    ///

    /// NOTE: Sentinel value must be included in the input to this function for the

    /// decoding to complete

    pub fn feed(&mut self, data: u8) -> Result<Option<usize>, DecodeError> {

        match self.state.feed(data) {

            Err(_) => Err(DecodeError::InvalidFrame {

                decoded_bytes: self.dest_idx,

            }),

            Ok(DecodeResult::NoData) => Ok(None),

            Ok(DecodeResult::DataContinue(n)) => {

                add(self.dest, self.dest_idx, n)?;

                self.dest_idx += 1;

                Ok(None)

            }

            Ok(DecodeResult::DataComplete) => Ok(Some(self.dest_idx)),

        }

    }

    /// Push a slice of bytes into the streaming CobsDecoder. Return values mean:

    ///

    /// * Ok(None) - State machine okay, more data needed

    /// * Ok(Some((N, M))) - A message of N bytes was successfully decoded,

    ///     using M bytes from `data` (and earlier data)

    /// * Err([DecodeError]) - Message decoding failed

    ///

    /// NOTE: Sentinel value must be included in the input to this function for the

    /// decoding to complete

    pub fn push(&mut self, data: &[u8]) -> Result<Option<(usize, usize)>, DecodeError> {

        for (consumed_idx, d) in data.iter().enumerate() {

            let opt_decoded_bytes = self.feed(*d)?;

            if let Some(decoded_bytes_ct) = opt_decoded_bytes {

                // convert from index to number of bytes consumed

                return Ok(Some((decoded_bytes_ct, consumed_idx + 1)));

            }

        }

        Ok(None)

    }

}

/// Decodes the `source` buffer into the `dest` buffer.

///

/// This function uses the typical sentinel value of 0.

///

/// # Failures

///

/// This will return `Err(())` if there was a decoding error. Otherwise,

/// it will return `Ok(n)` where `n` is the length of the decoded message.

pub fn decode(source: &[u8], dest: &mut [u8]) -> Result<usize, DecodeError> {

    if source.is_empty() {

        return Err(DecodeError::EmptyFrame);

    }

    let mut dec = CobsDecoder::new(dest);

    // Did we decode a message, using some or all of the buffer?

    if let Some((d_used, _s_used)) = dec.push(source)? {

        return Ok(d_used);

    }

    // If we consumed the entire buffer, but did NOT get a message,

    // AND the message did not end with a zero, try providing one to

    // complete the decoding.

    if source.last() != Some(&0) {

        // Explicitly push sentinel of zero

        if let Some((d_used, _s_used)) = dec.push(&[0])? {

            return Ok(d_used);

        }

    }

    // Nope, no early message, no missing terminator, just failed to decode

    Err(DecodeError::InvalidFrame {

        decoded_bytes: dec.dest_idx,

    })

}

/// A report of the source and destination bytes used during in-place decoding

#[derive(Debug, Copy, Clone)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

#[cfg_attr(feature = "defmt", derive(defmt::Format))]

pub struct DecodeReport {

    // The number of source bytes used, NOT INCLUDING the sentinel byte,

    // if there was one.

    pub src_used: usize,

    // The number of bytes of the source buffer that now include the

    // decoded result

    pub dst_used: usize,

}

/// Decodes a message in-place.

///

/// This is the same function as [decode_in_place], but provides a report

/// of both the number of source bytes consumed as well as the size of the

/// destination used.

pub fn decode_in_place_report(buf: &mut [u8]) -> Result<DecodeReport, DecodeError> {

    let mut source_index = 0;

    let mut dest_index = 0;

    if buf.is_empty() {

        return Err(DecodeError::EmptyFrame);

    }

    // Stop at the first terminator, if any

    let src_end = if let Some(end) = buf.iter().position(|b| *b == 0) {

        end

    } else {

        buf.len()

    };

    while source_index < src_end {

        let code = buf[source_index];

        if source_index + code as usize > src_end && code != 1 {

            return Err(DecodeError::InvalidFrame {

                decoded_bytes: dest_index,

            });

        }

        source_index += 1;

        for _ in 1..code {

            *buf.get_mut(dest_index)

                .ok_or(DecodeError::TargetBufTooSmall)? = buf[source_index];

            source_index += 1;

            dest_index += 1;

        }

        if 0xFF != code && source_index < src_end {

            *buf.get_mut(dest_index)

                .ok_or(DecodeError::TargetBufTooSmall)? = 0;

            dest_index += 1;

        }

    }

    Ok(DecodeReport {

        dst_used: dest_index,

        src_used: source_index,

    })

}

/// Decodes a message in-place.

///

/// This is the same function as [decode], but replaces the encoded message

/// with the decoded message instead of writing to another buffer.

///

/// The returned `usize` is the number of bytes used for the DECODED value,

/// NOT the number of source bytes consumed during decoding.

pub fn decode_in_place(buff: &mut [u8]) -> Result<usize, DecodeError> {

    decode_in_place_report(buff).map(|r| r.dst_used)

}

/// Decodes the `source` buffer into the `dest` buffer using an arbitrary sentinel value.

///

/// This is done by XOR-ing each byte of the source message with the chosen sentinel value,

/// which transforms the message into the same message encoded with a sentinel value of 0.

/// Then the regular decoding transformation is performed.

///

/// The returned `usize` is the number of bytes used for the DECODED value,

/// NOT the number of source bytes consumed during decoding.

pub fn decode_with_sentinel(

    source: &[u8],

    dest: &mut [u8],

    sentinel: u8,

) -> Result<usize, DecodeError> {

    for (x, y) in source.iter().zip(dest.iter_mut()) {

        *y = *x ^ sentinel;

    }

    decode_in_place(dest)

}

/// Decodes a message in-place using an arbitrary sentinel value.

///

/// The returned `usize` is the number of bytes used for the DECODED value,

/// NOT the number of source bytes consumed during decoding.

pub fn decode_in_place_with_sentinel(buff: &mut [u8], sentinel: u8) -> Result<usize, DecodeError> {

    for x in buff.iter_mut() {

        *x ^= sentinel;

    }

    decode_in_place(buff)

}

#[cfg(feature = "alloc")]

/// Decodes the `source` buffer into a vector.

pub fn decode_vec(source: &[u8]) -> Result<alloc::vec::Vec<u8>, DecodeError> {

    let mut decoded = alloc::vec![0; source.len()];

    let n = decode(source, &mut decoded[..])?;

    decoded.truncate(n);

    Ok(decoded)

}

#[cfg(feature = "alloc")]

/// Decodes the `source` buffer into a vector with an arbitrary sentinel value.

pub fn decode_vec_with_sentinel(

    source: &[u8],

    sentinel: u8,

) -> Result<alloc::vec::Vec<u8>, DecodeError> {

    let mut decoded = alloc::vec![0; source.len()];

    let n = decode_with_sentinel(source, &mut decoded[..], sentinel)?;

    decoded.truncate(n);

    Ok(decoded)

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn decode_malforemd() {

        let malformed_buf: [u8; 32] = [

            68, 69, 65, 68, 66, 69, 69, 70, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,

            0, 0, 0, 0, 0, 0,

        ];

        let mut dest_buf: [u8; 32] = [0; 32];

        if let Err(DecodeError::InvalidFrame { decoded_bytes }) =

            decode(&malformed_buf, &mut dest_buf)

        {

            assert_eq!(decoded_bytes, 7);

        } else {

            panic!("decoding worked when it should not have");

        }

    }

    #[test]

    fn decode_empty() {

        #[cfg(feature = "alloc")]

        matches!(decode_vec(&[]).unwrap_err(), DecodeError::EmptyFrame);

        matches!(

            decode_in_place(&mut []).unwrap_err(),

            DecodeError::EmptyFrame

        );

        matches!(

            decode(&[], &mut [0; 256]).unwrap_err(),

            DecodeError::EmptyFrame

        );

    }

    #[test]

    #[cfg(feature = "alloc")]

    fn decode_target_buf_too_small() {

        let encoded = &[3, 10, 11, 2, 12];

        let expected_decoded_len = 4;

        for i in 0..expected_decoded_len - 1 {

            let mut dest = alloc::vec![0; i];

            let result = decode(encoded, &mut dest);

            assert_eq!(result, Err(DecodeError::TargetBufTooSmall));

        }

    }

}