//! Bit-level reading and writing shared by the probes: an MSB-first bit reader with
//! `u(n)` and Exp-Golomb `ue(v)`, the Annex B byte-stream split of ITU-T H.265 and H.266
//! (and the start-code form ISO/IEC 23094-2 uses), and the removal of emulation-prevention
//! bytes. The writer is the corpus generator's.

use super::Refusal;

/// An MSB-first reader over a byte slice. Every read past the end is
/// [`Refusal::Truncated`].
pub struct BitReader<'a> {
    data: &'a [u8],
    pos: usize,
}

impl<'a> BitReader<'a> {
    pub fn new(data: &'a [u8]) -> Self {
        Self { data, pos: 0 }
    }
    /// The bit position from the start.
    pub fn position(&self) -> usize {
        self.pos
    }
    /// Bits left to read.
    pub fn remaining(&self) -> usize {
        self.data.len() * 8 - self.pos
    }
    pub fn byte_aligned(&self) -> bool {
        self.pos.is_multiple_of(8)
    }
    /// `u(n)` for `n` ≤ 64.
    pub fn u(&mut self, n: u32) -> Result<u64, Refusal> {
        debug_assert!(n <= 64);
        if self.remaining() < n as usize {
            return Err(Refusal::Truncated);
        }
        let mut v: u64 = 0;
        for _ in 0..n {
            let byte = self.data[self.pos / 8];
            let bit = (byte >> (7 - (self.pos % 8))) & 1;
            v = (v << 1) | u64::from(bit);
            self.pos += 1;
        }
        Ok(v)
    }
    pub fn flag(&mut self) -> Result<bool, Refusal> {
        Ok(self.u(1)? == 1)
    }
    /// Skips `n` bits.
    pub fn skip(&mut self, n: usize) -> Result<(), Refusal> {
        if self.remaining() < n {
            return Err(Refusal::Truncated);
        }
        self.pos += n;
        Ok(())
    }
    /// `ue(v)` (ITU-T H.265 §9.2, H.266 §9.2). A code with more than 31 leading zero bits
    /// has no value a syntax element of these standards may take: [`Refusal::ReservedValue`].
    pub fn ue(&mut self) -> Result<u64, Refusal> {
        let mut zeros = 0u32;
        while self.u(1)? == 0 {
            zeros += 1;
            if zeros > 31 {
                return Err(Refusal::ReservedValue);
            }
        }
        if zeros == 0 {
            return Ok(0);
        }
        Ok((1u64 << zeros) - 1 + self.u(zeros)?)
    }
    /// Skips to the next byte boundary, requiring every skipped bit to equal `bit`, or any
    /// value when `bit` is `None`.
    pub fn align(&mut self, bit: Option<u64>) -> Result<(), Refusal> {
        while !self.byte_aligned() {
            let b = self.u(1)?;
            if let Some(want) = bit
                && b != want
            {
                return Err(Refusal::ReservedValue);
            }
        }
        Ok(())
    }
}

/// An MSB-first bit writer (the corpus generator's).
#[derive(Default)]
pub struct BitWriter {
    out: Vec<u8>,
    acc: u8,
    nbits: u32,
}

impl BitWriter {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn u(&mut self, n: u32, v: u64) {
        for i in (0..n).rev() {
            self.acc = (self.acc << 1) | ((v >> i) & 1) as u8;
            self.nbits += 1;
            if self.nbits == 8 {
                self.out.push(self.acc);
                self.acc = 0;
                self.nbits = 0;
            }
        }
    }
    pub fn flag(&mut self, b: bool) {
        self.u(1, u64::from(b));
    }
    pub fn ue(&mut self, v: u64) {
        let x = v + 1;
        let len = 64 - x.leading_zeros();
        self.u(len - 1, 0);
        self.u(len, x);
    }
    pub fn byte_aligned(&self) -> bool {
        self.nbits == 0
    }
    /// Pads with `bit` to the next byte boundary.
    pub fn align(&mut self, bit: u64) {
        while self.nbits != 0 {
            self.u(1, bit);
        }
    }
    /// `rbsp_trailing_bits()`: a one, then zeros to the byte boundary.
    pub fn trailing(&mut self) {
        self.u(1, 1);
        self.align(0);
    }
    pub fn bytes(&mut self, b: &[u8]) {
        for &x in b {
            self.u(8, u64::from(x));
        }
    }
    pub fn finish(mut self) -> Vec<u8> {
        self.align(0);
        self.out
    }
}

