// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use anyhow::ensure;
use byteorder::{BigEndian, WriteBytesExt};
use fidl_fuchsia_media::FormatDetails;
use std::path::Path;
use std::{fs, mem};
use stream_processor_test::*;

const SEPARATE_SPS_PPS: bool = true;
// written to vector as big-endian
const NAL_START_CODE: u32 = 1;

/// Represents an H264 elementary stream.
pub struct H264Stream {
    data: Vec<u8>,
}

impl H264Stream {
    /// Adds data to raw H264 elementary stream
    pub fn append(&mut self, data: &mut Vec<u8>) {
        self.data.append(data);
    }

    /// Constructs an H264 elementary stream from a file with raw elementary stream data.
    pub fn from_file(filename: impl AsRef<Path>) -> Result<Self> {
        Ok(H264Stream::from(fs::read(filename)?))
    }

    /// Returns an iterator over H264 NALs that does not copy.
    pub fn nal_iter(&self) -> impl Iterator<Item = H264Nal<'_>> {
        H264NalIter { data: &self.data, pos: 0 }
    }

    /// Returns an iterator over H264 chunks (pictures and anything preceding each picture) that
    /// does not copy.
    //
    // TODO(https://fxbug.dev/42084549): Make H264ChunkIter capable of iterating NALs or iterating pictures with any
    // non-picture NALs in front of each picture prepended (the current behavior).
    fn chunk_iter(&self) -> impl Iterator<Item = H264Chunk<'_>> {
        H264ChunkIter::create(&self.data)
    }
}

impl From<Vec<u8>> for H264Stream {
    fn from(data: Vec<u8>) -> Self {
        Self { data }
    }
}

impl ElementaryStream for H264Stream {
    fn format_details(&self, version_ordinal: u64) -> FormatDetails {
        FormatDetails {
            format_details_version_ordinal: Some(version_ordinal),
            mime_type: Some(String::from("video/h264")),
            ..Default::default()
        }
    }

    fn is_access_units(&self) -> bool {
        true
    }

    fn stream<'a>(&'a self) -> Box<dyn Iterator<Item = ElementaryStreamChunk> + 'a> {
        Box::new(self.chunk_iter().map(|chunk| ElementaryStreamChunk {
            start_access_unit: true,
            known_end_access_unit: true,
            data: chunk.data.to_vec(),
            significance: match chunk.nal_kind {
                Some(H264NalKind::IDR) | Some(H264NalKind::NonIDR) => {
                    Significance::Video(VideoSignificance::Picture)
                }
                _ => Significance::Video(VideoSignificance::NotPicture),
            },
            timestamp: None,
        }))
    }
}

pub struct H264SeiItuT35 {
    pub country_code: u8,
    pub country_code_extension: u8,
    pub payload: Vec<u8>,
}

impl H264SeiItuT35 {
    pub const COUNTRY_CODE_UNITED_STATES: u8 = 0xb5;

    pub fn as_bytes(&self) -> Result<Vec<u8>> {
        const ITU_T35_PAYLOAD_TYPE: u8 = 4;

        let mut bytes = vec![];
        bytes.write_u32::<BigEndian>(NAL_START_CODE)?;
        bytes.write_u8(SEI_CODE)?;
        bytes.write_u8(ITU_T35_PAYLOAD_TYPE)?;
        bytes.write_u8(u8::try_from(self.payload_size())?)?;
        bytes.write_u8(self.country_code)?;
        bytes.write_u8(self.country_code_extension)?;
        bytes.append(&mut self.payload.clone());
        Ok(bytes)
    }

    fn payload_size(&self) -> usize {
        // At the time this is called, the country_code and country_code_extension bytes haven't
        // yet been written.
        mem::size_of::<u8>() + mem::size_of::<u8>() + self.payload.len()
    }
}

/// Represents an AVCC stream.
pub struct H264AVCCStream {
    // NAL units with 4-byte length prefixes but no start codes.
    nal_units: Vec<(H264NalKind, Vec<u8>)>,
    oob_data: Vec<u8>,
}

