znippy-plugin-media 0.9.0

//! Media metadata plugin for znippy — images, audio and video.
//!
//! Media files (JPEG, PNG, MP3, FLAC, MP4, MKV, …) are *already compressed*, so
//! znippy stores their blob **verbatim** (never re-codes the codec). What this
//! plugin adds is a rich, queryable metadata sub-index: for every media file we
//! parse the container/header and contribute Arrow columns (dimensions, color,
//! EXIF, duration, codecs, sample rate, tags, …) that land in the lookup
//! sub-index alongside the rest — so DuckDB/Polars can `SELECT … WHERE
//! width > 1920` over a `.znippy` archive.
//!
//! Everything here is **pure Rust** (airgap-friendly, no C/ffmpeg):
//! - images: [`image`] (dimensions + color type) + [`exif`](kamadak-exif) (EXIF)
//! - audio:  [`lofty`] (duration / sample rate / channels / bitrate + tags)
//! - video:  [`mp4`] (MP4 / MOV) and [`matroska`] (MKV / WebM)
//!
//! On any parse failure the extractor returns `None` for that file — it never
//! panics and never blocks compression.

mod native;
pub use native::NativeMediaPlugin;

use std::io::Cursor;

/// The on-disk DenseUnion discriminant this plugin writes. Continues the
/// `type_id` registry past the last skeleton (24).
pub const MEDIA_TYPE_ID: i8 = 25;

/// All file extensions the media handler claims (lower-case, with dot).
pub const MEDIA_EXTENSIONS: &[&str] = &[
    // images
    ".jpg", ".jpeg", ".png", ".webp", ".gif", ".tif", ".tiff", ".bmp",
    // audio
    ".mp3", ".flac", ".ogg", ".oga", ".opus", ".wav", ".m4a", ".aac",
    // video
    ".mp4", ".m4v", ".mov", ".mkv", ".webm",
];

/// What sort of media a row describes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MediaKind {
    Image,
    Audio,
    Video,
}

impl MediaKind {
    pub fn as_str(self) -> &'static str {
        match self {
            MediaKind::Image => "image",
            MediaKind::Audio => "audio",
            MediaKind::Video => "video",
        }
    }
}

/// Unified, mostly-optional metadata record for one media file. A single Arrow
/// schema covers image/audio/video; fields not relevant to a given kind stay
/// `None` and serialize to null cells.
#[derive(Debug, Default, Clone, PartialEq)]
pub struct MediaMeta {
    pub kind: Option<MediaKind>,
    /// Container / codec short name (e.g. `png`, `mp3`, `mp4`, `webm`).
    pub format: Option<String>,
    pub width: Option<u32>,
    pub height: Option<u32>,
    /// Image color type / bit-depth (e.g. `rgb8`, `rgba8`, `l16`).
    pub color: Option<String>,
    pub duration_ms: Option<u32>,
    pub sample_rate: Option<u32>,
    pub channels: Option<u32>,
    /// Audio/overall bitrate in kbit/s (lofty) or bit/s (mp4) — see extractor.
    pub bitrate: Option<u32>,
    pub audio_codec: Option<String>,
    pub video_codec: Option<String>,
    /// Frames per second, kept as text to preserve fractional rates (`29.970`).
    pub framerate: Option<String>,
    pub title: Option<String>,
    pub artist: Option<String>,
    pub album: Option<String>,
    /// EXIF camera (make + model).
    pub camera: Option<String>,
    /// EXIF capture timestamp (DateTimeOriginal, else DateTime).
    pub datetime: Option<String>,
    /// EXIF orientation tag (1..=8).
    pub orientation: Option<u32>,
    /// EXIF GPS position, `"<lat ref> / <lon ref>"` display form when present.
    pub gps: Option<String>,
}

/// Lower-case file extension (including the dot), e.g. `".png"`.
fn ext_of(path: &str) -> Option<String> {
    let fname = path.rsplit(['/', '\\']).next().unwrap_or(path);
    let dot = fname.rfind('.')?;
    Some(fname[dot..].to_ascii_lowercase())
}

