sharkyflac 0.2.0

use std::io::{self, Read, Seek, Write};

use byteorder::ReadBytesExt;

use crate::bits::Bitset;

pub trait BitRead: Read {
    /// Returns whether the reader is byte-aligned.
    fn is_byte_aligned(&self) -> bool;

    /// Retrieve the leftover bits. This byte-aligns the reader. The data is
    /// returned as a `u8` followed by the number of bits read.
    fn take_leftover(&mut self) -> Option<(u8, u8)>;

    /// Discard the leftover bits. This makes future reads byte aligned.
    fn discard_leftover(&mut self);

    /// Read a single bit
    fn read_bit(&mut self) -> io::Result<bool>;

    /// Read `n` bits (n ≤ 64) MSB-first.
    fn read_bits(&mut self, n: u32) -> io::Result<u64>;

    /// Read `n` bits sign-extended to an `i64`. `n` must be in the range
    /// `1..=64`. Reading zero bits returns `0`.
    fn read_signed(&mut self, n: u8) -> io::Result<i64>;
}

/// A bit reader over a byte stream. Starts reading a byte boundary
#[derive(Debug)]
pub struct BitReader<R: Read> {
    inner:     R,
    buf:       u8,
    bits_left: u8,
}

impl<R: Read> BitReader<R> {
    /// Create a new [`BitReader`].
    pub const fn new(inner: R) -> Self {
        Self {
            inner,
            buf: 0,
            bits_left: 0,
        }
    }

    /// Get the inner reader
    #[inline]
    pub fn into_inner(self) -> R {
        self.inner
    }
}

impl<R: Read> Read for BitReader<R> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }

        // Fast Path: If we are on a byte boundary, delegate straight to the inner
        // reader.
        if self.bits_left == 0 {
            return self.inner.read(buf);
        }

        // Slow Path: We are misaligned. Stitch together bits to construct bytes.
        let mut bytes_read = 0;
        for byte in buf.iter_mut() {
            match self.read_bits(8) {
                Ok(val) => {
                    *byte = val as u8;
                    bytes_read += 1;
                }
                Err(e) => {
                    // Standard `Read` behavior: if we hit EOF but already managed
                    // to read some bytes, return the count instead of the error.
                    if bytes_read > 0 && e.kind() == io::ErrorKind::UnexpectedEof {
                        break;
                    }
                    return Err(e);
                }
            }
        }

        Ok(bytes_read)
    }
}

impl<R: Read + Seek> Seek for BitReader<R> {
    #[inline]
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        self.bits_left = 0;
        self.inner.seek(pos)
    }

    #[inline]
    fn seek_relative(&mut self, offset: i64) -> io::Result<()> {
        self.bits_left = 0;
        self.inner.seek_relative(offset)
    }
}

impl<R: Read> BitRead for BitReader<R> {
    /// Returns whether the reader is byte-aligned.
    #[inline]
    fn is_byte_aligned(&self) -> bool {
        self.bits_left == 0
    }

    /// Read one bit
    #[inline]
    fn read_bit(&mut self) -> io::Result<bool> {
        if self.bits_left == 0 {
            self.buf = self.inner.read_u8()?;
            self.bits_left = 8;
        }

        let bit = self.buf.get_bit_msb(8 - u32::from(self.bits_left)); // MSB-first
        self.bits_left -= 1;
        Ok(bit)
    }

    /// Read `n` bits (n ≤ 64) MSB-first.
    #[inline]
    fn read_bits(&mut self, mut n: u32) -> io::Result<u64> {
        if n == 0 {
            return Ok(0);
        }
        assert!(n <= 64);

        let mut value: u64 = 0;
        let mut bit_count = 0;

        while n > 0 {
            if self.bits_left == 0 {
                match self.inner.read_u8() {
                    Ok(byte) => {
                        self.buf = byte;
                        self.bits_left = 8;
                    }
                    // recover trailing bits in the case of unaligned reads
                    Err(e) => {
                        self.buf = value as u8;
                        self.bits_left = bit_count as u8;
                        return Err(e);
                    }
                }
            }

            let take = n.min(u32::from(self.bits_left));
            let shift = u32::from(self.bits_left) - take;

            // Extract the top `take` bits from the current buffer byte
            let chunk = u64::from(self.buf) >> shift;
            let mask = (1u64 << take) - 1;
            value = (value << take) | (chunk & mask);

            self.bits_left -= take as u8;
            bit_count += take;
            n -= take;
        }

        Ok(value)
    }

