transotf/src/av/format.rs

use crate::av::as_ptr::{AsMutPtr, AsMutVoidPtr, AsPtr, AsVoidPtr};
use crate::av::avio::{AVIOWriter, AVIO};
use crate::av::dictionary::Dictionary;
use crate::av::options::Options;
use crate::av::packet::Packet;
use crate::av::stream::Stream;
use crate::av::{decoder_codec_from_name, verify_response};
use ffmpeg_sys_next::{
    av_interleaved_write_frame, av_read_frame, av_write_frame, av_write_trailer,
    avformat_alloc_context, avformat_alloc_output_context2, avformat_find_stream_info,
    avformat_init_output, avformat_new_stream, avformat_open_input, avformat_write_header,
    avio_open2, AVFormatContext, AVMediaType, AVIO_FLAG_WRITE,
};
use std::ffi::CString;
use std::fmt::{Debug, Error, Formatter};
use std::os::raw::{c_int, c_void};
use std::ptr::{null, null_mut};

#[derive(Clone)]
pub struct Format<T> {
    pub ctx: *mut AVFormatContext,
    pub avio: Option<AVIO<T>>,
}

impl<T> Debug for Format<T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        f.debug_struct("Format<_>").field("ctx", &self.ctx).finish()
    }
}

impl<T> Format<T> {
    pub fn flags(&self) -> c_int {
        unsafe { (*self.as_ptr()).flags }
    }

    pub fn set_flags(&self, flags: c_int) {
        unsafe { (*self.as_mut_ptr()).flags = flags }
    }
}

impl<T> Format<T> {
    pub fn stream(&self, stream_type: AVMediaType, index: Option<i32>) -> Option<Stream> {
        unsafe {
            let ctx = self.ctx;
            let mut stream = (*ctx).streams;

            for _ in 0..(*ctx).nb_streams {
                let curr_stream = *stream;
                if (*(*curr_stream).codecpar).codec_type == stream_type {
                    if index.is_none() || Some((*curr_stream).index) == index {
                        return Some(Stream(curr_stream));
                    }
                }

                stream = stream.add(1)
            }
        }

        None
    }

    pub fn as_mut_ptr(&self) -> *mut AVFormatContext {
        self.ctx
    }

    pub fn next_packet(&self) -> Option<Packet> {
        let packet = Packet::alloc();
        return if unsafe { av_read_frame(self.as_mut_ptr(), packet.as_mut_ptr()) } >= 0 {
            Some(packet)
        } else {
            None
        };
    }

    pub fn write_header(&self, options: Option<Dictionary>) -> Result<(), String> {
        verify_response("failed to write header to output", unsafe {
            avformat_write_header(
                self.as_mut_ptr(),
                if let Some(mut options) = options {
                    &mut options.as_mut_ptr()
                } else {
                    null_mut()
                },
            )
        })?;

        Ok(())
    }

    pub fn write_trailer(&self) -> Result<(), String> {
        verify_response("failed to write trailer to output", unsafe {
            av_write_trailer(self.as_mut_ptr())
        })?;

        Ok(())
    }

    pub fn write_packet_null(&self) -> Result<(), String> {
        verify_response("failed to write to output", unsafe {
            av_write_frame(self.as_mut_ptr(), null_mut())
        })?;

        Ok(())
    }

    pub fn write_packet_interleaved(&self, packet: &Packet) -> Result<(), String> {
        verify_response("failed to write interleaved to output", unsafe {
            av_interleaved_write_frame(self.as_mut_ptr(), packet.as_mut_ptr())
        })?;

        Ok(())
    }

    pub fn write_packet(&self, packet: &Packet) -> Result<(), String> {
        verify_response("failed to write to output", unsafe {
            av_write_frame(self.as_mut_ptr(), packet.as_mut_ptr())
        })?;

        Ok(())
    }

    pub fn new_stream(&self, codec: &str) -> Result<Stream, String> {
        let codec =
            decoder_codec_from_name(codec).ok_or(format!("No codec found with name {}", codec))?;
        let stream = unsafe { avformat_new_stream(self.ctx, codec) };

        Ok(Stream(stream))
    }

