sentryshot_util 0.1.2

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

use crate::{
    pixfmt::{
        frame_raw_height, PixelFormat, PIX_FMT_FLAG_BITSTREAM, PIX_FMT_FLAG_HWACCEL,
        PIX_FMT_FLAG_PAL,
    },
    Frame,
};
use std::num::NonZeroU16;
use thiserror::Error;

//#define FFALIGN(x, a) (((x)+(a)-1)&~((a)-1))
#[macro_export]
macro_rules! ffalign {
    ($x:expr, $a:expr) => {
        (($x) + ($a) - 1) & !(($a) - 1)
    };
}

pub fn image_fill_max_pixsteps(
    max_pixsteps: &mut [usize; 4],
    max_pixstep_comps: &mut [usize; 4],
    pix_fmt: Option<PixelFormat>,
) {
    let Some(pix_fmt) = pix_fmt else { return };
    for i in 0..pix_fmt.comps().len() {
        let comp = &pix_fmt.comps()[i];
        let plane = usize::from(comp.plane());
        let step = usize::from(comp.step());
        if step > max_pixsteps[plane] {
            max_pixsteps[plane] = step;
            max_pixstep_comps[plane] = i;
        }
    }
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum GetImageLinesizeError {
    #[error("max step")]
    MaxStep,
}

// Compute the size of an image line with format
// pix_fmt and width for the plane plane.
// Returns the computed size in bytes
fn av_get_image_linesize(
    pix_fmt: PixelFormat,
    width: NonZeroU16,
    plane: u8,
) -> Result<usize, GetImageLinesizeError> {
    // Max pixel step for each plane.
    let mut max_step = [0, 0, 0, 0];

    // The component for each plane which has the max pixel step.
    let mut max_step_comp = [0, 0, 0, 0];

    image_fill_max_pixsteps(&mut max_step, &mut max_step_comp, Some(pix_fmt));

    let plane = usize::from(plane);
    get_image_linesize(width, max_step[plane], max_step_comp[plane], pix_fmt)
}

fn get_image_linesize(
    width: NonZeroU16,
    max_step: usize,
    max_step_comp: usize,
    pix_fmt: PixelFormat,
) -> Result<usize, GetImageLinesizeError> {
    let s = if max_step_comp == 1 || max_step_comp == 2 {
        pix_fmt.log2_chroma_w()
    } else {
        0
    };

    let shifted_w = (usize::from(width.get()) + (1 << s) - 1) >> s;

    if shifted_w != 0 && max_step > usize::MAX / shifted_w {
        return Err(GetImageLinesizeError::MaxStep);
    }

    let linesize = max_step * shifted_w;

    if pix_fmt.flags() & PIX_FMT_FLAG_BITSTREAM != 0 {
        return Ok((linesize + 7) >> 3);
    }

    Ok(linesize)
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum ImageFillLinesizesError {
    #[error("hwaccel flag present")]
    HwaccelFlagPresent,

    #[error("get image linesize: {0}")]
    GetImageLinesize(#[from] GetImageLinesizeError),
}

fn image_fill_linesizes(
    linesizes: &mut [usize; 8],
    pix_fmt: PixelFormat,
    width: NonZeroU16,
) -> Result<(), ImageFillLinesizesError> {
    if pix_fmt.flags() & PIX_FMT_FLAG_HWACCEL != 0 {
        return Err(ImageFillLinesizesError::HwaccelFlagPresent);
    }

    // Max pixel step for each plane.
    let mut max_step = [0; 4];

    // The component for each plane which has the max pixel step.
    let mut max_step_comp = [0; 4];

    image_fill_max_pixsteps(&mut max_step, &mut max_step_comp, Some(pix_fmt));

    for i in 0..4 {
        linesizes[i] = get_image_linesize(width, max_step[i], max_step_comp[i], pix_fmt)?;
    }

    Ok(())
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum ImageFillPlaneSizesError {
    #[error("hwaccel flag present")]
    HwaccelFlagPresent,

    #[error("max size: {0}")]
    MaxSize(usize),
}

fn image_fill_plane_sizes(
    sizes: &mut [usize; 8],
    pix_fmt: PixelFormat,
    height: NonZeroU16,
    linesizes: [usize; 8],
) -> Result<(), ImageFillPlaneSizesError> {
    if pix_fmt.flags() & PIX_FMT_FLAG_HWACCEL != 0 {
        return Err(ImageFillPlaneSizesError::HwaccelFlagPresent);
    }

    if linesizes[0] > usize::MAX / usize::from(height.get()) {
        return Err(ImageFillPlaneSizesError::MaxSize(linesizes[0]));
    }

    sizes[0] = linesizes[0] * usize::from(height.get());

    if pix_fmt.flags() & PIX_FMT_FLAG_PAL != 0 {
        // palette is stored here as 256 32 bits words.
        sizes[1] = 256 * 4;
        return Ok(());
    }

    let mut has_plane = [false; 4];
    for comp in pix_fmt.comps() {
        has_plane[usize::from(comp.plane())] = true;
    }

    for i in 1..4_u8 {
        let i2 = usize::from(i);
        if !has_plane[i2] {
            break;
        }

        let h = usize::from(frame_raw_height(pix_fmt, height, i));

        if linesizes[i2] > usize::MAX / h {
            return Err(ImageFillPlaneSizesError::MaxSize(linesizes[i2]));
        }
        sizes[i2] = h * linesizes[i2];
    }

    Ok(())
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum ImageFillDataError {
    #[error("fill plane sizes: {0}")]
    FillPlaneSizes(#[from] ImageFillPlaneSizesError),

    #[error("size too big")]
    SizeTooBig,
}

// Fill plane data for an image with pixel format pix_fmt and height.
fn image_fill_data(
    data: &mut [Vec<u8>; 8],
    pix_fmt: PixelFormat,
    height: NonZeroU16,
    src: &[u8],
    linesizes: &[usize; 8],
) -> Result<usize, ImageFillDataError> {
    let mut linesizes1 = [0; 8];
    linesizes1[..4].copy_from_slice(&linesizes[..4]);

    let mut sizes = [0; 8];
    image_fill_plane_sizes(&mut sizes, pix_fmt, height, linesizes1)?;

    let mut ret = 0;
    for size in sizes {
        if size > usize::MAX - ret {
            return Err(ImageFillDataError::SizeTooBig);
        }
        ret += size;
    }

    let mut offset = 0;
    for i in 0..4 {
        let size = sizes[i];
        data[i].clear();
        data[i].extend_from_slice(&src[offset..offset + size]);
        offset += size;
    }

    Ok(ret)
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum ResetBufferError {
    #[error("check size: {0}")]
    CheckSize(#[from] ImageCheckSizeError),

    #[error("fill linesizes: {0}")]
    FillLinesizes(#[from] ImageFillLinesizesError),

    #[error("fill plane sizes: {0}")]
    FillPlaneSizes(#[from] ImageFillPlaneSizesError),

    #[error("size too big")]
    SizeToBig,

    #[error("fill data: {0}")]
    FillData(#[from] ImageFillDataError),

    #[error("pal unsupported")]
    PalUnsupported,
}

// Allocate an image with size w and h and pixel format pix_fmt, and
// fill pointers and linesizes accordingly.
pub(crate) fn image_alloc(
    frame: &mut Frame,
    width: NonZeroU16,
    height: NonZeroU16,
    pix_fmt: PixelFormat,
    align: u8,
) -> Result<(), ResetBufferError> {
    image_check_size(width, height)?;

    let linesizes = frame.linesize_mut();
    let w = if align > 7 {
        NonZeroU16::new(ffalign!(width.get(), 8)).expect("aligned value to be non-zero")
    } else {
        width
    };
    image_fill_linesizes(linesizes, pix_fmt, w)?;

    let mut linesizes1 = [0; 8];
    for i in 0..4 {
        linesizes[i] = ffalign!(linesizes[i], usize::from(align));
        linesizes1[i] = linesizes[i];
    }

    let mut sizes = [0; 8];
    image_fill_plane_sizes(&mut sizes, pix_fmt, height, linesizes1)?;

    let mut total_size: usize = align.into();
    #[allow(clippy::needless_range_loop)]
    for i in 0..4 {
        if total_size > usize::MAX - sizes[i] {
            return Err(ResetBufferError::SizeToBig);
        }
        total_size += sizes[i];
    }

    let buf = vec![0; total_size];
    let linesizes = linesizes.to_owned();
    let data = frame.data_mut();
    image_fill_data(data, pix_fmt, height, &buf, &linesizes)?;

    if pix_fmt.flags() & PIX_FMT_FLAG_PAL != 0 {
        return Err(ResetBufferError::PalUnsupported);
    };
    /*if (desc->flags & AV_PIX_FMT_FLAG_PAL) {
        avpriv_set_systematic_pal2((uint32_t*)pointers[1], pix_fmt);
        if (align < 4) {
            av_log(NULL, AV_LOG_ERROR, "Formats with a palette require a minimum alignment of 4\n");
            av_free(buf);
            return AVERROR(EINVAL);
        }
    }

    if (desc->flags & AV_PIX_FMT_FLAG_PAL && pointers[1] &&
        pointers[1] - pointers[0] > linesizes[0] * h) {
        /* zero-initialize the padding before the palette */
        memset(pointers[0] + linesizes[0] * h, 0,
               pointers[1] - pointers[0] - linesizes[0] * h);
    }*/

    Ok(())
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum ImageCheckSizeError {
    #[error("picture size {0}x{1} is invalid")]
    PictureSizeInvalid(NonZeroU16, NonZeroU16),

    #[error("picture size {0}x{1} exceeds specified max pixel count")]
    MaxPixels(NonZeroU16, NonZeroU16),
}

fn image_check_size(w: NonZeroU16, h: NonZeroU16) -> Result<(), ImageCheckSizeError> {
    image_check_size2(w, h, u32::MAX, None)
}

fn image_check_size2(
    w: NonZeroU16,
    h: NonZeroU16,
    max_pixels: u32,
    pix_fmt: Option<PixelFormat>,
) -> Result<(), ImageCheckSizeError> {
    let mut stride: usize = {
        if let Some(pix_fmt) = pix_fmt {
            match av_get_image_linesize(pix_fmt, w, 0) {
                Ok(v) => v,
                Err(_) => 8 * usize::from(w.get()),
            }
        } else {
            8 * usize::from(w.get())
        }
    };

    stride += 128 * 8;

    if w.get() == 0
        || h.get() == 0
        || u64::try_from(stride).expect("fit u64") * u64::from(h.get() + 128) >= u32::MAX.into()
    {
        return Err(ImageCheckSizeError::PictureSizeInvalid(w, h));
    }

    if max_pixels < u32::MAX && (u32::from(w.get()) * u32::from(h.get())) > max_pixels {
        return Err(ImageCheckSizeError::MaxPixels(w, h));
    }

    Ok(())
}

#[derive(Debug, Error)]
pub enum ImageFillArraysError {
    #[error("check size: {0}")]
    CheckSize(#[from] ImageCheckSizeError),

    #[error("fill linesizes: {0}")]
    FillLinesizes(#[from] ImageFillLinesizesError),

    #[error("fill data: {0}")]
    FillData(#[from] ImageFillDataError),
}

pub(crate) fn image_fill_arrays(
    dst_data: &mut [Vec<u8>; 8],
    dst_linesize: &mut [usize; 8],
    src: &[u8],
    pix_fmt: PixelFormat,
    width: NonZeroU16,
    height: NonZeroU16,
    align: u8,
) -> Result<(), ImageFillArraysError> {
    image_check_size(width, height)?;
    image_fill_linesizes(dst_linesize, pix_fmt, width)?;

    #[allow(clippy::needless_range_loop)]
    for i in 0..4 {
        dst_linesize[i] = ffalign!(dst_linesize[i], usize::from(align));
    }

    image_fill_data(dst_data, pix_fmt, height, src, dst_linesize)?;

    Ok(())
}

#[derive(Debug, Error)]
pub enum ImageBufferSizeError {
    #[error("check size: {0}")]
    CheckSize(#[from] ImageCheckSizeError),

    #[error("fill linesizes: {0}")]
    FillLinesize(#[from] ImageFillLinesizesError),

    #[error("fill plane sizes: {0}")]
    FillPlaneSizes(#[from] ImageFillPlaneSizesError),

    #[error("max size: {0}")]
    MaxSize(usize),
}

pub(crate) fn image_buffer_size(
    pix_fmt: PixelFormat,
    width: NonZeroU16,
    height: NonZeroU16,
    align: u8,
) -> Result<usize, ImageBufferSizeError> {
    image_check_size(width, height)?;

    let mut linesize = [0; 8];
    image_fill_linesizes(&mut linesize, pix_fmt, width)?;

    let mut aligned_linesize = [0; 8];
    for i in 0..4 {
        aligned_linesize[i] = ffalign!(linesize[i], usize::from(align));
    }

    let mut sizes = [0; 8];
    image_fill_plane_sizes(&mut sizes, pix_fmt, height, aligned_linesize)?;

    let mut ret = 0;
    for size in sizes {
        if size > usize::MAX - ret {
            return Err(ImageBufferSizeError::MaxSize(size));
        }
        ret += size;
    }

    Ok(ret)
}

#[derive(Debug, Error)]
pub enum ImageCopyToBufferError {
    #[error("buffer size: {0}")]
    BufferSize(#[from] ImageBufferSizeError),

    #[error("fill linesizes: {0}")]
    FillLinesizes(#[from] ImageFillLinesizesError),

    #[error("PAL pixel formats are unsupported")]
    PalUnsupported,
}

pub(crate) fn image_copy_to_buffer(
    frame: &Frame,
    dst: &mut Vec<u8>,
    align: u8,
) -> Result<(), ImageCopyToBufferError> {
    let size = image_buffer_size(frame.pix_fmt(), frame.width(), frame.height(), align)?;

    dst.resize(size, 0);

    let mut nb_planes = 0;
    for comp in frame.pix_fmt().comps() {
        nb_planes = std::cmp::max(comp.plane(), nb_planes);
    }

    nb_planes += 1;

    let src_linesize = frame.linesize();

    let mut dst_linesize = [0; 8];
    image_fill_linesizes(&mut dst_linesize, frame.pix_fmt(), frame.width())?;

    let mut dst_offset: usize = 0;

    for i in 0..nb_planes {
        let i2 = usize::from(i);
        let src = &frame.data()[i2];
        let src_size = src_linesize[i2];
        let dst_size = dst_linesize[i2];

        let mut src_offset = 0;
        for _ in 0..frame_raw_height(frame.pix_fmt(), frame.height(), i) {
            let dst_range = dst_offset..dst_offset + dst_size;
            let src_range = src_offset..src_offset + dst_size;
            dst[dst_range].copy_from_slice(&src[src_range]);

            dst_offset += ffalign!(dst_size, usize::from(align));
            src_offset += src_size;
        }
    }

    // "pix_fmt PAL stuff is black magic. I think only very few people understand it"
    if frame.pix_fmt().flags() & PIX_FMT_FLAG_PAL != 0 {
        return Err(ImageCopyToBufferError::PalUnsupported);
        //for i in 0..256 {}

        /*uint32_t *d32 = (uint32_t *)dst;

        for (i = 0; i<256; i++)
            AV_WL32(d32 + i, AV_RN32(src_data[1] + 4*i));*/
    }

    Ok(())
}