    /// Read `n` bits sign-extended to an `i64`. `n` must be in the range
    /// `1..=64`. Reading zero bits returns `0`.
    #[inline]
    fn read_signed(&mut self, n: u8) -> io::Result<i64> {
        debug_assert!((0..=64).contains(&n));
        if n == 0 {
            return Ok(0);
        }

        let raw = self.read_bits(u32::from(n))? as i64;
        let sh = 64 - n;
        Ok((raw << sh) >> sh)
    }

    /// Retrieve the leftover bits. This byte-aligns the reader. The data is
    /// returned as a `u8` followed by the number of bits read.
    #[inline]
    fn take_leftover(&mut self) -> Option<(u8, u8)> {
        if self.bits_left != 0 {
            Some((
                std::mem::take(&mut self.buf).get_bit_range_lsb(0, self.bits_left),
                std::mem::take(&mut self.bits_left),
            ))
        } else {
            None
        }
    }

    /// Discard the leftover bits. This makes future reads byte aligned.
    #[inline]
    fn discard_leftover(&mut self) {
        self.bits_left = 0;
    }
}

pub trait BitWrite: Write {
    fn is_byte_aligned(&self) -> bool;

    /// Pad remaining bits to zero and flush the current byte.
    /// No-op if already aligned.
    fn flush_bits(&mut self) -> io::Result<()>;

    /// Write the lowest `n` bits of `value`, MSB first. `n` must be ≤ 64.
    fn write_bits(&mut self, value: u64, n: u32) -> io::Result<()>;

    /// Write a single bit.
    #[inline]
    fn write_bit(&mut self, bit: bool) -> io::Result<()> {
        self.write_bits(u64::from(bit), 1)
    }

    /// Write `n` bits of a signed value, in two's complement MSB-first.
    #[inline]
    fn write_signed(&mut self, value: i64, n: u8) -> io::Result<()> {
        let mask = if n == 64 { u64::MAX } else { (1u64 << n) - 1 };
        self.write_bits(value as u64 & mask, u32::from(n))
    }
}

pub struct BitWriter<W: Write> {
    inner:     W,
    // accumulator populated msb first
    buf:       u8,
    bits_used: u8,
}

impl<W: Write> BitWriter<W> {
    pub const fn new(inner: W) -> Self {
        Self {
            inner,
            buf: 0,
            bits_used: 0,
        }
    }

    pub fn into_inner(self) -> W {
        self.inner
    }
}

impl<W: Write> Write for BitWriter<W> {
    #[inline]
    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
        if self.bits_used == 0 {
            return self.inner.write(data);
        }

        for &byte in data {
            self.write_bits(u64::from(byte), 8)?;
        }

        Ok(data.len())
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }
}

impl<W: Write> BitWrite for BitWriter<W> {
    #[inline]
    fn is_byte_aligned(&self) -> bool {
        self.bits_used == 0
    }

    #[inline]
    fn flush_bits(&mut self) -> io::Result<()> {
        if self.bits_used > 0 {
            // Remaining bits are already packed into the high end of buf;
            // the unused low bits are already zero from prior masking.
            self.inner.write_all(&[self.buf])?;
            self.buf = 0;
            self.bits_used = 0;
        }
        Ok(())
    }

    #[inline]
    fn write_bits(&mut self, value: u64, mut n: u32) -> io::Result<()> {
        debug_assert!(n <= 64);

        while n > 0 {
            let space = 8 - u32::from(self.bits_used);
            let take = n.min(space);
            // Extract the top `take` bits from value and shift them into position.
            let shift = n - take;
            let chunk = ((value >> shift) as u8) & ((1u8.unbounded_shr(take)).wrapping_sub(1));
            self.buf |= chunk << (space - take);
            self.bits_used += take as u8;
            n -= take;

            if self.bits_used == 8 {
                self.inner.write_all(&[self.buf])?;
                self.buf = 0;
                self.bits_used = 0;
            }
        }

        Ok(())
    }
}

