//! The YUV4MPEG2 source an encode reads: 8-bit 4:2:0, progressive, each frame's three
//! planes one after another ([platform-encode] §4.1).
//!
//! [platform-encode]: ../../../../extensions/platform-encode.md

use super::EncodeError;
use std::io::Read;

/// A stream's header.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Header {
    pub width: u32,
    pub height: u32,
    pub fps_num: u32,
    pub fps_den: u32,
    /// `XCOLORRANGE=FULL`.
    pub full_range: bool,
}

impl Header {
    /// Bytes of one frame.
    pub fn frame_len(&self) -> usize {
        let (w, h) = (self.width as usize, self.height as usize);
        w * h + 2 * w.div_ceil(2) * h.div_ceil(2)
    }
}

fn bad(detail: impl Into<String>) -> EncodeError {
    EncodeError::new("encode-input", detail)
}

/// A reader of one stream's frames, one at a time.
pub struct Reader<R: Read> {
    r: R,
    pub header: Header,
}

fn read_line<R: Read>(r: &mut R, max: usize) -> Result<Option<Vec<u8>>, EncodeError> {
    let mut line = Vec::new();
    let mut b = [0u8; 1];
    loop {
        match r.read(&mut b) {
            Ok(0) => return if line.is_empty() { Ok(None) } else { Err(bad("the source ends inside a header line")) },
            Ok(_) if b[0] == b'\n' => return Ok(Some(line)),
            Ok(_) => {
                line.push(b[0]);
                if line.len() > max {
                    return Err(bad("a header line longer than its bound"));
                }
            }
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
            Err(e) => return Err(bad(e.to_string())),
        }
    }
}

impl<R: Read> Reader<R> {
    /// Read the stream header. Refused: another signature, no `W`, `H` or `F`, a
    /// chroma other than 4:2:0, a depth above 8 bits, or interlace other than `p`.
    pub fn new(mut r: R) -> Result<Reader<R>, EncodeError> {
        let line = read_line(&mut r, 4096)?.ok_or_else(|| bad("an empty source"))?;
        let text = std::str::from_utf8(&line).map_err(|_| bad("a header that is not UTF-8"))?;
        let mut it = text.split(' ');
        if it.next() != Some("YUV4MPEG2") {
            return Err(bad("not a YUV4MPEG2 stream"));
        }
        let (mut w, mut h, mut rate, mut full_range) = (None, None, None, false);
        for p in it.filter(|p| !p.is_empty()) {
            let (tag, val) = p.split_at(1);
            match tag {
                "W" => w = val.parse::<u32>().ok(),
                "H" => h = val.parse::<u32>().ok(),
                "F" => {
                    rate = val.split_once(':').and_then(|(n, d)| Some((n.parse::<u32>().ok()?, d.parse::<u32>().ok()?))).filter(|&(n, d)| n > 0 && d > 0);
                    if rate.is_none() {
                        return Err(bad(format!("frame rate {val}")));
                    }
                }
                "I" if val != "p" && val != "?" => return Err(bad(format!("interlace {val}: progressive sources only"))),
                "C" if !matches!(val, "420" | "420jpeg" | "420paldv" | "420mpeg2") => return Err(bad(format!("chroma {val}: 8-bit 4:2:0 sources only"))),
                "X" if val == "COLORRANGE=FULL" => full_range = true,
                _ => {}
            }
        }
        let (Some(width), Some(height), Some((fps_num, fps_den))) = (w, h, rate) else {
            return Err(bad("the header lacks W, H or F"));
        };
        if width == 0 || height == 0 || width > 16384 || height > 16384 {
            return Err(bad(format!("a picture of {width}×{height}")));
        }
        Ok(Reader { r, header: Header { width, height, fps_num, fps_den, full_range } })
    }

    /// The next frame's planes, or `None` at the end of the stream.
    pub fn next_frame(&mut self) -> Result<Option<Vec<u8>>, EncodeError> {
        let Some(line) = read_line(&mut self.r, 4096)? else { return Ok(None) };
        if !line.starts_with(b"FRAME") {
            return Err(bad("a frame that does not begin with FRAME"));
        }
        let mut f = vec![0u8; self.header.frame_len()];
        self.r.read_exact(&mut f).map_err(|_| bad("the source ends inside a frame"))?;
        Ok(Some(f))
    }
}

/// Every frame of `bytes`.
pub fn read_all(bytes: &[u8]) -> Result<(Header, Vec<Vec<u8>>), EncodeError> {
    let mut r = Reader::new(bytes)?;
    let mut frames = Vec::new();
    while let Some(f) = r.next_frame()? {
        frames.push(f);
    }
    Ok((r.header, frames))
}

/// One stream of `frames`, each `w × h` 4:2:0 planes, at `fps_num / fps_den`.
pub fn write(w: u32, h: u32, fps_num: u32, fps_den: u32, frames: &[Vec<u8>]) -> Vec<u8> {
    let mut out = format!("YUV4MPEG2 W{w} H{h} F{fps_num}:{fps_den} Ip A1:1 C420jpeg\n").into_bytes();
    for f in frames {
        out.extend_from_slice(b"FRAME\n");
        out.extend_from_slice(f);
    }
    out
}

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

    #[test]
    fn a_written_stream_reads_back() {
        let frames = vec![(0..24u8).collect::<Vec<_>>(), vec![7; 24]];
        let s = write(4, 4, 30000, 1001, &frames);
        let (h, back) = read_all(&s).unwrap();
        assert_eq!(h, Header { width: 4, height: 4, fps_num: 30000, fps_den: 1001, full_range: false });
        assert_eq!(back, frames);
        assert!(read_all(b"YUV4MPEG2 W4 H4 F30:1 C444\n").is_err());
        assert!(read_all(b"YUV4MPEG2 W4 H4 F30:1 It\n").is_err());
        assert!(read_all(&s[..s.len() - 1]).is_err());
    }
}
