formaty/src/formatter/format.rs

use crate::byte_stream::{bit_mask, ByteStream, ByteStreamError};
use crate::formatter::printers::print_bytes_as_array;
use byteorder::{BigEndian, ByteOrder, LittleEndian, ReadBytesExt};
use num_bigint::{BigInt, BigUint};
use serde::Deserialize;
use std::fmt::{Display, Formatter, Write};
use std::io::Cursor;
use std::string::FromUtf8Error;

pub trait LastByte: ByteOrder {
    fn last_byte(buf: &mut Vec<u8>) -> Option<&mut u8>;

    fn big_int(buf: &[u8]) -> BigInt;

    fn big_uint(buf: &[u8]) -> BigUint;

    fn big_u_int_to_bytes(big_int: BigUint) -> Vec<u8>;

    fn big_int_to_bytes(big_int: BigInt) -> Vec<u8>;
}

impl LastByte for BigEndian {
    fn last_byte(buf: &mut Vec<u8>) -> Option<&mut u8> {
        buf.first_mut()
    }

    fn big_int(buf: &[u8]) -> BigInt {
        BigInt::from_signed_bytes_be(buf)
    }

    fn big_uint(buf: &[u8]) -> BigUint {
        BigUint::from_bytes_be(buf)
    }

    fn big_u_int_to_bytes(big_int: BigUint) -> Vec<u8> {
        big_int.to_bytes_be()
    }

    fn big_int_to_bytes(big_int: BigInt) -> Vec<u8> {
        big_int.to_signed_bytes_be()
    }
}

impl LastByte for LittleEndian {
    fn last_byte(buf: &mut Vec<u8>) -> Option<&mut u8> {
        buf.last_mut()
    }

    fn big_int(buf: &[u8]) -> BigInt {
        BigInt::from_signed_bytes_le(buf)
    }

    fn big_uint(buf: &[u8]) -> BigUint {
        BigUint::from_bytes_le(buf)
    }

    fn big_u_int_to_bytes(big_int: BigUint) -> Vec<u8> {
        big_int.to_bytes_le()
    }

    fn big_int_to_bytes(big_int: BigInt) -> Vec<u8> {
        big_int.to_signed_bytes_le()
    }
}

#[derive(Debug, Clone)]
pub enum FormatError {
    ByteSteamError(ByteStreamError),
    #[allow(dead_code)]
    NotSupported,
    StringParseError(FromUtf8Error),
}

impl From<ByteStreamError> for FormatError {
    fn from(e: ByteStreamError) -> Self {
        FormatError::ByteSteamError(e)
    }
}

impl From<FromUtf8Error> for FormatError {
    fn from(e: FromUtf8Error) -> Self {
        FormatError::StringParseError(e)
    }
}

impl Display for FormatError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            FormatError::ByteSteamError(e) => writeln!(f, "Byte stream error: {}", e),
            FormatError::NotSupported => write!(f, "Field type not supported"),
            FormatError::StringParseError(e) => write!(f, "String parse error: {}", e),
        }
    }
}

#[derive(Debug, Deserialize, Clone)]
#[serde(tag = "type")]
pub enum FieldType {
    /// Unsigned Int
    UInt {
        bit_width: usize,
        endianness: Endianness,
    },
    /// Unsigned Int
    Int {
        bit_width: usize,
        endianness: Endianness,
    },
    /// Single Precession Float
    Float { endianness: Endianness },
    /// Double Precession Float
    Double { endianness: Endianness },
    /// Null Terminated String Field
    String {
        max_len: usize,
        endianness: Endianness,
    },
    /// Fixed Byte Length Field
    Bytes {
        max_len: usize,
        endianness: Endianness,
    },
}

#[derive(Debug, Deserialize, Clone, PartialOrd, PartialEq, Copy)]
pub enum Endianness {
    LittleEndian,
    BigEndian,
}

#[derive(Debug, Deserialize, Clone)]
pub struct Field {
    /// Field Name
    pub name: String,
    /// Field Type
    pub field_type: FieldType,
    /// Flip Bit Order
    pub bit_flip: Option<bool>,
}

impl Field {
    fn format_int<T: LastByte>(
        byte_stream: &ByteStream,
        bit_ndx: usize,
        bit_width: usize,
    ) -> Result<(String, usize), FormatError> {
        let mut bytes = byte_stream.get_bytes::<T>(bit_ndx, bit_width)?;

        if let Some(last_byte) = T::last_byte(&mut bytes) {
            let last_bit = ((bit_width - 1) % 8) as u8;
            let sign_bit = (*last_byte >> last_bit) & 0x1 == 1;

            if sign_bit {
                // Sign extend
                *last_byte |= !bit_mask(last_bit + 1)
            }

            let big_int = T::big_int(&bytes);

            Ok((big_int.to_string(), bit_width))
        } else {
            Err(ByteStreamError::OutOfRange.into())
        }
    }

