markbase/markbase-fskit/src/fskit/frame_align.rs

use std::path::Path;

pub struct FrameAlignment {
    codec: CodecType,
    frame_size: usize,
    frame_boundary: usize,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CodecType {
    ProRes,
    H264,
    HEVC,
    DNxHD,
    DNxHR,
    XAVC,
    AVCIntra,
    Unknown,
}

impl FrameAlignment {
    pub fn detect(file_path: &Path) -> Option<Self> {
        let ext = file_path
            .extension()
            .and_then(|e| e.to_str())
            .map(|e| e.to_lowercase());

        let codec = match ext.as_deref() {
            Some("mov") => CodecType::ProRes,
            Some("mp4") => CodecType::H264,
            Some("m4v") => CodecType::H264,
            Some("mxf") => CodecType::DNxHD,
            Some("avi") => CodecType::Unknown,
            Some("mkv") => CodecType::Unknown,
            _ => CodecType::Unknown,
        };

        if codec == CodecType::Unknown {
            return None;
        }

        Some(FrameAlignment {
            codec,
            frame_size: 0,
            frame_boundary: 4096,
        })
    }

    pub fn is_aligned(&self, offset: usize, size: usize) -> bool {
        if self.frame_size == 0 {
            return offset % self.frame_boundary == 0 && size % self.frame_boundary == 0;
        }

        offset % self.frame_size == 0 && size % self.frame_size == 0
    }

    pub fn optimal_chunk_size(&self) -> usize {
        if self.frame_size == 0 {
            return self.frame_boundary;
        }

        self.frame_size
    }

    pub fn align_offset(&self, offset: usize) -> usize {
        if self.frame_size == 0 {
            let boundary = self.frame_boundary;
            (offset / boundary) * boundary
        } else {
            (offset / self.frame_size) * self.frame_size
        }
    }

    pub fn align_size(&self, size: usize) -> usize {
        if self.frame_size == 0 {
            let boundary = self.frame_boundary;
            ((size + boundary - 1) / boundary) * boundary
        } else {
            ((size + self.frame_size - 1) / self.frame_size) * self.frame_size
        }
    }

    pub fn get_codec(&self) -> CodecType {
        self.codec
    }
}

impl Default for FrameAlignment {
    fn default() -> Self {
        FrameAlignment {
            codec: CodecType::Unknown,
            frame_size: 0,
            frame_boundary: 4096,
        }
    }
}

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

    #[test]
    fn test_frame_alignment_prores() {
        let path = Path::new("/test/video.mov");
        let alignment = FrameAlignment::detect(path);
        assert!(alignment.is_some());
        let align = alignment.unwrap();
        assert_eq!(align.get_codec(), CodecType::ProRes);
    }

    #[test]
    fn test_frame_alignment_h264() {
        let path = Path::new("/test/video.mp4");
        let alignment = FrameAlignment::detect(path);
        assert!(alignment.is_some());
        let align = alignment.unwrap();
        assert_eq!(align.get_codec(), CodecType::H264);
    }

    #[test]
    fn test_frame_alignment_dnxhd() {
        let path = Path::new("/test/video.mxf");
        let alignment = FrameAlignment::detect(path);
        assert!(alignment.is_some());
        let align = alignment.unwrap();
        assert_eq!(align.get_codec(), CodecType::DNxHD);
    }

    #[test]
    fn test_frame_alignment_unknown() {
        let path = Path::new("/test/document.pdf");
        let alignment = FrameAlignment::detect(path);
        assert!(alignment.is_none());
    }

    #[test]
    fn test_alignment_boundary() {
        let align = FrameAlignment::default();
        assert!(align.is_aligned(0, 4096));
        assert!(align.is_aligned(4096, 4096));
        assert!(!align.is_aligned(100, 4096));
    }

    #[test]
    fn test_align_offset() {
        let align = FrameAlignment::default();
        assert_eq!(align.align_offset(0), 0);
        assert_eq!(align.align_offset(4096), 4096);
        assert_eq!(align.align_offset(5000), 4096);
        assert_eq!(align.align_offset(8192), 8192);
    }

    #[test]
    fn test_align_size() {
        let align = FrameAlignment::default();
        assert_eq!(align.align_size(0), 0);
        assert_eq!(align.align_size(4096), 4096);
        assert_eq!(align.align_size(5000), 8192);
        assert_eq!(align.align_size(100), 4096);
    }
}