//
// Copyright (C) 2022 Vivienne Watermeier <vwatermeier@igalia.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v2.0.
// If a copy of the MPL was not distributed with this file, You can obtain one at
// <https://mozilla.org/MPL/2.0/>.
//
// SPDX-License-Identifier: MPL-2.0

use gst::{glib, subclass::prelude::*};
use gst_rtp::{prelude::*, subclass::prelude::*};
use std::{
    cmp,
    collections::VecDeque,
    io::{Cursor, Read, Seek, SeekFrom, Write},
    sync::Mutex,
};

use bitstream_io::{BitReader, BitWriter};
use once_cell::sync::Lazy;

use crate::av1::common::{
    err_flow, leb128_size, write_leb128, ObuType, SizedObu, CLOCK_RATE, ENDIANNESS,
};

static CAT: Lazy<gst::DebugCategory> = Lazy::new(|| {
    gst::DebugCategory::new(
        "rtpav1pay",
        gst::DebugColorFlags::empty(),
        Some("RTP AV1 Payloader"),
    )
});

// TODO: properly handle `max_ptime` and `min_ptime`

/// Information about the OBUs intended to be grouped into one packet
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct PacketOBUData {
    obu_count: usize,
    payload_size: u32,
    last_obu_fragment_size: Option<u32>,
    omit_last_size_field: bool,
    ends_temporal_unit: bool,
}

