//! A stream under loss (extensions/opus-dred.md §5): which packets a seeded pattern drops, which
//! received packet's DRED rebuilds each lost 10 ms, the rebuilt features, and the digest a
//! verifier recomputes; and the presentation statement over it (§7), an entry of the player's
//! log of extensions/synthesis.md §10.3.
//!
//! The stream is a run of packets of one frame duration each, numbered from 0, so packet `n`
//! starts at `n·d` in units of 2.5 ms. After a run of lost packets, the first received packet
//! recovers the run: its DRED's newest feature vector ends `dred_offset` after its own start
//! (draft-ietf-mlcodec-opus-dred-07 §3), each older one 10 ms earlier, and the 10 ms half `h`
//! of lost frame `m`, which ends at `E = m·d + 4·(h + 1)`, takes row `r = ⌊(A − E)/4⌋` of it,
//! `A` the newest vector's end, when `0 ≤ r` and `r` is below 4 times its latent count.
//! Otherwise that half is not covered: no DRED rebuilds it.

use sha2::{Digest, Sha256};

use super::packet;
use super::payload;
use super::{Model, Reason, Row};

/// The loss pattern (§5.1): a two-state chain, lost or received, driven by a counter-mode
/// hash of its seed. Packet `n` draws `u_n`, the first two bytes (big-endian) of
/// SHA-256(`"wai:dred-loss\x01"` ‖ seed ‖ `n` as a big-endian `u32`); it is lost when
/// `u_n < enter` after a received packet (or at the start), and when `u_n < stay` after a
/// lost one. Thresholds are in units of `2^-16`, at most `2^16`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct LossModel {
    pub seed: [u8; 32],
    pub enter: u32,
    pub stay: u32,
}

pub const LOSS_DOMAIN: &[u8] = b"wai:dred-loss\x01";

impl LossModel {
    pub fn lost(&self, packets: usize) -> Vec<bool> {
        let mut prev = false;
        (0..packets)
            .map(|n| {
                let d = Sha256::new().chain_update(LOSS_DOMAIN).chain_update(self.seed).chain_update((n as u32).to_be_bytes()).finalize();
                let u = u16::from_be_bytes([d[0], d[1]]) as u32;
                prev = u < if prev { self.stay } else { self.enter };
                prev
            })
            .collect()
    }
}

/// One lost 10 ms half (§5.2): its frame, which half (0 the earlier), and the packet and row
/// that rebuild it, if any.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Slot {
    pub frame: usize,
    pub half: u8,
    pub source: Option<(usize, usize)>,
}

/// The DRED a received packet carries, as the stream reads it: its payload's integers and the
/// newest vector's end relative to the packet's start. `None` when the packet carries none, is
/// malformed, or its payload is refused: then it rebuilds nothing.
pub fn dred_of(p: &[u8]) -> Option<(payload::Payload, i64)> {
    let (bytes, frame_offset) = packet::find_dred(p).ok()??;
    let pl = payload::parse(bytes, None).ok()?;
    let off = pl.header.dred_offset(frame_offset);
    Some((pl, off))
}

/// Which row of which packet rebuilds each lost half (§5.2), for a stream of packets of `d`
/// ticks each (`None` where lost).
pub fn plan(packets: &[Option<Vec<u8>>], d: i64) -> Vec<Slot> {
    let mut slots = Vec::new();
    let mut m = 0usize;
    while m < packets.len() {
        if packets[m].is_some() {
            m += 1;
            continue;
        }
        let a = m;
        while m < packets.len() && packets[m].is_none() {
            m += 1;
        }
        // Lost run a..m; packet m (if any) recovers it.
        let rec = packets.get(m).and_then(|p| p.as_deref()).and_then(dred_of).map(|(pl, off)| (m, pl.latents.len(), m as i64 * d + off));
        for f in a..m {
            for h in 0..2u8 {
                let e = f as i64 * d + 4 * (h as i64 + 1);
                let source = rec.and_then(|(n, l, anchor)| {
                    let diff = anchor - e;
                    let r = diff.div_euclid(4);
                    (diff >= 0 && (r as usize) < 4 * l).then_some((n, r as usize))
                });
                slots.push(Slot { frame: f, half: h, source });
            }
        }
    }
    slots
}

/// The rows the slots name, each recovering packet's DRED decoded through as many latents as
/// its oldest slot needs.
pub fn rebuild(model: &Model, packets: &[Option<Vec<u8>>], slots: &[Slot]) -> Result<Vec<Option<Row>>, Reason> {
    let mut cache: Vec<(usize, Vec<Row>)> = Vec::new();
    let mut out = Vec::with_capacity(slots.len());
    for s in slots {
        let Some((n, r)) = s.source else {
            out.push(None);
            continue;
        };
        if !cache.iter().any(|(k, _)| *k == n) {
            let need = slots.iter().filter_map(|t| t.source).filter(|(k, _)| *k == n).map(|(_, r)| r / 4 + 1).max().unwrap_or(0);
            let (pl, _) = dred_of(packets[n].as_deref().expect("a source is received")).expect("a source carries DRED");
            cache.push((n, model.decode(&pl, need)?));
        }
        let rows = &cache.iter().find(|(k, _)| *k == n).unwrap().1;
        out.push(Some(rows[r]));
    }
    Ok(out)
}

