//! Finding a DRED payload in an Opus packet: the packet's framing (RFC 6716 §3.2), the
//! extensions in its padding (draft-ietf-mlcodec-opus-extension-06 §2), and the DRED extension
//! among them (draft-ietf-mlcodec-opus-dred-07 §3, §5).
//!
//! The extension iterator follows the extension draft, including the repeat mechanism (ID 2,
//! §2.3): it gives every extension instance with its frame index, in bitstream order. An
//! extension whose length or frame index would read past the packet ends the iteration, and
//! the instances already given stand (§2.7).

use super::Reason;

/// The DRED extension ID (draft §3, §5) and the experimental ID the draft says to use until
/// the IANA action (§3, §5), whose data starts with the two-byte prefix `'D'`, version.
pub const DRED_ID: u8 = 32;
pub const DRED_EXPERIMENTAL_ID: u8 = 126;
pub const DRED_EXPERIMENTAL_PREFIX: [u8; 2] = [b'D', 12];

/// A packet's framing (RFC 6716 §3.2): its TOC fields, frame count and the padding bytes.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Framing<'a> {
    pub config: u8,
    pub stereo: bool,
    pub frames: usize,
    /// One frame's duration in units of 2.5 ms (RFC 6716 §3.1, Table 2).
    pub frame_ticks: u32,
    pub frame_lens: Vec<usize>,
    pub padding: &'a [u8],
}

/// One frame's duration, in units of 2.5 ms, for a TOC configuration (RFC 6716 §3.1).
pub fn frame_ticks(config: u8) -> u32 {
    match config {
        0..=11 => [4, 8, 16, 24][(config % 4) as usize],
        12..=15 => [4, 8][(config % 2) as usize],
        _ => [1, 2, 4, 8][(config % 4) as usize],
    }
}

fn frame_len(p: &[u8], i: &mut usize) -> Option<usize> {
    let b0 = *p.get(*i)? as usize;
    *i += 1;
    if b0 < 252 {
        return Some(b0);
    }
    let b1 = *p.get(*i)? as usize;
    *i += 1;
    Some(4 * b1 + b0)
}

/// Parse a packet's framing (RFC 6716 §3.2), with the length rules of §3.4 (R1–R7). A packet
/// that breaks one is not a valid Opus packet, and nothing in it is read.
pub fn framing(p: &[u8]) -> Result<Framing<'_>, Reason> {
    let bad = Err(Reason::PacketMalformed);
    let Some(&toc) = p.first() else { return bad };
    let (config, stereo, code) = (toc >> 3, toc & 4 != 0, toc & 3);
    let ticks = frame_ticks(config);
    let body = p.len() - 1;
    let mut padding: &[u8] = &[];
    let lens = match code {
        0 => vec![body],
        1 => {
            if !body.is_multiple_of(2) {
                return bad;
            }
            vec![body / 2; 2]
        }
        2 => {
            let mut i = 1;
            let Some(n1) = frame_len(p, &mut i) else { return bad };
            if n1 > p.len() - i {
                return bad;
            }
            vec![n1, p.len() - i - n1]
        }
        _ => {
            let Some(&fc) = p.get(1) else { return bad };
            let (vbr, pad, m) = (fc & 0x80 != 0, fc & 0x40 != 0, (fc & 0x3f) as usize);
            if m == 0 || m as u32 * ticks > 48 {
                return bad;
            }
            let mut i = 2;
            let mut plen = 0usize;
            if pad {
                loop {
                    let Some(&b) = p.get(i) else { return bad };
                    i += 1;
                    if b == 255 {
                        plen += 254;
                    } else {
                        plen += b as usize;
                        break;
                    }
                }
            }
            if i + plen > p.len() {
                return bad;
            }
            let end = p.len() - plen;
            padding = &p[end..];
            let mut lens = Vec::with_capacity(m);
            if vbr {
                for _ in 0..m - 1 {
                    let Some(n) = frame_len(p, &mut i) else { return bad };
                    if i > end {
                        return bad;
                    }
                    lens.push(n);
                }
                let used: usize = lens.iter().sum();
                if i > end || used > end - i {
                    return bad;
                }
                lens.push(end - i - used);
            } else {
                if i > end || !(end - i).is_multiple_of(m) {
                    return bad;
                }
                lens = vec![(end - i) / m; m];
            }
            lens
        }
    };
    if lens.iter().any(|&n| n > 1275) {
        return bad;
    }
    Ok(Framing { config, stereo, frames: lens.len(), frame_ticks: ticks, frame_lens: lens, padding })
}

/// One extension instance: its ID, the frame it belongs to, and its data.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Ext<'a> {
    pub id: u8,
    pub frame: usize,
    pub data: &'a [u8],
}