#[allow(clippy::bool_assert_comparison)]
#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Cursor, ErrorKind, Read};

    fn create_reader(data: &[u8]) -> BitReader<Cursor<&[u8]>> {
        BitReader::new(Cursor::new(data))
    }

    #[test]
    fn read_single_bits() {
        // 0b10110010 = 178
        let mut reader = create_reader(&[0b1011_0010]);

        assert_eq!(reader.read_bit().unwrap(), true);
        assert_eq!(reader.read_bit().unwrap(), false);
        assert_eq!(reader.read_bit().unwrap(), true);
        assert_eq!(reader.read_bit().unwrap(), true);

        assert_eq!(reader.read_bit().unwrap(), false);
        assert_eq!(reader.read_bit().unwrap(), false);
        assert_eq!(reader.read_bit().unwrap(), true);
        assert_eq!(reader.read_bit().unwrap(), false);
    }

    #[test]
    fn read_bits_within_byte() {
        let mut reader = create_reader(&[0b1101_0101]);

        assert_eq!(reader.read_bits(4).unwrap(), 0b1101);
        assert_eq!(reader.read_bits(3).unwrap(), 0b010);
        assert_eq!(reader.read_bits(1).unwrap(), 0b1);
    }

    #[test]
    fn read_bits_across_byte_boundary() {
        // [0b1111_0000, 0b1010_1100]
        let mut reader = create_reader(&[0xF0, 0xAC]);

        assert_eq!(reader.read_bits(4).unwrap(), 0b1111);
        // Crosses boundary: takes 4 bits (0000) from byte 1, and 4 bits (1010) from
        // byte 2
        assert_eq!(reader.read_bits(8).unwrap(), 0b0000_1010);
        assert_eq!(reader.read_bits(4).unwrap(), 0b1100);
    }

    #[test]
    fn read_bits_64_bits() {
        // 8 bytes of 0xFF, followed by 1 byte of 0x00
        let data = [0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x00];
        let mut reader = create_reader(&data);

        // Read a full 64-bit integer
        assert_eq!(reader.read_bits(64).unwrap(), u64::MAX);
        // Read 8 bits from the remaining byte
        assert_eq!(reader.read_bits(8).unwrap(), 0x00);
    }

    #[test]
    fn discard_leftover() {
        let mut reader = create_reader(&[0b1111_1111, 0b0101_0101]);

        assert_eq!(reader.read_bits(3).unwrap(), 0b111);

        reader.discard_leftover();

        assert_eq!(reader.read_bits(8).unwrap(), 0b0101_0101);
    }

    #[test]
    fn read_trait_passthrough() {
        let data = [0x01, 0x02, 0x03, 0x04];
        let mut reader = create_reader(&data);
        let mut buf = [0u8; 3];

        let bytes_read = reader.read(&mut buf).unwrap();
        assert_eq!(bytes_read, 3);
        assert_eq!(buf, [0x01, 0x02, 0x03]);

        let cursor = reader.into_inner();
        assert_eq!(cursor.position(), 3);
    }

    #[test]
    fn read_trait_misaligned() {
        let data = [0b1111_0000, 0b1010_1111, 0b0000_1111];
        let mut reader = create_reader(&data);

        assert_eq!(reader.read_bits(4).unwrap(), 0b1111);

        let mut buf = [0u8; 2];
        let bytes_read = reader.read(&mut buf).unwrap();

        assert_eq!(bytes_read, 2);
        assert_eq!(buf[0], 0b0000_1010);
        assert_eq!(buf[1], 0b1111_0000);
    }

    #[test]
    fn read_bit_eof() {
        let mut reader = create_reader(&[0xAA]);

        assert_eq!(reader.read_bits(8).unwrap(), 0xAA);

        assert_eq!(
            reader.read_bit().unwrap_err().kind(),
            ErrorKind::UnexpectedEof
        );
        assert_eq!(reader.take_leftover(), None);
    }

    #[test]
    fn read_trait_partial_eof() {
        let data = [0b1111_0000, 0b1010_1111];
        let mut reader = create_reader(&data);

        // Misalign by 4 bits
        reader.read_bits(4).unwrap();

        let mut buf = [0u8; 2];
        let bytes_read = reader.read(&mut buf).unwrap();

        assert_eq!(bytes_read, 1);
        assert_eq!(buf[0], 0b0000_1010);

        assert_eq!(reader.take_leftover().unwrap(), (0b1111, 4));
    }

    #[test]
    fn seek() {
        let mut reader = create_reader(&[0b1011_0010, 0b0101_0101]);

        assert_eq!(reader.read_bit().unwrap(), true);
        assert_eq!(reader.read_bit().unwrap(), false);
        assert_eq!(reader.read_bit().unwrap(), true);
        assert_eq!(reader.read_bit().unwrap(), true);

        reader.seek(io::SeekFrom::Start(1)).unwrap();

        assert_eq!(reader.read_bit().unwrap(), false);
        assert_eq!(reader.read_bit().unwrap(), true);
        assert_eq!(reader.read_bit().unwrap(), false);
        assert_eq!(reader.read_bit().unwrap(), true);
        assert_eq!(reader.read_bit().unwrap(), false);
        assert_eq!(reader.read_bit().unwrap(), true);
    }
}

#[cfg(test)]
mod test_bit_writer {
    use super::*;
    use crate::bit_io::{BitRead, BitReader};
    use std::io::Cursor;