/// True if the media handler claims this path (by extension only — cheap, no read).
pub fn is_media_path(path: &str) -> bool {
    match ext_of(path) {
        Some(e) => MEDIA_EXTENSIONS.contains(&e.as_str()),
        None => false,
    }
}

/// Parse media metadata from an in-memory blob. Dispatches on the file
/// extension. Returns `None` if the extension isn't media or parsing fails.
pub fn extract_media_metadata(path: &str, data: &[u8]) -> Option<MediaMeta> {
    let ext = ext_of(path)?;
    match ext.as_str() {
        ".jpg" | ".jpeg" | ".png" | ".webp" | ".gif" | ".tif" | ".tiff" | ".bmp" => {
            extract_image(data)
        }
        ".mp3" | ".flac" | ".ogg" | ".oga" | ".opus" | ".wav" | ".m4a" | ".aac" => {
            extract_audio(data)
        }
        ".mp4" | ".m4v" | ".mov" => extract_mp4(data),
        ".mkv" | ".webm" => extract_matroska(data, &ext),
        _ => None,
    }
}

// ───────────────────────────── images ──────────────────────────────────────

fn extract_image(data: &[u8]) -> Option<MediaMeta> {
    use image::ImageDecoder;

    let reader = image::ImageReader::new(Cursor::new(data)).with_guessed_format().ok()?;
    let format = reader.format();
    let decoder = reader.into_decoder().ok()?;
    let (w, h) = decoder.dimensions();
    let color = decoder.color_type();
    drop(decoder);

    let mut m = MediaMeta {
        kind: Some(MediaKind::Image),
        format: format.map(|f| format!("{f:?}").to_ascii_lowercase()),
        width: Some(w),
        height: Some(h),
        color: Some(format!("{color:?}").to_ascii_lowercase()),
        ..Default::default()
    };

    read_exif_into(data, &mut m);
    Some(m)
}

/// Best-effort EXIF overlay. Populates camera / datetime / orientation / gps
/// when the image carries an EXIF segment; silently no-ops otherwise.
fn read_exif_into(data: &[u8], m: &mut MediaMeta) {
    use exif::{In, Tag};

    let exif = match exif::Reader::new().read_from_container(&mut Cursor::new(data)) {
        Ok(e) => e,
        Err(_) => return,
    };

    let text = |tag: Tag| -> Option<String> {
        exif.get_field(tag, In::PRIMARY).map(|f| f.display_value().to_string().trim().to_string())
    };
    let nonempty = |s: Option<String>| s.filter(|x| !x.is_empty());

    let make = nonempty(text(Tag::Make));
    let model = nonempty(text(Tag::Model));
    m.camera = match (make, model) {
        (Some(a), Some(b)) if b.starts_with(&a) => Some(b),
        (Some(a), Some(b)) => Some(format!("{a} {b}")),
        (Some(a), None) => Some(a),
        (None, Some(b)) => Some(b),
        (None, None) => None,
    };

    m.datetime = nonempty(text(Tag::DateTimeOriginal)).or_else(|| nonempty(text(Tag::DateTime)));

    m.orientation = exif
        .get_field(Tag::Orientation, In::PRIMARY)
        .and_then(|f| f.value.get_uint(0));

    let lat = nonempty(text(Tag::GPSLatitude));
    let lon = nonempty(text(Tag::GPSLongitude));
    if lat.is_some() || lon.is_some() {
        let latr = nonempty(text(Tag::GPSLatitudeRef)).unwrap_or_default();
        let lonr = nonempty(text(Tag::GPSLongitudeRef)).unwrap_or_default();
        m.gps = Some(format!(
            "{} {} / {} {}",
            lat.unwrap_or_default(),
            latr,
            lon.unwrap_or_default(),
            lonr
        ));
    }
}

// ───────────────────────────── audio ───────────────────────────────────────