impl Default for PacketOBUData {
    fn default() -> Self {
        PacketOBUData {
            payload_size: 1, // 1 byte is used for the aggregation header
            omit_last_size_field: true,
            obu_count: 0,
            last_obu_fragment_size: None,
            ends_temporal_unit: false,
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
struct ObuData {
    info: SizedObu,
    bytes: Vec<u8>,
    offset: usize,
    dts: Option<gst::ClockTime>,
    pts: Option<gst::ClockTime>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
struct State {
    /// Holds header information and raw bytes for all received OBUs,
    /// as well as DTS and PTS
    obus: VecDeque<ObuData>,

    /// Indicates that the first element in the Buffer is an OBU fragment,
    /// left over from the previous RTP packet
    open_obu_fragment: bool,

    /// Indicates the next constructed packet will be the first in its sequence
    /// (Corresponds to `N` field in the aggregation header)
    first_packet_in_seq: bool,

    /// The last observed DTS if upstream does not provide DTS for each OBU
    last_dts: Option<gst::ClockTime>,
    /// The last observed PTS if upstream does not provide PTS for each OBU
    last_pts: Option<gst::ClockTime>,
}

#[derive(Debug, Default)]
pub struct RTPAv1Pay {
    state: Mutex<State>,
}

impl Default for State {
    fn default() -> Self {
        Self {
            obus: VecDeque::new(),
            open_obu_fragment: false,
            first_packet_in_seq: true,
            last_dts: None,
            last_pts: None,
        }
    }
}

impl RTPAv1Pay {
    fn reset(&self, state: &mut State) {
        gst::debug!(CAT, imp: self, "resetting state");

        *state = State::default();
    }

    /// Parses new OBUs, stores them in the state,
    /// and constructs and sends new RTP packets when appropriate.
    fn handle_new_obus(
        &self,
        state: &mut State,
        data: &[u8],
        dts: Option<gst::ClockTime>,
        pts: Option<gst::ClockTime>,
    ) -> Result<gst::BufferList, gst::FlowError> {
        let mut reader = Cursor::new(data);

        while reader.position() < data.len() as u64 {
            let obu_start = reader.position();
            let obu = SizedObu::parse(&mut BitReader::endian(&mut reader, ENDIANNESS))
                .map_err(err_flow!(self, buf_read))?;

            // tile lists and temporal delimiters should not be transmitted,
            // see section 5 of the RTP AV1 spec
            match obu.obu_type {
                // completely ignore tile lists
                ObuType::TileList => {
                    gst::log!(CAT, imp: self, "ignoring tile list OBU");
                    reader
                        .seek(SeekFrom::Current(
                            (obu.header_len + obu.leb_size + obu.size) as i64,
                        ))
                        .map_err(err_flow!(self, buf_read))?;
                }

                // keep these OBUs around for now so we know where temporal units end
                ObuType::TemporalDelimiter => {
                    if obu.size != 0 {
                        gst::element_imp_error!(
                            self,
                            gst::ResourceError::Read,
                            ["temporal delimiter OBUs should have empty payload"]
                        );
                        return Err(gst::FlowError::Error);
                    }
                    state.obus.push_back(ObuData {
                        info: obu,
                        bytes: Vec::new(),
                        offset: 0,
                        dts,
                        pts,
                    });
                }

                _ => {
                    let bytes_total = (obu.header_len + obu.size) as usize;
                    let mut bytes = vec![0; bytes_total];

                    // read header
                    reader
                        .seek(SeekFrom::Start(obu_start))
                        .map_err(err_flow!(self, buf_read))?;
                    reader
                        .read_exact(&mut bytes[0..(obu.header_len as usize)])
                        .map_err(err_flow!(self, buf_read))?;

                    // skip size field
                    bytes[0] &= !2_u8; // set `has_size_field` to 0
                    reader
                        .seek(SeekFrom::Current(obu.leb_size as i64))
                        .map_err(err_flow!(self, buf_read))?;

                    // read OBU bytes
                    reader
                        .read_exact(&mut bytes[(obu.header_len as usize)..bytes_total])
                        .map_err(err_flow!(self, buf_read))?;

                    state.obus.push_back(ObuData {
                        info: obu,
                        bytes,
                        offset: 0,
                        dts,
                        pts,
                    });
                }
            }
        }

        let mut list = gst::BufferList::new();
        {
            let list = list.get_mut().unwrap();
            while let Some(packet_data) = self.consider_new_packet(state, false) {
                let buffer = self.generate_new_packet(state, packet_data)?;
                list.add(buffer);
            }
        }

        Ok(list)
    }

    /// Look at the size the currently stored OBUs would require,
    /// as well as their temportal IDs to decide if it is time to construct a
    /// new packet, and what OBUs to include in it.
    ///
    /// If `true` is passed for `force`, packets of any size will be accepted,
    /// which is used in flushing the last OBUs after receiving an EOS for example.
    fn consider_new_packet(&self, state: &mut State, force: bool) -> Option<PacketOBUData> {
        gst::trace!(
            CAT,
            imp: self,
            "{} new packet, currently storing {} OBUs",
            if force { "forcing" } else { "considering" },
            state.obus.len()
        );

        let payload_limit = gst_rtp::calc_payload_len(self.obj().mtu(), 0, 0);

        // Create information about the packet that can be created now while iterating over the
        // OBUs and return this if a full packet can indeed be created now.
        let mut packet = PacketOBUData::default();
        let mut pending_bytes = 0;
        let mut required_ids = None::<(u8, u8)>;

        // figure out how many OBUs we can fit into this packet
        for obu in &state.obus {
            // for OBUs with extension headers, spatial and temporal IDs must be equal
            // to all other such OBUs in the packet
            let matching_obu_ids = |obu: &SizedObu, required_ids: &mut Option<(u8, u8)>| -> bool {
                if let Some((sid, tid)) = *required_ids {
                    sid == obu.spatial_id && tid == obu.temporal_id
                } else {
                    *required_ids = Some((obu.spatial_id, obu.temporal_id));
                    true
                }
            };

            let current = &obu.info;

            // should this packet be finished here?
            if current.obu_type == ObuType::TemporalDelimiter {
                // ignore the temporal delimiter, it is not supposed to be transmitted,
                // it will be skipped later when building the packet
                gst::log!(CAT, imp: self, "ignoring temporal delimiter OBU");

                if packet.obu_count > 0 {
                    packet.ends_temporal_unit = true;
                    if packet.obu_count > 3 {
                        packet.payload_size += pending_bytes;
                        packet.omit_last_size_field = false;
                    }
                    return Some(packet);
                }

                continue;
            } else if packet.payload_size >= payload_limit
                || (packet.obu_count > 0 && current.obu_type == ObuType::SequenceHeader)
                || !matching_obu_ids(current, &mut required_ids)
            {
                if packet.obu_count > 3 {
                    packet.payload_size += pending_bytes;
                    packet.omit_last_size_field = false;
                }
                return Some(packet);
            }

            // would the full OBU fit?
            if packet.payload_size + pending_bytes + current.full_size() <= payload_limit {
                packet.obu_count += 1;
                packet.payload_size += current.partial_size() + pending_bytes;
                pending_bytes = current.leb_size;
            }
            // would it fit without the size field?
            else if packet.obu_count < 3
                && packet.payload_size + pending_bytes + current.partial_size() <= payload_limit
            {
                packet.obu_count += 1;
                packet.payload_size += current.partial_size() + pending_bytes;

                return Some(packet);
            }
            // otherwise consider putting an OBU fragment
            else {
                let leb_size = if packet.obu_count < 3 {
                    0
                } else {
                    // assume the biggest possible OBU fragment,
                    // so if anything the size field will be smaller than expected
                    leb128_size(payload_limit - packet.payload_size) as u32
                };

                // is there even enough space to bother?
                if packet.payload_size + pending_bytes + leb_size + current.header_len
                    < payload_limit
                {
                    packet.obu_count += 1;
                    packet.last_obu_fragment_size =
                        Some(payload_limit - packet.payload_size - pending_bytes - leb_size);
                    packet.payload_size = payload_limit;
                    packet.omit_last_size_field = leb_size == 0;
                } else if packet.obu_count > 3 {
                    packet.payload_size += pending_bytes;
                }

                return Some(packet);
            }
        }

        if force && packet.obu_count > 0 {
            if packet.obu_count > 3 {
                packet.payload_size += pending_bytes;
                packet.omit_last_size_field = false;
            }
            Some(packet)
        } else {
            // if we ran out of OBUs with space in the packet to spare, wait a bit longer
            None
        }
    }

    /// Given the information returned by consider_new_packet(), construct and return
    /// new RTP packet, filled with those OBUs.
    fn generate_new_packet(
        &self,
        state: &mut State,
        packet: PacketOBUData,
    ) -> Result<gst::Buffer, gst::FlowError> {
        gst::log!(
            CAT,
            imp: self,
            "constructing new RTP packet with {} OBUs",
            packet.obu_count
        );

        // prepare the outgoing buffer
        let mut outbuf =
            gst::Buffer::new_rtp_with_sizes(packet.payload_size, 0, 0).map_err(|err| {
                gst::element_imp_error!(
                    self,
                    gst::ResourceError::Write,
                    ["Failed to allocate output buffer: {}", err]
                );

                gst::FlowError::Error
            })?;

        {
            // this block enforces that outbuf_mut is dropped before pushing outbuf
            let first_obu = state.obus.front().unwrap();
            if let Some(dts) = first_obu.dts {
                state.last_dts = Some(
                    state
                        .last_dts
                        .map_or(dts, |last_dts| cmp::max(last_dts, dts)),
                );
            }
            if let Some(pts) = first_obu.pts {
                state.last_pts = Some(
                    state
                        .last_pts
                        .map_or(pts, |last_pts| cmp::max(last_pts, pts)),
                );
            }

            let outbuf_mut = outbuf
                .get_mut()
                .expect("Failed to get mutable reference to outbuf");
            outbuf_mut.set_dts(state.last_dts);
            outbuf_mut.set_pts(state.last_pts);

            let mut rtp = gst_rtp::RTPBuffer::from_buffer_writable(outbuf_mut)
                .expect("Failed to create RTPBuffer");
            rtp.set_marker(packet.ends_temporal_unit);

            let payload = rtp
                .payload_mut()
                .expect("Failed to get mutable reference to RTP payload");
            let mut writer = Cursor::new(payload);

            {
                // construct aggregation header
                let w = if packet.omit_last_size_field && packet.obu_count < 4 {
                    packet.obu_count
                } else {
                    0
                };

                let aggr_header: [u8; 1] = [
                    (state.open_obu_fragment as u8) << 7 |                  // Z
                    ((packet.last_obu_fragment_size.is_some()) as u8) << 6 |  // Y
                    (w as u8) << 4 |                                        // W
                    (state.first_packet_in_seq as u8) << 3                  // N
                ; 1];

                writer
                    .write(&aggr_header)
                    .map_err(err_flow!(self, aggr_header_write))?;

                state.first_packet_in_seq = false;
            }

            // append OBUs to the buffer
            for _ in 1..packet.obu_count {
                let obu = loop {
                    let obu = state.obus.pop_front().unwrap();

                    if let Some(dts) = obu.dts {
                        state.last_dts = Some(
                            state
                                .last_dts
                                .map_or(dts, |last_dts| cmp::max(last_dts, dts)),
                        );
                    }
                    if let Some(pts) = obu.pts {
                        state.last_pts = Some(
                            state
                                .last_pts
                                .map_or(pts, |last_pts| cmp::max(last_pts, pts)),
                        );
                    }

                    // Drop temporal delimiter from here
                    if obu.info.obu_type != ObuType::TemporalDelimiter {
                        break obu;
                    }
                };

                write_leb128(
                    &mut BitWriter::endian(&mut writer, ENDIANNESS),
                    obu.info.size + obu.info.header_len,
                )
                .map_err(err_flow!(self, leb_write))?;
                writer
                    .write(&obu.bytes[obu.offset..])
                    .map_err(err_flow!(self, obu_write))?;
            }
            state.open_obu_fragment = false;

            {
                let last_obu = loop {
                    let obu = state.obus.front_mut().unwrap();

                    if let Some(dts) = obu.dts {
                        state.last_dts = Some(
                            state
                                .last_dts
                                .map_or(dts, |last_dts| cmp::max(last_dts, dts)),
                        );
                    }
                    if let Some(pts) = obu.pts {
                        state.last_pts = Some(
                            state
                                .last_pts
                                .map_or(pts, |last_pts| cmp::max(last_pts, pts)),
                        );
                    }

                    // Drop temporal delimiter from here
                    if obu.info.obu_type != ObuType::TemporalDelimiter {
                        break obu;
                    }
                    let _ = state.obus.pop_front().unwrap();
                };

                // do the last OBU separately
                // in this instance `obu_size` includes the header length
                let obu_size = if let Some(size) = packet.last_obu_fragment_size {
                    state.open_obu_fragment = true;
                    size
                } else {
                    last_obu.bytes.len() as u32 - last_obu.offset as usize as u32
                };

                if !packet.omit_last_size_field {
                    write_leb128(&mut BitWriter::endian(&mut writer, ENDIANNESS), obu_size)
                        .map_err(err_flow!(self, leb_write))?;
                }

                // if this OBU is not a fragment, handle it as usual
                if packet.last_obu_fragment_size.is_none() {
                    writer
                        .write(&last_obu.bytes[last_obu.offset..])
                        .map_err(err_flow!(self, obu_write))?;
                    let _ = state.obus.pop_front().unwrap();
                }
                // otherwise write only a slice, and update the element
                // to only contain the unwritten bytes
                else {
                    writer
                        .write(
                            &last_obu.bytes[last_obu.offset..last_obu.offset + obu_size as usize],
                        )
                        .map_err(err_flow!(self, obu_write))?;

                    let new_size = last_obu.bytes.len() as u32 - last_obu.offset as u32 - obu_size;
                    last_obu.info = SizedObu {
                        size: new_size,
                        header_len: 0,
                        leb_size: leb128_size(new_size) as u32,
                        is_fragment: true,
                        ..last_obu.info
                    };
                    last_obu.offset += obu_size as usize;
                }
            }
        }

        gst::log!(
            CAT,
            imp: self,
            "generated RTP packet of size {}",
            outbuf.size()
        );

        Ok(outbuf)
    }
}

#[glib::object_subclass]
impl ObjectSubclass for RTPAv1Pay {
    const NAME: &'static str = "GstRtpAv1Pay";
    type Type = super::RTPAv1Pay;
    type ParentType = gst_rtp::RTPBasePayload;
}

impl ObjectImpl for RTPAv1Pay {}

impl GstObjectImpl for RTPAv1Pay {}

impl ElementImpl for RTPAv1Pay {
    fn metadata() -> Option<&'static gst::subclass::ElementMetadata> {
        static ELEMENT_METADATA: Lazy<gst::subclass::ElementMetadata> = Lazy::new(|| {
            gst::subclass::ElementMetadata::new(
                "RTP AV1 payloader",
                "Codec/Payloader/Network/RTP",
                "Payload AV1 as RTP packets",
                "Vivienne Watermeier <vwatermeier@igalia.com>",
            )
        });

        Some(&*ELEMENT_METADATA)
    }

    fn pad_templates() -> &'static [gst::PadTemplate] {
        static PAD_TEMPLATES: Lazy<Vec<gst::PadTemplate>> = Lazy::new(|| {
            let sink_pad_template = gst::PadTemplate::new(
                "sink",
                gst::PadDirection::Sink,
                gst::PadPresence::Always,
                &gst::Caps::builder("video/x-av1")
                    .field("parsed", true)
                    .field("stream-format", "obu-stream")
                    .field("alignment", "obu")
                    .build(),
            )
            .unwrap();

            let src_pad_template = gst::PadTemplate::new(
                "src",
                gst::PadDirection::Src,
                gst::PadPresence::Always,
                &gst::Caps::builder("application/x-rtp")
                    .field("media", "video")
                    .field("payload", gst::IntRange::new(96, 127))
                    .field("clock-rate", CLOCK_RATE as i32)
                    .field("encoding-name", "AV1")
                    .build(),
            )
            .unwrap();

            vec![src_pad_template, sink_pad_template]
        });

        PAD_TEMPLATES.as_ref()
    }

    fn change_state(
        &self,
        transition: gst::StateChange,
    ) -> Result<gst::StateChangeSuccess, gst::StateChangeError> {
        gst::debug!(CAT, imp: self, "changing state: {}", transition);

        if matches!(transition, gst::StateChange::ReadyToPaused) {
            let mut state = self.state.lock().unwrap();
            self.reset(&mut state);
        }

        let ret = self.parent_change_state(transition);

        if matches!(transition, gst::StateChange::PausedToReady) {
            let mut state = self.state.lock().unwrap();
            self.reset(&mut state);
        }

        ret
    }
}

impl RTPBasePayloadImpl for RTPAv1Pay {
    fn set_caps(&self, _caps: &gst::Caps) -> Result<(), gst::LoggableError> {
        self.obj().set_options("video", true, "AV1", CLOCK_RATE);

        gst::debug!(CAT, imp: self, "setting caps");

        Ok(())
    }

    fn handle_buffer(&self, buffer: gst::Buffer) -> Result<gst::FlowSuccess, gst::FlowError> {
        gst::trace!(CAT, imp: self, "received buffer of size {}", buffer.size());

        let mut state = self.state.lock().unwrap();

        if buffer.flags().contains(gst::BufferFlags::DISCONT) {
            gst::debug!(CAT, imp: self, "buffer discontinuity");
            self.reset(&mut state);
        }

        let dts = buffer.dts();
        let pts = buffer.pts();

        let buffer = buffer.into_mapped_buffer_readable().map_err(|_| {
            gst::element_imp_error!(
                self,
                gst::ResourceError::Read,
                ["Failed to map buffer readable"]
            );

            gst::FlowError::Error
        })?;

        let list = self.handle_new_obus(&mut state, buffer.as_slice(), dts, pts)?;
        drop(state);

        if !list.is_empty() {
            self.obj().push_list(list)
        } else {
            Ok(gst::FlowSuccess::Ok)
        }
    }

    fn sink_event(&self, event: gst::Event) -> bool {
        gst::log!(CAT, imp: self, "sink event: {}", event.type_());

        match event.view() {
            gst::EventView::Eos(_) => {
                // flush all remaining OBUs
                let mut list = gst::BufferList::new();
                {
                    let mut state = self.state.lock().unwrap();
                    let list = list.get_mut().unwrap();

                    while let Some(packet_data) = self.consider_new_packet(&mut state, true) {
                        match self.generate_new_packet(&mut state, packet_data) {
                            Ok(buffer) => list.add(buffer),
                            Err(_) => break,
                        }
                    }

                    self.reset(&mut state);
                }
                if !list.is_empty() {
                    let _ = self.obj().push_list(list);
                }
            }
            gst::EventView::FlushStop(_) => {
                let mut state = self.state.lock().unwrap();
                self.reset(&mut state);
            }
            _ => (),
        }

        self.parent_sink_event(event)
    }
}

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

    #[test]
    fn test_consider_new_packet() {
        gst::init().unwrap();

        let base_obu = SizedObu {
            has_extension: false,
            has_size_field: true,
            leb_size: 1,
            header_len: 1,
            is_fragment: false,
            ..SizedObu::default()
        };

        let input_data = [
            (
                false, // force argument
                State {
                    // payloader state
                    obus: VecDeque::from(vec![
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Padding,
                                size: 3,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3],
                            ..ObuData::default()
                        },
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Frame,
                                size: 4,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3, 4],
                            ..ObuData::default()
                        },
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Frame,
                                size: 5,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3, 4, 5],
                            ..ObuData::default()
                        },
                        ObuData {
                            // last two OBUs should not be counted
                            info: SizedObu {
                                obu_type: ObuType::TemporalDelimiter,
                                size: 0,
                                ..base_obu
                            },
                            ..ObuData::default()
                        },
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Frame,
                                size: 10,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
                            ..ObuData::default()
                        },
                    ]),
                    ..State::default()
                },
            ),
            (
                true,
                State {
                    obus: VecDeque::from(vec![
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::TemporalDelimiter,
                                size: 0,
                                ..base_obu
                            },
                            ..ObuData::default()
                        },
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Frame,
                                size: 7,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3, 4, 5, 6, 7],
                            ..ObuData::default()
                        },
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Padding,
                                size: 6,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3, 4, 5, 6],
                            ..ObuData::default()
                        },
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Frame,
                                size: 9,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3, 4, 5, 6, 7, 8, 9],
                            ..ObuData::default()
                        },
                        ObuData {
                            info: SizedObu {
                                obu_type: ObuType::Frame,
                                size: 3,
                                ..base_obu
                            },
                            bytes: vec![1, 2, 3],
                            ..ObuData::default()
                        },
                    ]),
                    ..State::default()
                },
            ),
            (
                false,
                State {
                    obus: VecDeque::from(vec![ObuData {
                        info: SizedObu {
                            obu_type: ObuType::Frame,
                            size: 4,
                            ..base_obu
                        },
                        bytes: vec![1, 2, 3, 4],
                        ..ObuData::default()
                    }]),
                    ..State::default()
                },
            ),
        ];

        let results = [
            (
                Some(PacketOBUData {
                    obu_count: 3,
                    payload_size: 18,
                    last_obu_fragment_size: None,
                    omit_last_size_field: true,
                    ends_temporal_unit: true,
                }),
                State {
                    obus: VecDeque::from(vec![
                        input_data[0].1.obus[0].clone(),
                        input_data[0].1.obus[1].clone(),
                        input_data[0].1.obus[2].clone(),
                        input_data[0].1.obus[4].clone(),
                    ]),
                    ..input_data[0].1
                },
            ),
            (
                Some(PacketOBUData {
                    obu_count: 4,
                    payload_size: 34,
                    last_obu_fragment_size: None,
                    omit_last_size_field: false,
                    ends_temporal_unit: false,
                }),
                State {
                    obus: {
                        let mut copy = input_data[1].1.obus.clone();
                        copy.pop_front().unwrap();
                        copy
                    },
                    ..input_data[1].1
                },
            ),
            (None, input_data[2].1.clone()),
        ];

        let element = <RTPAv1Pay as ObjectSubclass>::Type::new();
        let pay = element.imp();
        for idx in 0..input_data.len() {
            println!("running test {}...", idx);

            let mut state = pay.state.lock().unwrap();
            *state = input_data[idx].1.clone();

            assert_eq!(
                pay.consider_new_packet(&mut state, input_data[idx].0),
                results[idx].0,
            );
            assert_eq!(
                state
                    .obus
                    .iter()
                    .filter(|o| o.info.obu_type != ObuType::TemporalDelimiter)
                    .cloned()
                    .collect::<Vec<_>>(),
                results[idx].1.obus.iter().cloned().collect::<Vec<_>>()
            );
            assert_eq!(state.open_obu_fragment, results[idx].1.open_obu_fragment);
            assert_eq!(
                state.first_packet_in_seq,
                results[idx].1.first_packet_in_seq
            );
        }
    }
}