    fn make_writer() -> BitWriter<Vec<u8>> {
        BitWriter::new(Vec::new())
    }

    #[test]
    fn write_single_bits() {
        // 0b1011_0010 = 0xB2 — same constant as read_single_bits in BitReader
        let mut w = make_writer();
        for bit in [true, false, true, true, false, false, true, false] {
            w.write_bit(bit).unwrap();
        }
        assert_eq!(w.into_inner(), [0xB2]);
    }

    #[test]
    fn write_bits_within_byte() {
        // 4 bits 0b1101 | 3 bits 0b010 | 1 bit 0b1 = 0b1101_0101 = 0xD5
        // Mirrors read_bits_within_byte.
        let mut w = make_writer();
        w.write_bits(0b1101, 4).unwrap();
        w.write_bits(0b010, 3).unwrap();
        w.write_bits(0b1, 1).unwrap();
        assert_eq!(w.into_inner(), [0xD5]);
    }

    #[test]
    fn write_bits_across_byte_boundary() {
        // [0xF0, 0xAC] — mirrors read_bits_across_byte_boundary.
        let mut w = make_writer();
        w.write_bits(0b1111, 4).unwrap();
        w.write_bits(0b0000_1010, 8).unwrap();
        w.write_bits(0b1100, 4).unwrap();
        assert_eq!(w.into_inner(), [0xF0, 0xAC]);
    }

    #[test]
    fn write_u64_max() {
        // 8 bytes of 0xFF — mirrors read_bits_64_bits.
        let mut w = make_writer();
        w.write_bits(u64::MAX, 64).unwrap();
        assert_eq!(w.into_inner(), [0xFF; 8]);
    }

    #[test]
    fn flush_bits_zero_pads_remaining_bits() {
        // Write 3 bits (0b101), then flush. The remaining 5 bits must be 0.
        // Result: 0b101_00000 = 0xA0
        let mut w = make_writer();
        w.write_bits(0b101, 3).unwrap();
        w.flush_bits().unwrap();
        assert_eq!(w.into_inner(), [0xA0]);
    }

    #[test]
    fn flush_bits_is_noop_when_aligned() {
        let mut w = make_writer();
        w.write_bits(0xFF, 8).unwrap();
        w.flush_bits().unwrap(); // should not emit an extra byte
        assert_eq!(w.into_inner(), [0xFF]);
    }

    #[test]
    fn write_signed_negative_value() {
        // -1 in 8-bit two's complement is 0xFF.
        let mut w = make_writer();
        w.write_signed(-1, 8).unwrap();
        assert_eq!(w.into_inner(), [0xFF]);
    }

    #[test]
    fn write_signed_positive_value() {
        // 127 in 8-bit signed = 0x7F.
        let mut w = make_writer();
        w.write_signed(127, 8).unwrap();
        assert_eq!(w.into_inner(), [0x7F]);
    }

    #[test]
    fn write_signed_5bit_negative() {
        // -1 in 5 bits = 0b11111 = 0x1F, shifted to top of byte: 0b1111_1000 = 0xF8
        // (because we only write 5 bits then the byte is incomplete)
        let mut w = make_writer();
        w.write_signed(-1, 5).unwrap();
        w.flush_bits().unwrap();
        assert_eq!(w.into_inner(), [0xF8]);
    }

    #[test]
    fn round_trip_arbitrary_bit_pattern() {
        let mut w = make_writer();
        w.write_bits(0b1011, 4).unwrap();
        w.write_bits(0b00101, 5).unwrap();
        w.write_bits(0b111, 3).unwrap();
        w.write_bits(0b0100, 4).unwrap();
        let bytes = w.into_inner();

        let mut r = BitReader::new(Cursor::new(bytes));
        assert_eq!(r.read_bits(4).unwrap(), 0b1011);
        assert_eq!(r.read_bits(5).unwrap(), 0b00101);
        assert_eq!(r.read_bits(3).unwrap(), 0b111);
        assert_eq!(r.read_bits(4).unwrap(), 0b0100);
    }

    #[test]
    fn round_trip_signed_values() {
        let values: &[(i64, u8)] = &[
            (0, 5),
            (1, 5),
            (-1, 5),
            (15, 5),
            (-16, 5),
            (127, 8),
            (-128, 8),
        ];

        for &(value, bits) in values {
            let mut w = make_writer();
            w.write_signed(value, bits).unwrap();
            w.flush_bits().unwrap();
            let bytes = w.into_inner();

            let mut r = BitReader::new(Cursor::new(bytes));
            let decoded = r.read_signed(bits).unwrap();
            assert_eq!(
                decoded, value,
                "round-trip failed for value={value} bits={bits}"
            );
        }
    }
}