/// The reconstruction buffer (§5.3): for each slot in order, `frame` (`u32`), `half` (`u8`),
/// whether it is covered (`u8`), and for a covered one the packet and row (`u32` each) and the
/// row's 20 features (`i32`); little-endian.
pub fn buffer(slots: &[Slot], rows: &[Option<Row>]) -> Vec<u8> {
    let mut b = Vec::new();
    for (s, r) in slots.iter().zip(rows) {
        b.extend_from_slice(&(s.frame as u32).to_le_bytes());
        b.push(s.half);
        match (s.source, r) {
            (Some((n, row)), Some(v)) => {
                b.push(1);
                b.extend_from_slice(&(n as u32).to_le_bytes());
                b.extend_from_slice(&(row as u32).to_le_bytes());
                for x in v {
                    b.extend_from_slice(&x.to_le_bytes());
                }
            }
            _ => b.push(0),
        }
    }
    b
}

pub fn digest(b: &[u8]) -> [u8; 32] {
    Sha256::digest(b).into()
}

/// The lost runs, each with the packet after it (which recovers it when received) and the range
/// of its slots: a run ends where the next slot's frame skips a received packet.
pub fn recoveries(slots: &[Slot]) -> Vec<(usize, std::ops::Range<usize>)> {
    let mut out = Vec::new();
    let mut i = 0;
    while i < slots.len() {
        let mut j = i + 1;
        while j < slots.len() && slots[j].frame <= slots[j - 1].frame + 1 {
            j += 1;
        }
        out.push((slots[j - 1].frame + 1, i..j));
        i = j;
    }
    out
}

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

    #[test]
    fn the_loss_pattern_is_its_seeds() {
        let m = LossModel { seed: [7; 32], enter: 1 << 14, stay: 1 << 15 };
        let a = m.lost(500);
        assert_eq!(a, m.lost(500));
        let n = a.iter().filter(|&&x| x).count();
        assert!(n > 50 && n < 250, "{n}");
        assert!(LossModel { enter: 0, ..m }.lost(100).iter().all(|&x| !x));
        assert!(LossModel { enter: 1 << 16, stay: 1 << 16, ..m }.lost(100).iter().all(|&x| x));
        assert_ne!(a, LossModel { seed: [8; 32], ..m }.lost(500));
    }

    #[test]
    fn slots_follow_the_anchor() {
        // A packet of 20 ms whose DRED ends at its own end (offset 8) and holds 3 latents (12 rows).
        let h = payload::Header { q0: 2, dq: 0, extended: None, indicator: 8, qmax: 15 };
        let mut l = [0i32; payload::LATENT_DIM];
        for (k, v) in l.iter_mut().enumerate() {
            if payload::coded(super::super::tables::LATENT_DECAY[k][2], super::super::tables::LATENT_P0[k][2]) {
                *v = 3 - (k % 7) as i32;
            }
        }
        let pl = payload::Payload { header: h, state: [0; payload::STATE_DIM], latents: vec![(l, 2); 3] };
        assert_eq!(payload::parse(&payload::write(&pl), None).unwrap(), pl);
        let pk = packet::build(9 << 3, &[], &packet::dred_padding(&payload::write(&pl), true));
        let pk = Some(pk);
        // Frames 0 and 5 received; 1..=4 lost; 5 recovers.
        let packets = vec![pk.clone(), None, None, None, None, pk.clone()];
        let s = plan(&packets, 8);
        assert_eq!(s.len(), 8);
        // A = 5·8 + 8 = 48. Frame 4 half 1 ends at 40: r = 2; half 0 at 36: r = 3.
        assert_eq!(s[7], Slot { frame: 4, half: 1, source: Some((5, 2)) });
        assert_eq!(s[6], Slot { frame: 4, half: 0, source: Some((5, 3)) });
        // Frame 1 half 1 ends at 16: r = 8; half 0 at 12: r = 9 (rows 0..12).
        assert_eq!(s[1].source, Some((5, 8)));
        assert_eq!(s[0].source, Some((5, 9)));
        // Fewer latents: frame 1 is not covered by 2 latents (rows 0..8).
        let pl2 = payload::Payload { latents: pl.latents[..2].to_vec(), ..pl.clone() };
        let pk2 = Some(packet::build(9 << 3, &[], &packet::dred_padding(&payload::write(&pl2), false)));
        let s = plan(&[pk.clone(), None, None, None, None, pk2], 8);
        assert_eq!((s[0].source, s[1].source, s[2].source), (None, None, Some((5, 7))));
        // A run at the end has no recovering packet.
        let s = plan(&[pk.clone(), None], 8);
        assert_eq!(s, vec![Slot { frame: 1, half: 0, source: None }, Slot { frame: 1, half: 1, source: None }]);
        assert_eq!(recoveries(&plan(&packets, 8)), vec![(5, 0..8)]);
    }
}
