transotf/src/av/xcoder.rs

use crate::av::dictionary::Dictionary;
use crate::av::stream::Stream;
use crate::av::{verify_response, Rational};
use crate::av_err2str;
use ffmpeg_sys_next::AVPixelFormat::AV_PIX_FMT_NONE;
use ffmpeg_sys_next::{
    av_malloc, avcodec_open2, AVCodecContext, AVColorPrimaries, AVColorRange,
    AVColorTransferCharacteristic, AVPixelFormat, AVRational, AVSampleFormat, AVERROR, AVERROR_EOF,
    EAGAIN,
};
use num_rational::Ratio;
use std::ffi::CStr;
use std::os::raw::c_int;
use std::ptr::null_mut;

pub trait XCoder: Sized {
    fn from_name(name: &str) -> Option<Self>;
    fn as_mut_ptr(&self) -> *mut AVCodecContext;

    #[inline]
    fn as_ptr(&self) -> *const AVCodecContext {
        self.as_mut_ptr()
    }

    fn name(&self) -> String {
        unsafe {
            if let Some(x) = (*(*self.as_ptr()).av_class).item_name {
                CStr::from_ptr(x(self.as_mut_ptr().cast()))
                    .to_str()
                    .unwrap_or("")
                    .to_string()
            } else {
                String::new()
            }
        }
    }

    fn configure(&self, stream: &Stream) -> Result<(), String>;

    fn open(&self) -> Result<(), String> {
        let ctx = self.as_mut_ptr();
        let resp = unsafe { avcodec_open2(ctx, (*ctx).codec, null_mut()) };

        if resp < 0 {
            return Err(format!(
                "Failed configuring coder from stream (code {}: {})",
                resp,
                av_err2str(resp)
            ));
        }

        Ok(())
    }

    #[inline]
    fn time_base(&self) -> Ratio<i32> {
        unsafe { (*self.as_ptr()).time_base.to_num() }
    }

    #[inline]
    fn set_time_base<R: Rational>(&self, ratio: R) {
        unsafe { (*self.as_mut_ptr()).time_base = ratio.to_av() }
    }

    #[inline]
    fn pixel_format(&self) -> AVPixelFormat {
        unsafe { (*self.as_ptr()).pix_fmt }
    }

    #[inline]
    fn set_pixel_format(&self, pixel_format: AVPixelFormat) {
        unsafe { (*self.as_mut_ptr()).pix_fmt = pixel_format }
    }

    #[inline]
    fn width(&self) -> c_int {
        unsafe { (*self.as_ptr()).width }
    }

    #[inline]
    fn set_width(&self, width: c_int) {
        unsafe { (*self.as_mut_ptr()).width = width }
    }

    #[inline]
    fn height(&self) -> c_int {
        unsafe { (*self.as_ptr()).height }
    }

    #[inline]
    fn set_height(&self, height: c_int) {
        unsafe { (*self.as_mut_ptr()).height = height }
    }

    #[inline]
    fn bit_rate(&self) -> i64 {
        unsafe { (*self.as_mut_ptr()).bit_rate }
    }

    #[inline]
    fn set_bit_rate(&self, bit_rate: i64) {
        unsafe {
            (*self.as_mut_ptr()).bit_rate = bit_rate;
        }
    }

    #[inline]
    fn sample_aspect_ratio(&self) -> AVRational {
        unsafe { (*self.as_ptr()).sample_aspect_ratio }
    }

    #[inline]
    fn color_range(&self) -> AVColorRange {
        unsafe { (*self.as_ptr()).color_range }
    }

    #[inline]
    fn set_color_range(&self, color_range: AVColorRange) {
        unsafe { (*self.as_mut_ptr()).color_range = color_range }
    }

    #[inline]
    fn color_primaries(&self) -> AVColorPrimaries {
        unsafe { (*self.as_ptr()).color_primaries }
    }

