//! MPEG transport streams (media-identify §5.5): the packet size the sync byte keeps, the program
//! tables, each with its checksum, and what scrambling the packets state. The payloads are not read.

use super::{Candidate, Strength, Track};
use serde_json::json;
use std::collections::BTreeMap;

/// A packet's size and where its sync byte sits in it.
const LAYOUTS: [(usize, usize); 3] = [(188, 0), (192, 4), (204, 0)];
/// The packets a run of sync bytes needs.
const MIN_PACKETS: usize = 5;
/// The most programs read.
const MAX_PROGRAMS: usize = 64;

/// MPEG-2's checksum: the polynomial 0x04C11DB7, from all ones, not reflected.
pub(crate) fn crc32_mpeg(b: &[u8]) -> u32 {
    let mut c = 0xffff_ffffu32;
    for &x in b {
        c ^= (x as u32) << 24;
        for _ in 0..8 {
            c = if c & 0x8000_0000 != 0 { (c << 1) ^ 0x04c1_1db7 } else { c << 1 };
        }
    }
    c
}

pub(super) fn codec(stream_type: u8) -> Option<&'static str> {
    Some(match stream_type {
        0x1b => "avc",
        0x24 => "hevc",
        0x33 => "vvc",
        0x0f | 0x11 => "aac",
        0x03 | 0x04 => "mp3",
        0x81 => "ac3",
        0x87 => "eac3",
        _ => return None,
    })
}

pub(super) fn media(stream_type: u8) -> &'static str {
    match stream_type {
        0x01 | 0x02 | 0x10 | 0x1b | 0x24 | 0x33 => "video",
        0x03 | 0x04 | 0x0f | 0x11 | 0x81 | 0x87 => "audio",
        _ => "other",
    }
}

/// The first packet size whose sync byte holds, and the offset of the first packet.
fn layout(head: &[u8]) -> Option<(usize, usize, usize)> {
    for (size, sync) in LAYOUTS {
        for start in 0..size {
            let ok = (0..MIN_PACKETS).all(|k| head.get(start + k * size + sync) == Some(&0x47));
            if ok {
                return Some((size, sync, start));
            }
        }
    }
    None
}

/// A section being gathered across packets.
struct Gathered {
    bytes: Vec<u8>,
    at: u64,
}

pub fn walk(head: &[u8], size: u64) -> Option<Candidate> {
    let (psize, sync, start) = layout(head)?;
    let mut c = Candidate::new("mpeg-ts", "container");
    c.see((start + sync) as u64, 1, "sync");
    c.fact("packet_size", psize as u64);
    let mut packets = 0u64;
    let mut scrambled = 0u64;
    let mut pmts: BTreeMap<u64, u64> = BTreeMap::new();
    let mut gather: BTreeMap<u64, Gathered> = BTreeMap::new();
    let (mut pat_done, mut pmts_done) = (false, BTreeMap::<u64, bool>::new());
    let mut pos = start;
    let mut stopped = None;
    let mut streams: Vec<(u64, u8, bool)> = Vec::new();
    while let Some(p) = head.get(pos + sync..pos + sync + 188) {
        if p[0] != 0x47 {
            stopped = Some(pos + sync);
            break;
        }
        packets += 1;
        let pid = (((p[1] & 0x1f) as u64) << 8) | p[2] as u64;
        let pusi = p[1] & 0x40 != 0;
        if p[3] & 0xc0 != 0 {
            scrambled += 1;
        }
        let afc = (p[3] >> 4) & 3;
        let mut i = 4usize;
        if afc & 2 != 0 {
            i += 1 + p[4] as usize;
        }
        let wanted = pid == 0 || pmts.values().any(|&v| v == pid);
        if afc & 1 != 0 && i < 188 && wanted {
            if pusi {
                let pointer = p[i] as usize;
                i += 1 + pointer;
                if i < 188 {
                    gather.insert(pid, Gathered { bytes: Vec::new(), at: (pos + sync + i) as u64 });
                }
            }
            if let (Some(g), true) = (gather.get_mut(&pid), i < 188) {
                g.bytes.extend_from_slice(&p[i..]);
                let total = (g.bytes.len() >= 3).then(|| 3 + ((((g.bytes[1] & 0x0f) as usize) << 8) | g.bytes[2] as usize));
                if let Some(n) = total.filter(|&n| g.bytes.len() >= n) {
                    let at = g.at;
                    let sec = g.bytes[..n].to_vec();
                    gather.remove(&pid);
                    if n >= 12 && crc32_mpeg(&sec) == 0 {
                        if pid == 0 && sec[0] == 0 && !pat_done {
                            pat_done = true;
                            c.see(at, n as u64, "pat");
                            for e in sec[8..n - 4].as_chunks::<4>().0.iter().take(MAX_PROGRAMS) {
                                let program = ((e[0] as u64) << 8) | e[1] as u64;
                                let map = (((e[2] & 0x1f) as u64) << 8) | e[3] as u64;
                                if program != 0 {
                                    pmts.insert(program, map);
                                    pmts_done.insert(map, false);
                                }
                            }
                        } else if sec[0] == 2 && pmts_done.get(&pid) == Some(&false) {
                            pmts_done.insert(pid, true);
                            c.see(at, n as u64, "pmt");
                            let info = (((sec[10] & 0x0f) as usize) << 8) | sec[11] as usize;
                            let has_ca = |d: &[u8]| {
                                let mut j = 0;
                                while j + 2 <= d.len() {
                                    if d[j] == 0x09 {
                                        return true;
                                    }
                                    j += 2 + d[j + 1] as usize;
                                }
                                false
                            };
                            let mut j = 12 + info;
                            let program_ca = sec.get(12..j).is_some_and(has_ca);
                            while j + 5 <= n - 4 {
                                let ty = sec[j];
                                let es = (((sec[j + 1] & 0x1f) as u64) << 8) | sec[j + 2] as u64;
                                let dl = (((sec[j + 3] & 0x0f) as usize) << 8) | sec[j + 4] as usize;
                                let ca = program_ca || sec.get(j + 5..j + 5 + dl).is_some_and(has_ca);
                                streams.push((es, ty, ca));
                                j += 5 + dl;
                            }
                        }
                    }
                }
            }
        }
        pos += psize;
    }
    for (es, ty, ca) in streams {
        c.track(Track {
            track: Some(es),
            media: media(ty),
            stated: format!("stream_type:0x{ty:02x}"),
            codec: codec(ty),
            protected: ca.then(|| json!({"scheme": null, "original": null})),
            ..Track::default()
        });
    }
    c.fact("packets_read", packets);
    c.fact("scrambled_packets", scrambled);
    c.fact("programs", pmts.len() as u64);
    c.walked_to = Some(pos as u64);
    match stopped {
        Some(at) => c.stop(at as u64, "invalid"),
        None if (pos as u64) < size && (head.len() as u64) == size => c.stop(pos as u64, "overrun"),
        None if (head.len() as u64) < size => c.stop(pos as u64, "limit"),
        None => {}
    }
    c.strength = Strength::Structure;
    Some(c)
}