    pub fn avio_inner(&self) -> Option<&T> {
        self.avio.as_ref().map(|x| x.inner())
    }

    pub fn avio_inner_mut(&mut self) -> Option<&mut T> {
        self.avio.as_mut().map(|x| x.inner_mut())
    }

    pub fn init_output(&self, mut options: Dictionary) -> Result<(), String> {
        verify_response("Failed to init output", unsafe {
            avformat_init_output(self.as_mut_ptr(), &mut options.as_mut_ptr())
        })?;
        Ok(())
    }
}

impl<T: AVIOWriter> Format<T> {
    pub fn output_avio<'a>(avio: AVIO<T>, format: &str) -> Result<Format<T>, String> {
        let mut context = unsafe { avformat_alloc_context() };
        let format_c = CString::new(format).unwrap();

        verify_response("Failed to open output with AVIO", unsafe {
            avformat_alloc_output_context2(&mut context, null_mut(), format_c.as_ptr(), null_mut())
        })?;

        unsafe {
            (*context).pb = avio.as_mut_ptr();
        }

        Ok(Format {
            ctx: context,
            avio: Some(avio),
        })
    }
}

type FS = ();

impl Format<FS> {
    pub fn output(
        path: &str,
        format: Option<&str>,
        options: Option<Dictionary>,
    ) -> Result<Format<FS>, String> {
        let mut context = unsafe { avformat_alloc_context() };

        let path_c = CString::new(path).unwrap();
        let format_c = if let Some(fmt) = format {
            CString::new(fmt).unwrap()
        } else {
            CString::default()
        };

        let resp = unsafe {
            avformat_alloc_output_context2(
                &mut context,
                null_mut(),
                if format.is_none() {
                    null()
                } else {
                    format_c.as_ptr()
                },
                path_c.as_ptr(),
            )
        };

        if resp < 0 {
            return Err(format!(
                "Failed to open output (path: {}, format: {:?})",
                path, format
            ));
        }

        verify_response("Failed to open output", unsafe {
            avio_open2(
                &mut (*context).pb,
                path_c.as_ptr(),
                AVIO_FLAG_WRITE,
                null(),
                if let Some(mut options) = options {
                    options.disown();
                    &mut options.as_mut_ptr()
                } else {
                    null_mut()
                },
            )
        })?;

        Ok(Format {
            ctx: context,
            avio: None,
        })
    }

    pub fn open(path: &str, options: Option<Dictionary>) -> Result<Format<FS>, String> {
        unsafe {
            let mut context = avformat_alloc_context();
            let c_path = CString::new(path).unwrap();
            verify_response(
                "failed to open file",
                avformat_open_input(
                    &mut context,
                    c_path.as_ptr(),
                    null_mut(),
                    if let Some(mut options) = options {
                        &mut options.as_mut_ptr()
                    } else {
                        null_mut()
                    },
                ),
            )?;
            verify_response(
                "failed to find stream info",
                avformat_find_stream_info(context, null_mut()),
            )?;

            Ok(Format {
                ctx: context,
                avio: None,
            })
        }
    }
}

impl<T> AsPtr<AVFormatContext> for Format<T> {
    #[inline]
    fn as_ptr(&self) -> *const AVFormatContext {
        self.ctx
    }
}

impl<T> AsMutPtr<AVFormatContext> for Format<T> {
    #[inline]
    fn as_mut_ptr(&self) -> *mut AVFormatContext {
        self.ctx
    }
}

impl<T> AsVoidPtr for Format<T> {
    #[inline]
    fn as_void_ptr(&self) -> *const c_void {
        self.as_ptr().cast()
    }
}

impl<T> AsMutVoidPtr for Format<T> {
    #[inline]
    fn as_mut_void_ptr(&self) -> *mut c_void {
        self.as_mut_ptr().cast()
    }
}

impl<T> Options for Format<T> {}