impl H264AVCCStream {
    pub fn from_annexb_stream(annexb: H264Stream) -> Result<Self> {
        let mut oob_data = Vec::new();

        let mut found_pps = false;
        for nal in annexb.nal_iter() {
            if nal.kind == H264NalKind::SPS {
                let without_start = remove_nal_start(nal.data);
                ensure!(oob_data.is_empty(), "Too many SPS NAL units");
                let profile = without_start[1];
                let profile_compatibility = without_start[2];
                let level = without_start[3];
                let nalu_size_minus1 = 3;
                let num_sps_nals = 1;
                oob_data.extend(&[
                    1,
                    profile,
                    profile_compatibility,
                    level,
                    0xfc | nalu_size_minus1, // reserved | nalu length size
                    0xe0 | num_sps_nals,     // reserved | Num  SPS nals
                ]);

                oob_data.extend((without_start.len() as u16).to_be_bytes());
                oob_data.extend(without_start);
            } else if nal.kind == H264NalKind::PPS {
                ensure!(!found_pps, "Too many PPS NAL units");
                found_pps = true;
                let without_start = remove_nal_start(nal.data);
                let num_pps_nals = 1;
                oob_data.push(num_pps_nals);
                oob_data.extend((without_start.len() as u16).to_be_bytes());
                oob_data.extend(without_start);
            }
        }

        let nal_units = annexb
            .nal_iter()
            .filter_map(|nal_unit| {
                match nal_unit.kind {
                    H264NalKind::SPS | H264NalKind::PPS => return None,
                    _ => (),
                }
                let without_start = remove_nal_start(nal_unit.data);
                let mut chunk_data = Vec::new();
                chunk_data.extend((without_start.len() as u32).to_be_bytes());
                chunk_data.extend(without_start);
                Some((nal_unit.kind, chunk_data))
            })
            .collect();
        Ok(Self { nal_units, oob_data })
    }
}

impl ElementaryStream for H264AVCCStream {
    fn format_details(&self, version_ordinal: u64) -> FormatDetails {
        FormatDetails {
            format_details_version_ordinal: Some(version_ordinal),
            mime_type: Some(String::from("video/h264")),
            oob_bytes: Some(self.oob_data.clone()),
            ..Default::default()
        }
    }

    fn is_access_units(&self) -> bool {
        true
    }

    fn stream<'a>(&'a self) -> Box<dyn Iterator<Item = ElementaryStreamChunk> + 'a> {
        Box::new(self.nal_units.iter().map(|chunk| ElementaryStreamChunk {
            start_access_unit: true,
            known_end_access_unit: true,
            data: chunk.1.clone(),
            significance: match chunk.0 {
                H264NalKind::IDR | H264NalKind::NonIDR => {
                    Significance::Video(VideoSignificance::Picture)
                }
                _ => Significance::Video(VideoSignificance::NotPicture),
            },
            timestamp: None,
        }))
    }
}

#[derive(Debug)]
pub struct H264Nal<'a> {
    pub kind: H264NalKind,
    pub data: &'a [u8],
}

// Returns the NAL unit with a 3-byte or 4-byte start code removed.
fn remove_nal_start<'a>(input: &'a [u8]) -> &'a [u8] {
    let start_len = if input[2] == 1 { 3 } else { 4 };
    assert!(input[start_len - 1] == 1);
    &input[start_len..]
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum H264NalKind {
    NonIDR,
    IDR,
    SEI,
    SPS,
    PPS,
    Unknown,
}

const NON_IDR_PICTURE_CODE: u8 = 1;
const IDR_PICTURE_CODE: u8 = 5;
const SEI_CODE: u8 = 6;
const SPS_CODE: u8 = 7;
const PPS_CODE: u8 = 8;

impl H264NalKind {
    fn from_header(header: u8) -> Self {
        let kind = header & 0x1f;
        match kind {
            kind if kind == IDR_PICTURE_CODE => H264NalKind::IDR,
            kind if kind == NON_IDR_PICTURE_CODE => H264NalKind::NonIDR,
            kind if kind == SPS_CODE => H264NalKind::SPS,
            kind if kind == PPS_CODE => H264NalKind::PPS,
            kind if kind == SEI_CODE => H264NalKind::SEI,
            _ => H264NalKind::Unknown,
        }
    }
}

struct H264NalStart<'a> {
    /// Position in the h264 stream of the start.
    pos: usize,
    /// All the data from the start of the NAL onward.
    data: &'a [u8],
    kind: H264NalKind,
}

/// An iterator over NALs in an H264 stream.
struct H264NalIter<'a> {
    data: &'a [u8],
    pos: usize,
}

impl<'a> H264NalIter<'a> {
    fn next_nal(&self, pos: usize) -> Option<H264Nal<'a>> {
        // This won't need to search if pos already at a start code.
        let nal_start = self.next_nal_start(pos)?;
        // We search 3 bytes after the found nal's start, because that will
        // ensure we don't just find the same start code again.
        match self.next_nal_start(nal_start.pos + 3) {
            Some(next_start) => Some(H264Nal {
                kind: nal_start.kind,
                data: &nal_start.data[0..(next_start.pos - nal_start.pos)],
            }),
            None => Some(H264Nal { kind: nal_start.kind, data: nal_start.data }),
        }
    }

    fn next_nal_start(&self, pos: usize) -> Option<H264NalStart<'a>> {
        // This search size will find 3 and 4 byte start codes, and the
        // header value.
        const NAL_SEARCH_SIZE: usize = 5;

        let data = self.data.get(pos..)?;
        data.windows(NAL_SEARCH_SIZE).enumerate().find_map(|(i, candidate)| match candidate {
            [0, 0, 0, 1, h] | [0, 0, 1, h, _] => Some(H264NalStart {
                pos: pos + i,
                data: data.get(i..).expect("Getting slice starting where we just matched"),
                kind: H264NalKind::from_header(*h),
            }),
            _ => None,
        })
    }
}