    fn format_uint<T: LastByte>(
        byte_stream: &ByteStream,
        bit_ndx: usize,
        bit_width: usize,
    ) -> Result<(String, usize), FormatError> {
        let bytes = byte_stream.get_bytes::<T>(bit_ndx, bit_width)?;

        let big_int = T::big_uint(&bytes);
        Ok((big_int.to_string(), bit_width))
    }

    fn format_float<T: LastByte>(
        byte_stream: &ByteStream,
        bit_ndx: usize,
    ) -> Result<(String, usize), FormatError> {
        let bytes = byte_stream.get_bytes::<T>(bit_ndx, 32)?;
        let mut cursor = Cursor::new(bytes);

        Ok((cursor.read_f32::<T>().unwrap().to_string(), 4))
    }

    fn format_double<T: LastByte>(
        byte_stream: &ByteStream,
        bit_ndx: usize,
    ) -> Result<(String, usize), FormatError> {
        let bytes = byte_stream.get_bytes::<T>(bit_ndx, 64)?;
        let mut cursor = Cursor::new(bytes);

        Ok((cursor.read_f64::<T>().unwrap().to_string(), 4))
    }

    fn format_string<T: LastByte>(
        byte_stream: &ByteStream,
        mut bit_ndx: usize,
        max_byte_len: usize,
    ) -> Result<(String, usize), FormatError> {
        let mut string_bytes = Vec::new();

        for _ in 0..max_byte_len {
            let byte = byte_stream.get_bytes::<T>(bit_ndx, 8)?[0];

            if byte == 0 {
                break;
            }

            string_bytes.push(byte);
            bit_ndx += 8;
        }

        let byte_count = string_bytes.len();
        Ok((String::from_utf8(string_bytes)?, byte_count))
    }

    fn format_bytes<T: LastByte>(
        byte_stream: &ByteStream,
        bit_ndx: usize,
        max_byte_len: usize,
    ) -> Result<(String, usize), FormatError> {
        let data_remaining = byte_stream.len() - bit_ndx / 8;

        let width = if max_byte_len < data_remaining {
            max_byte_len * 8
        } else {
            data_remaining * 8
        };

        let data = byte_stream.get_bytes::<T>(bit_ndx, width)?;

        Ok((print_bytes_as_array(&data), width))
    }

    fn format_data(
        &self,
        byte_stream: &mut ByteStream,
        bit_ndx: usize,
    ) -> Result<(String, usize), FormatError> {
        let global_bit_flip = byte_stream.get_reverse_bits();

        if let Some(bit_flip) = self.bit_flip {
            byte_stream.set_reverse_bits(bit_flip);
        }

        let fmt = match self.field_type {
            FieldType::UInt {
                bit_width,
                endianness,
            } => match endianness {
                Endianness::LittleEndian => {
                    Self::format_uint::<LittleEndian>(byte_stream, bit_ndx, bit_width)
                }
                Endianness::BigEndian => {
                    Self::format_uint::<BigEndian>(byte_stream, bit_ndx, bit_width)
                }
            },
            FieldType::Int {
                bit_width,
                endianness,
            } => match endianness {
                Endianness::LittleEndian => {
                    Self::format_int::<LittleEndian>(byte_stream, bit_ndx, bit_width)
                }
                Endianness::BigEndian => {
                    Self::format_int::<BigEndian>(byte_stream, bit_ndx, bit_width)
                }
            },
            FieldType::Float { endianness } => match endianness {
                Endianness::LittleEndian => {
                    Self::format_float::<LittleEndian>(byte_stream, bit_ndx)
                }
                Endianness::BigEndian => Self::format_float::<BigEndian>(byte_stream, bit_ndx),
            },
            FieldType::Double { endianness } => match endianness {
                Endianness::LittleEndian => {
                    Self::format_double::<LittleEndian>(byte_stream, bit_ndx)
                }
                Endianness::BigEndian => Self::format_double::<BigEndian>(byte_stream, bit_ndx),
            },
            FieldType::String {
                max_len,
                endianness,
            } => match endianness {
                Endianness::LittleEndian => {
                    Self::format_string::<LittleEndian>(byte_stream, bit_ndx, max_len)
                }
                Endianness::BigEndian => {
                    Self::format_string::<LittleEndian>(byte_stream, bit_ndx, max_len)
                }
            },
            FieldType::Bytes {
                max_len,
                endianness,
            } => match endianness {
                Endianness::LittleEndian => {
                    Self::format_bytes::<LittleEndian>(byte_stream, bit_ndx, max_len)
                }
                Endianness::BigEndian => {
                    Self::format_bytes::<LittleEndian>(byte_stream, bit_ndx, max_len)
                }
            },
        }?;

        byte_stream.set_reverse_bits(global_bit_flip);

        Ok(fmt)
    }
}

#[derive(Debug, Deserialize, Clone)]
pub struct Format {
    /// Format Name
    pub name: String,
    /// Flip bits
    pub bit_flip: bool,
    /// Elements of the format
    pub fields: Vec<Field>,
}

