sentryshot_util 0.1.2

// SPDX-License-Identifier: MPL-2.0+

use crate::{
    imgutils::{
        image_alloc, image_buffer_size, image_copy_to_buffer, image_fill_arrays,
        ImageBufferSizeError, ImageCopyToBufferError, ImageFillArraysError, ResetBufferError,
    },
    pixfmt::PixelFormat,
    ColorRange,
};
use std::num::NonZeroU16;

pub const NUM_DATA_POINTERS: usize = 8;

// Vectors of the picture/channel planes.
pub type FrameData = [Vec<u8>; NUM_DATA_POINTERS];

// The sizes of each line in the `FrameData`
pub type Linesize = [usize; NUM_DATA_POINTERS];

/// This structure describes decoded (raw) audio or video data.
pub struct Frame {
    data: FrameData,
    linesize: Linesize,

    // The coded dimensions (in pixels) of the video frame, * i.e. the size of
    // the rectangle that contains some well-defined values.
    width: NonZeroU16,
    height: NonZeroU16,

    // Presentation timestamp in time_base units (time when frame should be shown to user).
    pts: i64,

    pix_fmt: PixelFormat,

    //MPEG vs JPEG YUV range
    color_range: ColorRange,
}

impl Default for Frame {
    fn default() -> Self {
        Self::new()
    }
}

impl Frame {
    #[must_use]
    #[allow(clippy::unwrap_used, clippy::missing_panics_doc)]
    pub fn new() -> Self {
        Self {
            data: FrameData::default(),
            linesize: Linesize::default(),
            width: NonZeroU16::new(1).unwrap(),
            height: NonZeroU16::new(1).unwrap(),
            pts: 0,
            pix_fmt: PixelFormat::YUV420P,
            color_range: ColorRange::UNSPECIFIED,
        }
    }

    pub fn from_raw(
        src: &[u8],
        pix_fmt: PixelFormat,
        width: NonZeroU16,
        height: NonZeroU16,
        align: u8,
    ) -> Result<Self, ImageFillArraysError> {
        let mut frame = Self::new();
        frame.width = width;
        frame.height = height;
        frame.pix_fmt = pix_fmt;

        image_fill_arrays(
            &mut frame.data,
            &mut frame.linesize,
            src,
            pix_fmt,
            width,
            height,
            align,
        )?;
        Ok(frame)
    }

    #[must_use]
    pub fn data(&self) -> &FrameData {
        &self.data
    }
    pub fn data_mut(&mut self) -> &mut FrameData {
        &mut self.data
    }

    #[must_use]
    pub fn linesize(&self) -> &Linesize {
        &self.linesize
    }
    pub fn linesize_mut(&mut self) -> &mut Linesize {
        &mut self.linesize
    }

    pub fn line_and_data_mut(&mut self) -> (&mut Linesize, &mut FrameData) {
        (&mut self.linesize, &mut self.data)
    }

    #[must_use]
    pub fn width(&self) -> NonZeroU16 {
        self.width
    }
    pub fn set_width(&mut self, width: NonZeroU16) {
        self.width = width;
    }

    #[must_use]
    pub fn height(&self) -> NonZeroU16 {
        self.height
    }
    pub fn set_height(&mut self, width: NonZeroU16) {
        self.height = width;
    }

    #[must_use]
    pub fn pts(&self) -> i64 {
        self.pts
    }
    pub fn set_pts(&mut self, pts: i64) {
        self.pts = pts;
    }

    #[must_use]
    pub fn pix_fmt(&self) -> PixelFormat {
        self.pix_fmt
    }
    pub fn set_pix_fmt(&mut self, pix_fmt: PixelFormat) {
        self.pix_fmt = pix_fmt;
    }

    #[must_use]
    pub fn color_range(&self) -> ColorRange {
        self.color_range
    }
    pub fn set_color_range(&mut self, color_range: ColorRange) {
        self.color_range = color_range;
    }

    /// Copy image data from an image into a buffer. `buffer_size()`
    /// can be used to compute the required size for the buffer to fill.
    pub fn copy_to_buffer(
        &self,
        dst: &mut Vec<u8>,
        align: u8,
    ) -> Result<(), ImageCopyToBufferError> {
        image_copy_to_buffer(self, dst, align)
    }

    pub fn buffer_size(&self, align: u8) -> Result<usize, ImageBufferSizeError> {
        image_buffer_size(self.pix_fmt(), self.width(), self.height(), align)
    }

    pub fn reset_buffer(
        &mut self,
        width: NonZeroU16,
        height: NonZeroU16,
        pix_fmt: PixelFormat,
        align: u8,
    ) -> Result<(), ResetBufferError> {
        image_alloc(self, width, height, pix_fmt, align)
    }
}

impl std::fmt::Display for Frame {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "width={}, height={}, pix_fmt={}",
            self.width, self.height, self.pix_fmt
        )
    }
}

impl std::fmt::Debug for Frame {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "width={}, height={}, pix_fmt={}, linesize={:?}",
            self.width, self.height, self.pix_fmt, self.linesize
        )
    }
}

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

    const YUV420P_FRAME: &[u8] = include_bytes!("../../testdata/yuv420p_64x64.bin");
    #[test]
    fn test_raw() {
        let frame = Frame::from_raw(
            YUV420P_FRAME,
            PixelFormat::YUV420P,
            NonZeroU16::new(64).unwrap(),
            NonZeroU16::new(64).unwrap(),
            1,
        )
        .unwrap();

        let mut raw = Vec::new();
        frame.copy_to_buffer(&mut raw, 1).unwrap();

        assert_eq!(YUV420P_FRAME, raw);
    }
}