/// Skip one extension's payload (the length rules of §2): the data and the length-indicator
/// bytes, or `None` when they would read past `len`. `trailing` is the bytes the short
/// extensions repeated after an `L = 0` long one keep (§2.3).
fn skip_payload(p: &[u8], at: usize, id_byte: u8, trailing: usize) -> Option<(usize, usize)> {
    let (id, l) = (id_byte >> 1, (id_byte & 1) as usize);
    let len = p.len() - at;
    if (id == 0 && l == 1) || id == 2 {
        Some((at, 0))
    } else if id > 0 && id < 32 {
        if len < l {
            return None;
        }
        Some((at + l, 0))
    } else if l == 0 {
        if len < trailing {
            return None;
        }
        Some((p.len() - trailing, 0))
    } else {
        let (mut i, mut bytes, mut header) = (at, 0usize, 0usize);
        loop {
            let &lacing = p.get(i)?;
            i += 1;
            header += 1;
            bytes += lacing as usize;
            if lacing != 255 {
                break;
            }
        }
        if bytes > p.len() - i {
            return None;
        }
        Some((i + bytes, header))
    }
}

/// Every extension instance in `padding`, for a packet of `frames` frames, in bitstream order
/// (draft-ietf-mlcodec-opus-extension-06 §2). The iteration stops at the first extension that
/// would read past the padding or name a frame past the last; what came before it stands.
pub fn extensions(padding: &[u8], frames: usize) -> Vec<Ext<'_>> {
    let p = padding;
    let mut out = Vec::new();
    let (mut at, mut frame) = (0usize, 0usize);
    // The extensions the next repeat repeats start here; the last long extension among them;
    // the payload bytes of the short extensions after it.
    let mut repeat_from = 0usize;
    let mut last_long: Option<usize> = None;
    let mut trailing = 0usize;
    while at < p.len() {
        let start = at;
        let id_byte = p[at];
        let (id, l) = (id_byte >> 1, id_byte & 1);
        let Some((next, header)) = skip_payload(p, at + 1, id_byte, 0) else { return out };
        at = next;
        match id {
            1 => {
                let inc = if l == 0 { 1 } else { p[start + 1] as usize };
                if inc == 0 {
                    continue;
                }
                frame += inc;
                if frame >= frames {
                    return out;
                }
                repeat_from = at;
                last_long = None;
                trailing = 0;
            }
            2 => {
                // Repeat the non-padding extensions coded since `repeat_from` (up to this
                // indicator) for every later frame, each payload coded here in turn.
                let src: Vec<(usize, u8)> = {
                    let mut v = Vec::new();
                    let mut j = repeat_from;
                    while j < start {
                        let b = p[j];
                        let (n, _) = skip_payload(p, j + 1, b, 0).expect("skipped once already");
                        if b > 3 {
                            v.push((n, b));
                        }
                        j = n;
                    }
                    v
                };
                for f in frame + 1..frames {
                    for &(src_end, b) in &src {
                        let mut b = b;
                        if l == 0 && f + 1 >= frames && Some(src_end) == last_long {
                            b &= !1;
                        }
                        let Some((n, hdr)) = skip_payload(p, at, b, trailing) else { return out };
                        out.push(Ext { id: b >> 1, frame: f, data: &p[at + hdr..n] });
                        at = n;
                    }
                }
                repeat_from = at;
                last_long = None;
                if l == 0 {
                    frame += 1;
                    if frame >= frames {
                        return out;
                    }
                }
            }
            0 => {}
            _ => {
                if id >= 32 {
                    last_long = Some(at);
                    trailing = 0;
                } else {
                    trailing += l as usize;
                }
                out.push(Ext { id, frame, data: &p[start + 1 + header..at] });
            }
        }
    }
    out
}

/// The DRED payload a packet carries, and the offset of its frame within the packet in units
/// of 2.5 ms: the first extension instance that is DRED (ID 32 with data, or ID 126 whose data
/// starts with the experimental prefix and has more after it), as the draft's reference decoder
/// finds it. `Ok(None)` for a packet without one.
pub fn find_dred(p: &[u8]) -> Result<Option<(&[u8], i64)>, Reason> {
    let f = framing(p)?;
    for e in extensions(f.padding, f.frames) {
        let frame_offset = (e.frame as u32 * f.frame_ticks) as i64;
        if e.id == DRED_ID && !e.data.is_empty() {
            return Ok(Some((e.data, frame_offset)));
        }
        if e.id == DRED_EXPERIMENTAL_ID && e.data.len() > 2 && e.data[..2] == DRED_EXPERIMENTAL_PREFIX {
            return Ok(Some((&e.data[2..], frame_offset)));
        }
    }
    Ok(None)
}

/// A code-3 packet of one frame with `padding` (RFC 6716 §3.2.5), for the corpus generator.
pub fn build(toc: u8, frame: &[u8], padding: &[u8]) -> Vec<u8> {
    let mut p = vec![(toc & !3) | 3, 0x40 | 1];
    let mut n = padding.len();
    while n >= 255 {
        p.push(255);
        n -= 254;
    }
    p.push(n as u8);
    p.extend_from_slice(frame);
    p.extend_from_slice(padding);
    p
}

