#[derive(Clone, Debug)]
pub enum ByteStreamError {
    OutOfRange,
}

pub const fn bit_mask(mask: u8) -> u8 {
    match mask {
        0 => 0x00,
        1 => 0x01,
        2 => 0x03,
        3 => 0x07,
        4 => 0x0f,
        5 => 0x1f,
        6 => 0x3f,
        7 => 0x7f,
        _ => 0xff,
    }
}

#[derive(Default, Clone, Debug)]
pub struct ByteStream {
    data: Vec<u8>,
}

impl ByteStream {
    pub fn get_bytes(&self, bit_ndx: usize, bit_count: usize) -> Result<Vec<u8>, ByteStreamError> {
        let byte_ndx = bit_ndx / 8;
        let bits_before = (bit_ndx % 8) as u8;
        let bits_in_last_byte = ((bit_ndx + bit_count - bits_before as usize) % 8) as u8;
        let byte_count = if bit_count >= 8 { bit_count / 8 } else { 1 };
        let bytes_needed = if bits_before != 0 {
            byte_count + 1
        } else {
            byte_count
        };

        if bytes_needed > self.data.len() || bytes_needed + byte_ndx > self.data.len() {
            return Err(ByteStreamError::OutOfRange);
        }

        let mut byte_stream = self.data[byte_ndx..byte_ndx + bytes_needed].to_vec();

        if bits_before != 0 {
            let mut carry: u8 = 0;
            for _ in 0..bits_before {
                for k in (0..bytes_needed).rev() {
                    let next: u8 = (byte_stream[k] & 0x01) << 7;
                    byte_stream[k] = (byte_stream[k] >> 1) | carry;

                    carry = next;
                }
            }
        }

        if bytes_needed > byte_count {
            byte_stream.pop();
        }

        if bits_in_last_byte != 0 {
            *byte_stream.last_mut().unwrap() &= bit_mask(bits_in_last_byte);
        }

        Ok(byte_stream)
    }
}

impl From<&[u8]> for ByteStream {
    fn from(slice: &[u8]) -> Self {
        ByteStream::from(slice.to_vec())
    }
}

impl From<Vec<u8>> for ByteStream {
    fn from(vec: Vec<u8>) -> Self {
        ByteStream { data: vec }
    }
}

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

    #[test]
    fn test_get_bytes_no_shift() {
        let bytes: Vec<u8> = vec![0xff, 0x00, 0x55];
        let bit_stream = ByteStream::from(bytes.clone());

        let new_bytes = bit_stream.get_bytes(0, bytes.len() * 8).unwrap();
        assert_eq!(bytes, new_bytes);
    }

    #[test]
    fn test_get_bytes_with_shift_in_byte() {
        let bytes: Vec<u8> = vec![0xff, 0x00, 0x55];
        let bit_stream = ByteStream::from(bytes.clone());

        let new_bytes = bit_stream.get_bytes(4, 4).unwrap();
        assert_eq!(vec![0x0f], new_bytes);
    }

    #[test]
    fn test_get_bytes_with_shift_across_bytes() {
        let bytes: Vec<u8> = vec![0xff, 0x55];
        let bit_stream = ByteStream::from(bytes.clone());

        let new_bytes = bit_stream.get_bytes(4, 8).unwrap();
        assert_eq!(vec![0x5f], new_bytes);
    }

    #[test]
    fn test_get_bytes_with_shift_across_bytes_odd() {
        let bytes: Vec<u8> = vec![0xff, 0x55];
        let bit_stream = ByteStream::from(bytes.clone());

        let new_bytes = bit_stream.get_bytes(7, 2).unwrap();
        assert_eq!(vec![0x03], new_bytes);
    }

    #[test]
    fn test_get_bytes_one_byte() {
        let bytes: Vec<u8> = vec![0xff, 0x55];
        let bit_stream = ByteStream::from(bytes.clone());

        let new_bytes = bit_stream.get_bytes(8, 8).unwrap();
        assert_eq!(vec![0x55], new_bytes);
    }

    #[test]
    fn test_get_bytes_3_bits() {
        let bytes: Vec<u8> = vec![0x55];
        let bit_stream = ByteStream::from(bytes.clone());

        let new_bytes = bit_stream.get_bytes(0, 3).unwrap();
        assert_eq!(vec![0x05], new_bytes);
    }
}