transotf/src/av/avio.rs

use crate::av::as_ptr::{AsMutPtr, AsPtr};
use ffmpeg_sys_next::{av_malloc, avio_alloc_context, avio_flush, AVIOContext};
use std::cmp::min;
use std::os::raw::c_void;
use std::vec::Vec;

#[derive(Clone)]
pub struct AVIO<T> {
    ctx: *mut AVIOContext,
    inner: Box<AVIOInner<T>>,
}

impl<T> AVIO<T> {
    pub fn flush(&self) {
        unsafe {
            avio_flush(self.as_mut_ptr());
        }
    }
}

#[derive(Clone)]
struct AVIOInner<T> {
    data: Box<T>,
    read: Option<fn(&mut T, i32) -> Vec<u8>>,
    write: Option<fn(&mut T, Vec<u8>)>,
    seek: Option<fn(&mut T, i64, i32)>,
}

unsafe extern "C" fn read_packet_with_inner(
    ctx: *mut c_void,
    buffer: *mut u8,
    buffer_size: i32,
) -> i32 {
    let mut avio: Box<AVIOInner<Box<dyn AVIOReader>>> = Box::from_raw(ctx.cast());
    let items = avio.data.read(buffer_size);
    let len = min(items.len(), buffer_size as usize);
    buffer.copy_from(items.as_ptr().cast(), len);
    len as i32
}

unsafe extern "C" fn write_packet_with_inner(ctx: *mut c_void, buffer: *mut u8, size: i32) -> i32 {
    let mut avio: Box<AVIOInner<Box<dyn AVIOWriter>>> = Box::from_raw(ctx.cast());
    let mut byte_buffer = vec![0; size as usize];
    buffer.copy_to(byte_buffer.as_mut_ptr(), size as usize);
    let write_fn = avio.write.unwrap();
    write_fn(&mut avio.data, byte_buffer);
    std::mem::forget(avio);
    size
}

unsafe extern "C" fn seek_with_inner(ctx: *mut c_void, offset: i64, whence: i32) -> i64 {
    println!("ok");
    let mut avio: Box<AVIOInner<Box<dyn AVIOSeekable>>> = Box::from_raw(ctx.cast());
    avio.data.seek(offset, whence);

    0
}

impl<T> AVIO<T> {
    fn new(
        buffer_size: usize,
        obj: T,
        read: Option<fn(&mut T, i32) -> Vec<u8>>,
        write: Option<fn(&mut T, std::vec::Vec<u8>)>,
        seek: Option<fn(&mut T, i64, i32)>,
    ) -> AVIO<T> {
        let mut opaque = Box::new(AVIOInner {
            data: Box::new(obj),
            read,
            write,
            seek,
        });
        let ctx = unsafe {
            let buf: *mut u8 = av_malloc(buffer_size).cast();
            let opaque_c: *mut AVIOInner<T> = &mut *opaque;

            avio_alloc_context(
                buf,
                buffer_size as i32,
                if write.is_some() { 1 } else { 0 },
                opaque_c.cast(),
                if read.is_some() {
                    Some(read_packet_with_inner)
                } else {
                    None
                },
                if write.is_some() {
                    Some(write_packet_with_inner)
                } else {
                    None
                },
                if seek.is_some() {
                    Some(seek_with_inner)
                } else {
                    None
                },
            )
        };

        AVIO { ctx, inner: opaque }
    }

    pub fn inner(&self) -> &T {
        self.inner.data.as_ref()
    }

    pub fn inner_mut(&mut self) -> &mut T {
        self.inner.data.as_mut()
    }
}

impl<T> AsMutPtr<AVIOContext> for AVIO<T> {
    fn as_mut_ptr(&self) -> *mut AVIOContext {
        self.ctx
    }
}

impl<T> AsPtr<AVIOContext> for AVIO<T> {
    fn as_ptr(&self) -> *const AVIOContext {
        self.ctx
    }
}

impl<T: AVIOReader + AVIOWriter + AVIOSeekable> AVIO<T> {
    pub fn duplex_with_seek(item: T) -> AVIO<T> {
        AVIO::new(4096, item, Some(T::read), Some(T::write), Some(T::seek))
    }
}

impl<T: AVIOReader + AVIOWriter> AVIO<T> {
    pub fn duplex(item: T) -> AVIO<T> {
        AVIO::new(4096, item, Some(T::read), Some(T::write), None)
    }
}

impl<T: AVIOReader + AVIOSeekable> AVIO<T> {
    pub fn reader_with_seek(item: T) -> AVIO<T> {
        AVIO::new(4096, item, Some(T::read), None, Some(T::seek))
    }
}

impl<T: AVIOReader> AVIO<T> {
    pub fn reader(item: T) -> AVIO<T> {
        AVIO::new(4096, item, Some(T::read), None, None)
    }
}

impl<T: AVIOWriter + AVIOSeekable> AVIO<T> {
    pub fn writer_with_seek(item: T) -> AVIO<T> {
        AVIO::new(4096, item, None, Some(T::write), Some(T::seek))
    }
}

impl<T: AVIOWriter> AVIO<T> {
    pub fn writer(item: T) -> AVIO<T> {
        AVIO::new(4096, item, None, Some(T::write), None)
    }
}

pub trait AVIOSeekable {
    fn seek(&mut self, offset: i64, whence: i32);
}

pub trait AVIOWriter {
    fn write(&mut self, buffer: Vec<u8>);
}

pub trait AVIOReader {
    fn read(&mut self, length: i32) -> Vec<u8>;
}