//! Minimal H.264 Annex-B NAL scanning.
//!
//! A `foxglove.CompressedVideo` frame's `data` is an Annex-B bytestream (NAL
//! units separated by `00 00 01` / `00 00 00 01` start codes). We only need to
//! know whether an access unit contains an IDR slice (nal_unit_type 5), i.e.
//! whether the frame is a keyframe and therefore a clean cut/splice point.

/// True if the access unit contains an IDR slice NAL (type 5) — a keyframe.
pub fn is_keyframe(annexb: &[u8]) -> bool {
    nal_types(annexb).any(|t| t == 5)
}

/// Iterate the `nal_unit_type` (low 5 bits of the byte after each start code)
/// for every NAL unit in an Annex-B bytestream.
pub fn nal_types(data: &[u8]) -> NalTypes<'_> {
    NalTypes { data, pos: 0 }
}

pub struct NalTypes<'a> {
    data: &'a [u8],
    pos: usize,
}

impl Iterator for NalTypes<'_> {
    type Item = u8;

    fn next(&mut self) -> Option<u8> {
        let d = self.data;
        let mut i = self.pos;
        // Scan for a `00 00 01` start-code prefix (the 4-byte `00 00 00 01`
        // variant contains it at its second byte, so this finds both).
        while i + 3 <= d.len() {
            if d[i] == 0 && d[i + 1] == 0 && d[i + 2] == 1 {
                let nal_start = i + 3;
                if nal_start >= d.len() {
                    self.pos = d.len();
                    return None;
                }
                self.pos = nal_start + 1;
                return Some(d[nal_start] & 0x1f);
            }
            i += 1;
        }
        self.pos = d.len();
        None
    }
}

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

    #[test]
    fn detects_idr_3byte_startcode() {
        // 00 00 01 | 0x65 (forbidden_zero=0, nri=3, type=5 => IDR)
        let data = [0x00, 0x00, 0x01, 0x65, 0xaa, 0xbb];
        assert!(is_keyframe(&data));
    }

    #[test]
    fn detects_non_idr_with_4byte_startcode() {
        // 00 00 00 01 | 0x41 (type=1 => non-IDR slice)
        let data = [0x00, 0x00, 0x00, 0x01, 0x41, 0xaa];
        assert!(!is_keyframe(&data));
        assert_eq!(nal_types(&data).collect::<Vec<_>>(), vec![1]);
    }

    #[test]
    fn sps_pps_then_idr() {
        // SPS(7), PPS(8), IDR(5) — typical keyframe AU with inline param sets.
        let data = [
            0x00, 0x00, 0x01, 0x67, 0x10, // SPS
            0x00, 0x00, 0x01, 0x68, 0x20, // PPS
            0x00, 0x00, 0x01, 0x65, 0x30, // IDR
        ];
        assert_eq!(nal_types(&data).collect::<Vec<_>>(), vec![7, 8, 5]);
        assert!(is_keyframe(&data));
    }
}