impl Format {
    pub fn format_data(&self, data: &[u8]) -> Result<String, FormatError> {
        let mut format_str = String::new();
        let mut byte_stream = ByteStream::from(data);
        let mut bit_ndx: usize = 0;

        byte_stream.set_reverse_bits(self.bit_flip);

        for field in &self.fields {
            match field.format_data(&mut byte_stream, bit_ndx) {
                Ok((data_str, bit_width)) => {
                    bit_ndx += bit_width;
                    writeln!(format_str, "{}: {}", field.name, data_str).unwrap();
                }
                Err(e) => {
                    println!("Error formatting field: \"{}\": {}", field.name, e);
                    return Err(e);
                }
            }
        }

        Ok(format_str)
    }
}

#[cfg(test)]
mod tests {
    use crate::byte_stream::ByteStream;
    use crate::formatter::format::{Endianness, Field, FieldType};

    #[test]
    fn test_format_int_4_bits() {
        let field = Field {
            field_type: FieldType::Int {
                bit_width: 4,
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };

        for i in 0i8..7i8 {
            let mut byte_vec = ::new();
            byte_vec.push(i as u8);
            byte_vec.push((-i) as u8);

            let byte_steam = ByteStream::from(byte_vec);
            let (pos_output, _) = field.format_data(&byte_steam, 0).unwrap();
            let (neg_output, _) = field.format_data(&byte_steam, 8).unwrap();

            assert_eq!(pos_output, i.to_string());
            assert_eq!(neg_output, (-i).to_string());
        }
    }

    #[test]
    fn test_format_int_5_bits() {
        let field = Field {
            field_type: FieldType::Int {
                bit_width: 5,
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };
        let byte_steam = ByteStream::from(vec![0x1B]);
        let (output, _) = field.format_data(&byte_steam, 0).unwrap();

        assert_eq!(output, "-5")
    }

    #[test]
    fn test_format_int_16_bits() {
        let field = Field {
            field_type: FieldType::Int {
                bit_width: 16,
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(vec![0xFC, 0xA5]);
        let (output, _) = field.format_data(&byte_steam, 0).unwrap();

        assert_eq!(output, "-23044")
    }

    #[test]
    fn test_format_int_16_bits_not_aligned() {
        let field = Field {
            field_type: FieldType::Int {
                bit_width: 16,
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(vec![0xC0, 0x5F, 0x0A]);
        let (output, _) = field.format_data(&byte_steam, 4).unwrap();

        assert_eq!(output, "-23044")
    }

    #[test]
    fn test_format_float() {
        let field = Field {
            field_type: FieldType::Float {
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(b"\x52\x58\xd2\xc3".to_vec());
        let (output, _) = field.format_data(&byte_steam, 0).unwrap();

        assert_eq!(output, "-420.69")
    }

    #[test]
    fn test_format_double() {
        let field = Field {
            field_type: FieldType::Double {
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(b"\xD7\xA3\x70\x3D\x0A\x4B\x7A\xC0".to_vec());
        let (output, _) = field.format_data(&byte_steam, 0).unwrap();

        assert_eq!(output, "-420.69")
    }

    #[test]
    fn test_format_float_err() {
        let field = Field {
            field_type: FieldType::Double {
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(b"\x3D\x70\xA3\xD7".to_vec());

        assert!(field.format_data(&byte_steam, 0).is_err())
    }

    #[test]
    fn test_format_string() {
        let field = Field {
            field_type: FieldType::String { max_len: 16 },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(b"Hello World!\0".to_vec());

        let (output, _) = field.format_data(&byte_steam, 0).unwrap();

        assert_eq!(output, "Hello World!".to_string())
    }

    #[test]
    fn test_format_bytes() {
        let field = Field {
            field_type: FieldType::Bytes { max_len: 2 },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(vec![0xDE, 0xAD, 0xBE, 0xEF]);

        let (output, _) = field.format_data(&byte_steam, 0).unwrap();

        assert_eq!(output, "[222, 173]")
    }

    #[test]
    fn test_format_bytes_max_len_bigger_than_data() {
        let field = Field {
            field_type: FieldType::Bytes { max_len: 64 },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(vec![0xDE, 0xAD]);

        let (output, _) = field.format_data(&byte_steam, 0).unwrap();

        assert_eq!(output, "[222, 173]")
    }

    #[test]
    fn test_ccsds_apid_issue() {
        let field = Field {
            field_type: FieldType::UInt {
                bit_width: 11,
                endianness: Endianness::LittleEndian,
            },
            name: "test".to_string(),
            bit_flip: None,
        };

        let byte_steam = ByteStream::from(vec![
            0xe0, 0xa1, 0xc0, 0x00, 0x05, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
        ]);

        let (output, _) = field.format_data(&byte_steam, 5).unwrap();

        assert_eq!(output, "1295")
    }
}