impl<'a> Iterator for H264NalIter<'a> {
    type Item = H264Nal<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        let nal = self.next_nal(self.pos);
        self.pos += nal.as_ref().map(|n| n.data.len()).unwrap_or(0);
        nal
    }
}

#[derive(Debug)]
pub struct H264Chunk<'a> {
    /// Only NonIDR, IDR are used, or not set if last data in file needs to be emitted but IDR or
    /// NonIDR wasn't last in the file.
    pub nal_kind: Option<H264NalKind>,
    pub data: &'a [u8],
}

/// An iterator over groupings of NALs in an H264 stream.
struct H264ChunkIter<'a> {
    nal_iter: H264NalIter<'a>,
    chunk_start_pos: Option<usize>,
}

impl<'a> H264ChunkIter<'a> {
    fn create(data: &'a [u8]) -> H264ChunkIter<'a> {
        H264ChunkIter { nal_iter: H264NalIter { data: data, pos: 0 }, chunk_start_pos: None }
    }
}

impl<'a> Iterator for H264ChunkIter<'a> {
    type Item = H264Chunk<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        let chunk = loop {
            let maybe_nal = self.nal_iter.next();
            match maybe_nal {
                Some(nal) => {
                    if let None = self.chunk_start_pos {
                        self.chunk_start_pos = Some(self.nal_iter.pos - nal.data.len());
                    }
                    let queue_chunk = match nal.kind {
                        // picture
                        H264NalKind::IDR | H264NalKind::NonIDR => true,
                        // not picture
                        _ => SEPARATE_SPS_PPS,
                    };
                    if queue_chunk {
                        break Some(H264Chunk {
                            nal_kind: Some(nal.kind),
                            data: &self.nal_iter.data
                                [self.chunk_start_pos.unwrap()..self.nal_iter.pos],
                        });
                    } else {
                        continue;
                    }
                }
                None => {
                    if self.chunk_start_pos.expect("Zero NALs?") == self.nal_iter.data.len() {
                        // done
                        break None;
                    } else {
                        break Some(H264Chunk {
                            nal_kind: None,
                            data: &self.nal_iter.data[self.chunk_start_pos.unwrap()..],
                        });
                    }
                }
            }
        };
        *self.chunk_start_pos.as_mut().unwrap() +=
            chunk.as_ref().map(|c| c.data.len()).unwrap_or(0);
        chunk
    }
}

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

    #[fuchsia::test]
    fn test_nal_iter() {
        let sps_pps_idr = [
            0x00, 0x00, 0x00, 0x01, 0x67, 0x4d, 0x00, 0x28, 0xf4, 0x03, 0xc0, 0x11, 0x3f, 0x2a,
            0x00, 0x00, 0x00, 0x01, 0x68, 0xee, 0x38, 0x80, 0x00, 0x00, 0x00, 0x01, 0x65, 0x88,
            0x84, 0x01,
        ];

        let iter = H264NalIter { data: &sps_pps_idr, pos: 0 };

        let result: Vec<H264Nal<'_>> = iter.collect();

        assert_eq!(result[0].kind, H264NalKind::SPS);
        assert_eq!(result[1].kind, H264NalKind::PPS);
        assert_eq!(result[2].kind, H264NalKind::IDR);
    }

    #[fuchsia::test]
    fn test_chunk_iter() {
        let sps_pps_idr = [
            0x00, 0x00, 0x00, 0x01, 0x67, 0x4d, 0x00, 0x28, 0xf4, 0x03, 0xc0, 0x11, 0x3f, 0x2a,
            0x00, 0x00, 0x00, 0x01, 0x68, 0xee, 0x38, 0x80, 0x00, 0x00, 0x00, 0x01, 0x65, 0x88,
            0x84, 0x01,
        ];

        let iter = H264ChunkIter::create(&sps_pps_idr);

        let result: Vec<H264Chunk<'_>> = iter.collect();

        if SEPARATE_SPS_PPS {
            assert_eq!(result[0].nal_kind.expect("nal"), H264NalKind::SPS);
            assert_eq!(result[1].nal_kind.expect("nal"), H264NalKind::PPS);
            assert_eq!(result[2].nal_kind.expect("nal"), H264NalKind::IDR);
        } else {
            assert_eq!(result[0].nal_kind.expect("nal"), H264NalKind::IDR);
        }
    }
}