/// One NAL unit of an Annex B byte stream: its bytes after the start code, emulation
/// prevention still in place, trailing zero bytes removed.
#[derive(Debug, Clone, Copy)]
pub struct NalSpan<'a> {
    pub bytes: &'a [u8],
}

/// Splits an Annex B byte stream (ITU-T H.265 and H.266 Annex B; ISO/IEC 23094-2's byte
/// stream) into NAL units. The stream must begin, after any leading zero bytes, with a
/// start code prefix `0x000001`; otherwise [`Refusal::NoSignature`]. A NAL unit ends where
/// the next start code prefix begins; the zero bytes before a prefix belong to no NAL unit.
pub fn annex_b_units(stream: &[u8]) -> Result<Vec<NalSpan<'_>>, Refusal> {
    let mut i = 0;
    while i < stream.len() && stream[i] == 0 {
        i += 1;
    }
    if i < 2 || i >= stream.len() || stream[i] != 1 {
        return Err(Refusal::NoSignature);
    }
    let mut starts = Vec::new();
    let mut j = 0;
    while j + 2 < stream.len() {
        if stream[j] == 0 && stream[j + 1] == 0 && stream[j + 2] == 1 {
            starts.push(j + 3);
            j += 3;
        } else {
            j += 1;
        }
    }
    let mut units = Vec::with_capacity(starts.len());
    for (k, &s) in starts.iter().enumerate() {
        let mut end = if k + 1 < starts.len() { starts[k + 1] - 3 } else { stream.len() };
        while end > s && stream[end - 1] == 0 {
            end -= 1;
        }
        units.push(NalSpan { bytes: &stream[s..end] });
    }
    Ok(units)
}

/// Removes emulation-prevention bytes: in every `0x000003`, the `0x03` (ITU-T H.265 and
/// H.266 §7.4.2, `emulation_prevention_three_byte`).
pub fn unescape(nal: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(nal.len());
    let mut zeros = 0;
    for &b in nal {
        if zeros >= 2 && b == 3 {
            zeros = 0;
            continue;
        }
        if b == 0 {
            zeros += 1;
        } else {
            zeros = 0;
        }
        out.push(b);
    }
    out
}

/// Inserts emulation-prevention bytes (the generator's inverse of [`unescape`]).
pub fn escape(rbsp: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(rbsp.len() + rbsp.len() / 64);
    let mut zeros = 0;
    for &b in rbsp {
        if zeros >= 2 && b <= 3 {
            out.push(3);
            zeros = 0;
        }
        if b == 0 {
            zeros += 1;
        } else {
            zeros = 0;
        }
        out.push(b);
    }
    out
}

/// Prefixes each NAL unit with a four-byte start code `0x00000001`.
pub fn annex_b(units: &[Vec<u8>]) -> Vec<u8> {
    let mut out = Vec::new();
    for u in units {
        out.extend_from_slice(&[0, 0, 0, 1]);
        out.extend_from_slice(u);
    }
    out
}

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

    #[test]
    fn exp_golomb_round_trips() {
        let mut w = BitWriter::new();
        for v in [0u64, 1, 2, 3, 7, 254, 1 << 20, (1 << 32) - 2] {
            w.ue(v);
        }
        let b = w.finish();
        let mut r = BitReader::new(&b);
        for v in [0u64, 1, 2, 3, 7, 254, 1 << 20, (1 << 32) - 2] {
            assert_eq!(r.ue().unwrap(), v);
        }
    }

    #[test]
    fn emulation_prevention_round_trips() {
        let rbsp = [0u8, 0, 0, 0, 0, 1, 0, 0, 2, 0, 0, 3, 0, 0, 4];
        let e = escape(&rbsp);
        assert_eq!(e, [0, 0, 3, 0, 0, 3, 0, 1, 0, 0, 3, 2, 0, 0, 3, 3, 0, 0, 4]);
        assert_eq!(unescape(&e), rbsp);
    }

    #[test]
    fn annex_b_needs_a_start_code() {
        assert_eq!(annex_b_units(&[1, 2, 3]).unwrap_err(), Refusal::NoSignature);
        assert_eq!(annex_b_units(&[0, 1, 2]).unwrap_err(), Refusal::NoSignature);
        let u = annex_b_units(&[0, 0, 0, 1, 0x40, 1, 0, 0, 1, 0x42, 1, 0, 0]).unwrap();
        assert_eq!(u.len(), 2);
        assert_eq!(u[0].bytes, &[0x40, 1]);
        assert_eq!(u[1].bytes, &[0x42, 1]);
    }
}
