byteflow 0.2.1

use std::slice;

use num_traits::PrimInt;

use crate::{
    common::{Endianness, SerializerError},
    num_traits::SerializableInt,
};

pub struct Deserializer<'a, T: std::io::Read> {
    data: &'a mut T,
    endianness: Endianness,
    offset_in_bits: usize,
    buffer: u8,
}

impl<'a, T: std::io::Read> Deserializer<'a, T> {
    pub fn new(data: &'a mut T, endianness: Endianness) -> Deserializer<'a, T> {
        Self {
            data,
            endianness,
            offset_in_bits: 0,
            buffer: 0,
        }
    }

    pub fn offset_in_bits(&self) -> usize {
        self.offset_in_bits
    }

    pub fn skip(&mut self, mut bit_count: usize) -> Result<(), SerializerError> {
        let start_offset = self.offset_in_bits % 8;

        // Case 1: Start not byte-aligned
        if start_offset != 0 {
            let bits_to_consume = (8 - start_offset).min(bit_count);
            self.offset_in_bits += bits_to_consume;
            bit_count -= bits_to_consume;
        }

        // Case 2: When start aligned, read as many whole bytes as possible from data (and discard them)
        let bytes_to_skip = bit_count / 8;
        for _ in 0..bytes_to_skip {
            let mut buf = [0u8; 1];
            self.data.read_exact(&mut buf)?;
            self.offset_in_bits += 8;
        }

        // Case 3: If end is not byte-aligned
        let remaining_bits = bit_count % 8;
        if remaining_bits != 0 {
            self.data.read_exact(slice::from_mut(&mut self.buffer))?;
        }
        self.offset_in_bits += remaining_bits;

        Ok(())
    }

    fn read_u8(&mut self, bit_size: u8) -> Result<u8, SerializerError> {
        if bit_size > 8 {
            return Err(SerializerError::TooManyBits);
        }

        // If we read 8 aligned bits, we can read a byte
        if bit_size == 8 && self.offset_in_bits % 8 == 0 {
            self.data.read_exact(slice::from_mut(&mut self.buffer))?;
            self.offset_in_bits += 8;
            Ok(self.buffer)
        } else {
            let mut value = 0;

            for i in self.offset_in_bits..(self.offset_in_bits + bit_size as usize) {
                // If next bit is part of a new byte, read another byte from the reader
                if i % 8 == 0 {
                    self.data.read_exact(slice::from_mut(&mut self.buffer))?;
                }

                // Make space for the next bit by shifting left the previous bits
                value <<= 1;

                let bit_check_mask = 1 << (7 - (i % 8));

                // Set bit if it's set in the input
                if (self.buffer & bit_check_mask) != 0 {
                    value |= 1;
                }
            }

            self.offset_in_bits += bit_size as usize;
            Ok(value)
        }
    }

    fn sign_extend<V>(value: V, bit_size: u8) -> V
    where
        V: SerializableInt
            + num_traits::PrimInt
            + std::ops::Sub<V, Output = V>
            + std::ops::Shl<V, Output = V>,
    {
        // If zero needs 1 bit, it is signed, so check if it should be negative
        if V::zero().min_bit_count() == 1
            && value.leading_zeros() == V::bit_size() - bit_size as u32
            && bit_size != V::bit_size() as u8
        {
            let bits_to_add: V = ((V::one()
                << (V::from(V::bit_size()).unwrap() - V::from(bit_size).unwrap()))
                - V::one())
                << V::from(bit_size).unwrap();
            value | V::from(bits_to_add).unwrap()
        } else {
            value
        }
    }

    pub fn read<V>(&mut self, bit_size: u8) -> Result<V, SerializerError>
    where
        V: SerializableInt + std::ops::Shl<V, Output = V> + std::ops::BitOrAssign<V> + PrimInt,
    {
        if bit_size > 128 || bit_size as u32 > V::bit_size() {
            return Err(SerializerError::TooManyBits);
        } else if bit_size == 0 {
            return Err(SerializerError::TooFewBits);
        }

        // calculate how many bits can be accessed byte by byte
        let num_byte_bits = (bit_size / 8) * 8;

        match self.endianness {
            Endianness::LittleEndian => {
                let mut value = V::zero();
                // decode bytes starting with LSB
                for i in (0..num_byte_bits).step_by(8) {
                    value |= V::from_u8(self.read_u8(8)?) << V::from_u8(i);
                }

                // decode remaining bits (if they exist)
                if bit_size % 8 != 0 {
                    value |= V::from_u8(self.read_u8(bit_size % 8)?) << V::from_u8(num_byte_bits);
                }
                Ok(Self::sign_extend(value, bit_size))
            }
            Endianness::BigEndian => {
                let mut value = if bit_size % 8 != 0 {
                    V::from_u8(self.read_u8(bit_size % 8)?) << V::from_u8(num_byte_bits)
                } else {
                    V::zero()
                };

                // decode bytes starting with MSB
                for i in (8..=num_byte_bits).step_by(8) {
                    value |= V::from_u8(self.read_u8(8)?) << V::from_u8(num_byte_bits - i);
                }

                Ok(Self::sign_extend(value, bit_size))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use crate::common::Endianness;

    use super::{Deserializer, SerializerError};

    #[test]
    fn aligned_u8_read() -> Result<(), SerializerError> {
        let v = vec![
            0b00000000_u8,
            0b10101010_u8,
            0b01010101_u8,
            0b11111111_u8,
            0b00001111_u8,
            0b11110000_u8,
        ];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(deserializer.read::<u8>(8)?, 0b00000000);
        assert_eq!(deserializer.read::<u8>(8)?, 0b10101010);
        assert_eq!(deserializer.read::<u8>(8)?, 0b01010101);
        assert_eq!(deserializer.read::<u8>(8)?, 0b11111111);
        assert_eq!(deserializer.read::<u8>(8)?, 0b00001111);
        assert_eq!(deserializer.read::<u8>(8)?, 0b11110000);

        Ok(())
    }

    #[test]
    fn unaligned_u8_read() -> Result<(), SerializerError> {
        let v = vec![0b01010101_u8, 0b01010101_u8];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(deserializer.read::<u8>(5)?, 0b01010);
        assert_eq!(deserializer.read::<u8>(5)?, 0b10101);
        assert_eq!(deserializer.read::<u8>(3)?, 0b010);
        assert_eq!(deserializer.read::<u8>(1)?, 0b1);
        assert_eq!(deserializer.read::<u8>(2)?, 0b01);

        Ok(())
    }

    #[test]
    fn big_endian_different_unsigned_sizes() -> Result<(), SerializerError> {
        let v = vec![0b01010101_u8, 0b01010101_u8, 0b01010101_u8];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(deserializer.read::<u8>(4)?, 0b0101);
        assert_eq!(deserializer.read::<u16>(4)?, 0b0101);
        assert_eq!(deserializer.read::<u32>(4)?, 0b0101);
        assert_eq!(deserializer.read::<u64>(4)?, 0b0101);
        assert_eq!(deserializer.read::<u128>(4)?, 0b0101);
        assert_eq!(deserializer.read::<usize>(4)?, 0b0101);

        Ok(())
    }

    #[test]
    fn big_endian_different_signed_sizes() -> Result<(), SerializerError> {
        let v = vec![0b10101010_u8, 0b10101010_u8, 0b10101010_u8];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(deserializer.read::<i8>(4)?, -6);
        assert_eq!(deserializer.read::<i16>(4)?, -6);
        assert_eq!(deserializer.read::<i32>(4)?, -6);
        assert_eq!(deserializer.read::<i64>(4)?, -6);
        assert_eq!(deserializer.read::<i128>(4)?, -6);
        assert_eq!(deserializer.read::<isize>(4)?, -6);

        Ok(())
    }

    #[test]
    fn size_bigger_than_data() -> Result<(), SerializerError> {
        let v = vec![0b10101010_u8, 0b10101010_u8, 0b10101010_u8];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        let err = deserializer.read_u8(40).unwrap_err();
        let expected_err = SerializerError::TooManyBits;
        assert_eq!(err, expected_err);

        Ok(())
    }

    #[test]
    fn zero_bits_cannot_be_serialized() -> Result<(), SerializerError> {
        let v: Vec<u8> = Vec::new();

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(
            deserializer.read::<u8>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<u16>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<u32>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<u64>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<u128>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<usize>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<i8>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<i16>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<i32>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<i64>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<i128>(0).unwrap_err(),
            SerializerError::TooFewBits
        );
        assert_eq!(
            deserializer.read::<isize>(0).unwrap_err(),
            SerializerError::TooFewBits
        );

        Ok(())
    }

    #[test]
    fn value_out_of_range() -> Result<(), SerializerError> {
        let v: Vec<u8> = Vec::new();

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(
            deserializer.read::<u8>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<u16>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<u32>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<u64>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<u128>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<usize>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<i8>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<i16>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<i32>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<i64>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<i128>(129).unwrap_err(),
            SerializerError::TooManyBits
        );
        assert_eq!(
            deserializer.read::<isize>(129).unwrap_err(),
            SerializerError::TooManyBits
        );

        Ok(())
    }

    #[test]
    fn big_endian_multibyte() -> Result<(), SerializerError> {
        let v: Vec<u8> = vec![
            0xFF, 0x00, 0xFF, 0x00, 0xAA, 0x55, 0xFF, 0x00, 0xAA, 0x55, 0x33, 0xCC, 0xF0, 0x0F,
            0xFF, 0x00, 0xFF, 0x00, 0xAA, 0x55, 0xFF, 0x00, 0xAA, 0x55, 0x33, 0xCC, 0xF0, 0x0F,
        ];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(deserializer.read::<u16>(16)?, 0xFF00_u16);
        assert_eq!(deserializer.read::<u32>(32)?, 0xFF00AA55_u32);
        assert_eq!(deserializer.read::<u64>(64)?, 0xFF00AA5533CCF00F_u64);
        assert_eq!(deserializer.read::<i16>(16)?, -256i16); // Equivalent to 0xFF00
        assert_eq!(deserializer.read::<i32>(32)?, -16733611i32); // Equivalent to 0xFF00AA55
        assert_eq!(deserializer.read::<i64>(64)?, -71870311119917041i64); // Equivalent to 0xFF00AA5533CCF00F

        Ok(())
    }

    #[test]
    fn little_endian_multibyte() -> Result<(), SerializerError> {
        let v: Vec<u8> = vec![
            0x00, 0xFF, 0x55, 0xAA, 0x00, 0xFF, 0x0F, 0xF0, 0xCC, 0x33, 0x55, 0xAA, 0x00, 0xFF,
            0x00, 0xFF, 0x55, 0xAA, 0x00, 0xFF, 0x0F, 0xF0, 0xCC, 0x33, 0x55, 0xAA, 0x00, 0xFF,
        ];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::LittleEndian);

        assert_eq!(deserializer.read::<u16>(16)?, 0xFF00_u16);
        assert_eq!(deserializer.read::<u32>(32)?, 0xFF00AA55_u32);
        assert_eq!(deserializer.read::<u64>(64)?, 0xFF00AA5533CCF00F_u64);
        assert_eq!(deserializer.read::<i16>(16)?, -256i16); // Equivalent to 0xFF00
        assert_eq!(deserializer.read::<i32>(32)?, -16733611i32); // Equivalent to 0xFF00AA55
        assert_eq!(deserializer.read::<i64>(64)?, -71870311119917041i64); // Equivalent to 0xFF00AA5533CCF00F

        Ok(())
    }

    #[test]
    fn little_endian_multibyte_unaligned() -> Result<(), SerializerError> {
        let v: Vec<u8> = vec![
            0x00, 0xF_F, 0x55, 0xAA, 0x0_0, 0xFF, 0x0F, 0xF0, 0xCC, 0x33, 0x5_5, 0xAA, 0x0_0, 0xFF,
            0x00, 0xFF, 0x5_5, 0xAA, 0x00, 0xFF, 0x0F, 0xF0, 0xCC, 0x33, 0x55,
        ];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::LittleEndian);

        assert_eq!(deserializer.read::<u16>(12)?, 0xF00_u16);
        assert_eq!(deserializer.read::<u32>(24)?, 0xA05AF5_u32);
        assert_eq!(deserializer.read::<u64>(48)?, 0x35C30CFFF00F_u64);
        assert_eq!(deserializer.read::<i16>(16)?, -24486i16); // Equivalent to 0xA05A
        assert_eq!(deserializer.read::<i32>(32)?, -183504881i32); // Equivalent to 0xF50FF00F
        assert_eq!(deserializer.read::<i64>(64)?, 0x35C30CFFF00FA05Ai64);

        Ok(())
    }

    #[test]
    fn skip_bits() -> Result<(), SerializerError> {
        let v = vec![0b00110011_u8, 0b01000000_u8 | 0x01_u8, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xC0_u8 | 0b00000011_u8];

        let mut data = Cursor::new(v);

        let mut deserializer = Deserializer::new(&mut data, Endianness::BigEndian);

        assert_eq!(deserializer.read::<u8>(4)?, 0b0011);
        assert_eq!(deserializer.offset_in_bits(), 4);

        deserializer.skip(2)?;
        assert_eq!(deserializer.offset_in_bits(), 6);

        assert_eq!(deserializer.read::<u8>(2)?, 0b11);
        assert_eq!(deserializer.offset_in_bits(), 8);

        assert_eq!(deserializer.read::<u8>(3)?, 0b010);
        assert_eq!(deserializer.offset_in_bits(), 11);

        deserializer.skip(51)?;
        assert_eq!(deserializer.offset_in_bits(), 62);

        assert_eq!(deserializer.read::<u8>(2)?, 0b11);
        assert_eq!(deserializer.offset_in_bits(), 64);

        Ok(())
    }
}