fn extract_audio(data: &[u8]) -> Option<MediaMeta> {
    use lofty::prelude::*;
    use lofty::probe::Probe;

    let tagged =
        Probe::new(Cursor::new(data)).guess_file_type().ok()?.read().ok()?;
    let props = tagged.properties();

    let codec = format!("{:?}", tagged.file_type()).to_ascii_lowercase();
    let mut m = MediaMeta {
        kind: Some(MediaKind::Audio),
        format: Some(codec.clone()),
        audio_codec: Some(codec),
        duration_ms: u32::try_from(props.duration().as_millis()).ok(),
        sample_rate: props.sample_rate(),
        channels: props.channels().map(u32::from),
        bitrate: props.audio_bitrate().or_else(|| props.overall_bitrate()),
        ..Default::default()
    };

    if let Some(tag) = tagged.primary_tag().or_else(|| tagged.first_tag()) {
        m.title = tag.title().map(|c| c.into_owned());
        m.artist = tag.artist().map(|c| c.into_owned());
        m.album = tag.album().map(|c| c.into_owned());
    }

    Some(m)
}

// ───────────────────────────── video: MP4 / MOV ────────────────────────────

fn extract_mp4(data: &[u8]) -> Option<MediaMeta> {
    use mp4::TrackType;

    let size = data.len() as u64;
    let reader = mp4::Mp4Reader::read_header(Cursor::new(data), size).ok()?;

    let mut m = MediaMeta {
        kind: Some(MediaKind::Video),
        format: Some("mp4".to_string()),
        duration_ms: u32::try_from(reader.duration().as_millis()).ok(),
        ..Default::default()
    };

    for track in reader.tracks().values() {
        match track.track_type() {
            Ok(TrackType::Video) => {
                m.width = Some(u32::from(track.width()));
                m.height = Some(u32::from(track.height()));
                m.video_codec =
                    track.media_type().ok().map(|t| t.to_string().to_ascii_lowercase());
                let fr = track.frame_rate();
                if fr.is_finite() && fr > 0.0 {
                    m.framerate = Some(format!("{fr:.3}"));
                }
                let br = track.bitrate();
                if br > 0 {
                    m.bitrate = Some(br);
                }
            }
            Ok(TrackType::Audio) => {
                m.audio_codec =
                    track.media_type().ok().map(|t| t.to_string().to_ascii_lowercase());
                if let Ok(freq) = track.sample_freq_index() {
                    m.sample_rate = Some(freq.freq());
                }
            }
            _ => {}
        }
    }

    Some(m)
}

// ───────────────────────────── video: MKV / WebM ───────────────────────────

fn extract_matroska(data: &[u8], ext: &str) -> Option<MediaMeta> {
    use matroska::{Matroska, Settings};

    let mkv = Matroska::open(Cursor::new(data)).ok()?;

    let format = if ext == ".webm" { "webm" } else { "matroska" };
    let mut m = MediaMeta {
        kind: Some(MediaKind::Video),
        format: Some(format.to_string()),
        duration_ms: mkv.info.duration.and_then(|d| u32::try_from(d.as_millis()).ok()),
        title: mkv.info.title.clone(),
        ..Default::default()
    };

    if let Some(v) = mkv.video_tracks().next() {
        m.video_codec = Some(v.codec_id.to_ascii_lowercase());
        if let Settings::Video(vid) = &v.settings {
            m.width = u32::try_from(vid.pixel_width).ok();
            m.height = u32::try_from(vid.pixel_height).ok();
        }
        if let Some(dd) = v.default_duration {
            let nanos = dd.as_nanos();
            if nanos > 0 {
                let fps = 1.0e9 / nanos as f64;
                m.framerate = Some(format!("{fps:.3}"));
            }
        }
    }

    if let Some(a) = mkv.audio_tracks().next() {
        m.audio_codec = Some(a.codec_id.to_ascii_lowercase());
        if let Settings::Audio(au) = &a.settings {
            m.sample_rate = Some(au.sample_rate as u32);
            m.channels = u32::try_from(au.channels).ok();
        }
    }

    Some(m)
}

#[cfg(test)]
mod tests;