//! JPEG AI (ISO/IEC 6048-1) codestream parser — port of the reference
//! `bitstream_structure/`: `SOC … [marker | expGolomb-k0 size | payload] … EOC`,
//! big-endian markers, MSB-first bit reader for the size field. Pure integer.

pub const MARKER_SOC: u16 = 0xFF80;
pub const MARKER_EOC: u16 = 0xFF81;
pub const MARKER_PIH: u16 = 0xFF82;
pub const MARKER_TON: u16 = 0xFF83;
pub const MARKER_RDI: u16 = 0xFF84;
pub const MARKER_SOZ: u16 = 0xFF88;
pub const MARKER_SORP: u16 = 0xFF89;
pub const MARKER_SORS: u16 = 0xFF8A;
pub const MARKER_SOQ: u16 = 0xFF8B;
pub const MARKER_UDI: u16 = 0xFF8C;

const KNOWN: [u16; 8] = [MARKER_PIH, MARKER_TON, MARKER_RDI, MARKER_SOZ,
                         MARKER_SORP, MARKER_SORS, MARKER_SOQ, MARKER_UDI];

/// MSB-first bit reader (port of `StreamBitReader`, big-endian byte order).
struct BitReader<'a> {
    buf: &'a [u8],
    byte: usize,
    bit: u32,
}

impl<'a> BitReader<'a> {
    fn new(buf: &'a [u8]) -> Self {
        Self { buf, byte: 0, bit: 0 }
    }

    fn read_bit(&mut self) -> Option<u32> {
        let b = *self.buf.get(self.byte)?;
        let v = (b >> (7 - self.bit)) & 1;
        self.bit += 1;
        if self.bit == 8 {
            self.bit = 0;
            self.byte += 1;
        }
        Some(v as u32)
    }

    /// Bytes consumed, rounding a partial byte up (the size field is flushed
    /// to a byte boundary on the encode side).
    fn consumed(&self) -> usize {
        self.byte + usize::from(self.bit != 0)
    }

    /// Unsigned exp-Golomb, k = 0.
    fn ue(&mut self) -> Option<u64> {
        let mut zeros = 0u32;
        while self.read_bit()? == 0 {
            zeros += 1;
            if zeros > 63 {
                return None;
            }
        }
        let mut v: u64 = 1;
        for _ in 0..zeros {
            v = (v << 1) | self.read_bit()? as u64;
        }
        Some(v - 1)
    }
}

#[derive(Debug)]
pub struct Substream<'a> {
    pub marker: u16,
    pub payload: &'a [u8],
}

/// Walk `SOC … substreams … EOC`; unknown markers or truncation abort (None).
pub fn parse_codestream(buf: &[u8]) -> Option<Vec<Substream<'_>>> {
    let rd16 = |o: usize| -> Option<u16> {
        Some(u16::from_be_bytes([*buf.get(o)?, *buf.get(o + 1)?]))
    };
    if rd16(0)? != MARKER_SOC {
        return None;
    }
    let mut off = 2usize;
    let mut out = Vec::new();
    loop {
        let marker = rd16(off)?;
        off += 2;
        if marker == MARKER_EOC {
            return Some(out);
        }
        if !KNOWN.contains(&marker) {
            return None;
        }
        let mut br = BitReader::new(&buf[off..]);
        let size = br.ue()? as usize;
        off += br.consumed();
        let payload = buf.get(off..off + size)?;
        off += size;
        out.push(Substream { marker, payload });
    }
}

impl<'a> BitReader<'a> {
    /// Signed exp-Golomb, k = 0 (reference convention: `2·|d| − (d>0)`,
    /// even code → negative).
    fn se(&mut self) -> Option<i64> {
        let v = self.ue()?;
        let mag = ((v + 1) >> 1) as i64;
        Some(if v & 1 == 0 { -mag } else { mag })
    }
}