/// The padding of one DRED extension that takes the rest of the padding (`L = 0`), under the
/// experimental ID with its prefix, or under ID 32.
pub fn dred_padding(payload: &[u8], experimental: bool) -> Vec<u8> {
    let mut v = Vec::with_capacity(payload.len() + 3);
    if experimental {
        v.push(DRED_EXPERIMENTAL_ID << 1);
        v.extend_from_slice(&DRED_EXPERIMENTAL_PREFIX);
    } else {
        v.push(DRED_ID << 1);
    }
    v.extend_from_slice(payload);
    v
}

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

    /// The worked examples of draft-ietf-mlcodec-opus-extension-06 Appendix A, in a packet of
    /// three frames: each encoding gives the extensions its example lists.
    #[test]
    fn the_extension_drafts_examples_parse() {
        type Want = &'static [(u8, usize, &'static [u8])];
        let a: Want = &[(28, 0, b"a")];
        let ab: Want = &[(28, 0, b"a"), (28, 1, b"b")];
        let abc: Want = &[(28, 0, b"a"), (28, 1, b"b"), (28, 2, b"c")];
        let ex4: Want = &[(28, 0, b"a"), (28, 1, b"b"), (28, 2, b"c"), (29, 0, b"d")];
        let ex5: Want = &[(28, 0, b"a"), (28, 1, b"b"), (28, 2, b"c"), (29, 0, b"d"), (29, 1, b"")];
        let ex6: Want = &[(28, 0, b"a"), (28, 1, b"b"), (28, 2, b"c"), (29, 0, b"d"), (29, 1, b""), (29, 2, b"")];
        let ex7: Want = &[(28, 0, b"a"), (28, 1, b"b"), (28, 2, b"c"), (29, 0, b"d"), (29, 1, b""), (29, 2, b""), (120, 2, b"E0ex2")];
        let ex8: Want = &[(28, 0, b"a"), (28, 1, b"b"), (28, 2, b"c"), (29, 0, b"d"), (29, 1, b""), (120, 1, b"E0ex"), (29, 2, b""), (120, 2, b"E0ex2")];
        let cases: &[(&[u8], Want)] = &[
            (&[0x39, 0x61], a),
            (&[0x39, 0x61, 0x02, 0x39, 0x62], ab),
            (&[0x39, 0x61, 0x04, 0x62, 0x63], abc),
            (&[0x39, 0x61, 0x05, 0x62, 0x63, 0x3B, 0x64], ex4),
            (&[0x39, 0x61, 0x05, 0x62, 0x63, 0x3B, 0x64, 0x02, 0x3A], ex5),
            (&[0x39, 0x61, 0x05, 0x62, 0x63, 0x3B, 0x64, 0x02, 0x3A, 0x04], ex6),
            (&[0x39, 0x61, 0x05, 0x62, 0x63, 0x3B, 0x64, 0x02, 0x3A, 0x04, 0x04, 0x00], ex6),
            (&[0x39, 0x61, 0x05, 0x62, 0x63, 0x3B, 0x64, 0x02, 0x3A, 0x04, 0xF0, 0x45, 0x30, 0x65, 0x78, 0x32], ex7),
            (&[0x39, 0x61, 0x05, 0x62, 0x63, 0x3B, 0x64, 0x02, 0x3A, 0xF1, 0x04, 0x45, 0x30, 0x65, 0x78, 0x04, 0x45, 0x30, 0x65, 0x78, 0x32], ex8),
            (&[0x39, 0x61, 0x05, 0x62, 0x63, 0x3B, 0x64, 0x02, 0x3A, 0xF1, 0x04, 0x45, 0x30, 0x65, 0x78, 0x01, 0x04, 0x45, 0x30, 0x65, 0x78, 0x32], ex8),
        ];
        for (k, (bytes, want)) in cases.iter().enumerate() {
            let mut got: Vec<(u8, usize, Vec<u8>)> = extensions(bytes, 3).into_iter().map(|e| (e.id, e.frame, e.data.to_vec())).collect();
            let mut want: Vec<(u8, usize, Vec<u8>)> = want.iter().map(|&(i, f, d)| (i, f, d.to_vec())).collect();
            // Within a frame the order is the example's; across frames it is bitstream order.
            got.sort_by_key(|e| e.1);
            want.sort_by_key(|e| e.1);
            assert_eq!(got, want, "example {}", k + 1);
        }
    }

    #[test]
    fn a_dred_extension_is_found_in_a_packet() {
        let payload = [9u8, 8, 7, 6, 5];
        for exp in [true, false] {
            let p = build(9 << 3, &[1, 2, 3], &dred_padding(&payload, exp));
            let f = framing(&p).unwrap();
            assert_eq!((f.frames, f.frame_ticks, f.frame_lens.clone()), (1, 8, vec![3]));
            assert_eq!(find_dred(&p).unwrap(), Some((&payload[..], 0)));
        }
        // A long extension whose length passes the padding ends the iteration.
        let p = build(9 << 3, &[1], &[(DRED_ID << 1) | 1, 200, 1, 2]);
        assert_eq!(find_dred(&p).unwrap(), None);
        // An experimental extension with another prefix is not DRED.
        let p = build(9 << 3, &[1], &[DRED_EXPERIMENTAL_ID << 1, b'D', 11, 5, 5]);
        assert_eq!(find_dred(&p).unwrap(), None);
        assert_eq!(framing(&[]), Err(Reason::PacketMalformed));
    }
}