    #[inline]
    fn set_color_primaries(&self, color_primaries: AVColorPrimaries) {
        unsafe { (*self.as_mut_ptr()).color_primaries = color_primaries }
    }

    #[inline]
    fn color_trc(&self) -> AVColorTransferCharacteristic {
        unsafe { (*self.as_ptr()).color_trc }
    }

    #[inline]
    fn set_color_trc(&self, color_trc: AVColorTransferCharacteristic) {
        unsafe { (*self.as_mut_ptr()).color_trc = color_trc }
    }

    #[inline]
    fn channels(&self) -> i32 {
        unsafe { (*self.as_ptr()).channels }
    }

    #[inline]
    fn set_channels(&self, channels: i32) {
        unsafe { (*self.as_mut_ptr()).channels = channels }
    }

    #[inline]
    fn channel_layout(&self) -> u64 {
        unsafe { (*self.as_ptr()).channel_layout }
    }

    #[inline]
    fn set_channel_layout(&self, channel_layout: u64) {
        unsafe { (*self.as_mut_ptr()).channel_layout = channel_layout }
    }

    #[inline]
    fn sample_format(&self) -> AVSampleFormat {
        unsafe { (*self.as_ptr()).sample_fmt }
    }

    #[inline]
    fn set_sample_format(&self, sample_format: AVSampleFormat) {
        unsafe { (*self.as_mut_ptr()).sample_fmt = sample_format }
    }

    #[inline]
    fn frame_size(&self) -> i32 {
        unsafe { (*self.as_ptr()).frame_size }
    }

    #[inline]
    fn set_frame_size(&self, frame_size: i32) {
        unsafe { (*self.as_mut_ptr()).frame_size = frame_size }
    }

    #[inline]
    fn sample_rate(&self) -> i32 {
        unsafe { (*self.as_ptr()).sample_rate }
    }

    #[inline]
    fn set_sample_rate(&self, sample_rate: i32) {
        unsafe { (*self.as_mut_ptr()).sample_rate = sample_rate }
    }

    #[inline]
    fn extra_data(&self) -> Vec<u8> {
        unsafe {
            let extra_data_size = (*self.as_ptr()).extradata_size as usize;
            let mut extra_data = vec![0u8; extra_data_size];
            (*self.as_ptr())
                .extradata
                .copy_to(extra_data.as_mut_ptr(), extra_data_size);

            extra_data
        }
    }

    #[inline]
    fn set_extra_data(&self, extra_data: Vec<u8>) {
        unsafe {
            let new_extra_data: *mut u8 = av_malloc(extra_data.len()).cast();
            (*self.as_mut_ptr()).extradata = new_extra_data;
            (*self.as_mut_ptr()).extradata_size = extra_data.len() as i32;
        }
    }

    fn set_private_data(&self, key: &str, value: &str) {
        unsafe {
            Dictionary {
                ptr: (*self.as_ptr()).priv_data.cast(),
                owned: false,
            }
            .set(key, value)
        }
    }

    fn pixel_formats(&self) -> Vec<AVPixelFormat> {
        let mut pixel_formats = vec![];
        unsafe {
            let mut current = (*(*self.as_ptr()).codec).pix_fmts;
            while *current != AV_PIX_FMT_NONE {
                pixel_formats.push(*current);
                current = current.add(1);
            }
        }

        pixel_formats
    }

    fn sample_formats(&self) -> Vec<AVSampleFormat> {
        let mut sample_formats = vec![];
        unsafe {
            let mut current = (*(*self.as_ptr()).codec).sample_fmts;
            while *current != AVSampleFormat::AV_SAMPLE_FMT_NONE {
                sample_formats.push(*current);
                current = current.add(1);
            }
        }

        sample_formats
    }
}

pub(crate) fn process_return(error: &str, resp: i32) -> Result<bool, String> {
    if resp == AVERROR(EAGAIN) || resp == AVERROR_EOF {
        Ok(false)
    } else {
        verify_response(error, resp).and(Ok(true))
    }
}