/// Unpack an AE substream payload (SOZ/SORP/SORS/SOQ) into the mANS memory
/// region and per-thread END offsets, ready for `MansDecoder` — port of
/// `aemem.parse_substream`: `(num_threads − 1)` signed exp-Golomb deltas off
/// the mean thread size (byte-aligned), then the raw coder memory.
pub fn parse_ae_substream(payload: &[u8], num_threads: usize) -> Option<(&[u8], Vec<u32>)> {
    if num_threads == 0 {
        return None;
    }
    let mut br = BitReader::new(payload);
    let mut deltas = Vec::with_capacity(num_threads - 1);
    for _ in 0..num_threads - 1 {
        deltas.push(br.se()?);
    }
    let mem = &payload[br.consumed()..];
    let mean = (mem.len() / num_threads) as i64;
    let mut sizes: Vec<i64> = deltas.iter().map(|d| mean - d).collect();
    let used: i64 = sizes.iter().sum();
    sizes.push(mem.len() as i64 - used);
    let mut offsets = Vec::with_capacity(num_threads);
    let mut acc = 0i64;
    for s in sizes {
        if s <= 0 {
            return None;
        }
        acc += s;
        offsets.push(u32::try_from(acc).ok()?);
    }
    Some((mem, offsets))
}

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

    /// Encode-side mirror (marker BE + ue(k0) flushed to byte + payload).
    fn frame(marker: u16, payload: &[u8]) -> Vec<u8> {
        let n = payload.len() as u64 + 1;
        let len = 64 - n.leading_zeros();
        let mut bits = vec![0u8; (len - 1) as usize];
        for i in (0..len).rev() {
            bits.push(((n >> i) & 1) as u8);
        }
        let mut out = marker.to_be_bytes().to_vec();
        for chunk in bits.chunks(8) {
            let mut b = 0u8;
            for (i, &bit) in chunk.iter().enumerate() {
                b |= bit << (7 - i);
            }
            out.push(b);
        }
        out.extend_from_slice(payload);
        out
    }

    #[test]
    fn codestream_roundtrip_walks_all_substreams() {
        let mut cs = MARKER_SOC.to_be_bytes().to_vec();
        let z = vec![7u8; 300]; // multi-byte exp-Golomb size
        let r = vec![9u8; 5];
        cs.extend(frame(MARKER_PIH, &[1, 2, 3]));
        cs.extend(frame(MARKER_SOZ, &z));
        cs.extend(frame(MARKER_SORP, &r));
        cs.extend(MARKER_EOC.to_be_bytes());
        let subs = parse_codestream(&cs).expect("parse");
        assert_eq!(subs.len(), 3);
        assert_eq!((subs[0].marker, subs[0].payload.len()), (MARKER_PIH, 3));
        assert_eq!((subs[1].marker, subs[1].payload), (MARKER_SOZ, z.as_slice()));
        assert_eq!((subs[2].marker, subs[2].payload), (MARKER_SORP, r.as_slice()));
        // corrupt marker → refused
        let mut bad = cs.clone();
        bad[2] = 0xAA;
        assert!(parse_codestream(&bad).is_none());
    }

    /// Encode-side mirror of the thread-delta prefix: se(k0) per delta,
    /// flushed to a byte boundary, then the raw memory.
    fn frame_ae(deltas: &[i64], mem: &[u8]) -> Vec<u8> {
        let mut bits = Vec::new();
        for &d in deltas {
            let v = (2 * d.unsigned_abs()) - u64::from(d > 0);
            let n = v + 1;
            let len = 64 - n.leading_zeros();
            bits.extend(std::iter::repeat_n(0u8, (len - 1) as usize));
            for i in (0..len).rev() {
                bits.push(((n >> i) & 1) as u8);
            }
        }
        let mut out = Vec::new();
        for chunk in bits.chunks(8) {
            let mut b = 0u8;
            for (i, &bit) in chunk.iter().enumerate() {
                b |= bit << (7 - i);
            }
            out.push(b);
        }
        out.extend_from_slice(mem);
        out
    }

    #[test]
    fn ae_substream_thread_offsets() {
        // Single thread: payload IS the memory, one end offset.
        let (mem, off) = parse_ae_substream(&[5, 6, 7, 8], 1).unwrap();
        assert_eq!((mem, off.as_slice()), (&[5u8, 6, 7, 8][..], &[4u32][..]));

        // 3 threads over 90 bytes: mean 30, deltas +2/−5 → sizes 28, 35, 27.
        let memory = vec![0xABu8; 90];
        let payload = frame_ae(&[2, -5], &memory);
        let (mem, off) = parse_ae_substream(&payload, 3).unwrap();
        assert_eq!(mem, memory.as_slice());
        assert_eq!(off, vec![28, 63, 90]);

        // Oversized delta starves a thread → refused.
        let bad = frame_ae(&[31], &[0u8; 60]);
        assert!(parse_ae_substream(&bad, 2).is_none());
    }
}
