nordisk/znippy

plugins/native/znippy-plugin-media/src/native.rs

//! The `ArchiveTypePlugin` implementation registered in the znippy CLI.
//!
//! One Arrow schema spans image/audio/video; a per-file [`MediaMeta`] fills only
//! the columns relevant to its kind and the rest serialize to null cells. The
//! writer maps `ExtensionValue::U32` onto `UInt32` columns and everything else
//! onto `Utf8`, so numeric facets (width / duration / sample rate) stay typed and
//! queryable.

use std::collections::HashMap;

use znippy_common::plugin::{
    default_command, ArchiveTypePlugin, Col, ExtensionRow, ExtensionValue, Claim, HandlerSpec,
};

use crate::{MEDIA_EXTENSIONS, MEDIA_TYPE_ID, MediaMeta, extract_media_metadata};

/// Native media handler: images (dims/color/EXIF), audio (duration/rate/tags),
/// video (dims/codecs/framerate). Stores the blob verbatim, extracts metadata.
pub struct NativeMediaPlugin;

impl NativeMediaPlugin {
    pub fn new() -> Self {
        NativeMediaPlugin
    }
}

impl Default for NativeMediaPlugin {
    fn default() -> Self {
        Self::new()
    }
}

/// The media handler, stated once.
pub static MEDIA_SPEC: HandlerSpec = HandlerSpec {
    name: "media",
    aliases: &["image", "audio", "video"],
    type_id: MEDIA_TYPE_ID,
    ecosystem: "Media files (images / audio / video)",
    extensions: MEDIA_EXTENSIONS,
    // Case-insensitive (`IMG.JPG` is a photo), as `is_media_path`.
    claim: Claim::ExtensionsAnyCase,
    description: "Image / audio / video metadata (dimensions, EXIF, duration, codecs, \
                  tags) — pure-Rust, blob stored verbatim",
    fields: &[
        ("media_kind", Col::Utf8),
        ("format", Col::Utf8),
        ("width", Col::U32),
        ("height", Col::U32),
        ("color", Col::Utf8),
        ("duration_ms", Col::U32),
        ("sample_rate", Col::U32),
        ("channels", Col::U32),
        ("bitrate", Col::U32),
        ("audio_codec", Col::Utf8),
        ("video_codec", Col::Utf8),
        ("framerate", Col::Utf8),
        ("title", Col::Utf8),
        ("artist", Col::Utf8),
        ("album", Col::Utf8),
        ("camera", Col::Utf8),
        ("datetime", Col::Utf8),
        ("orientation", Col::U32),
        ("gps", Col::Utf8),
    ],
    commands: &[("probe", "Print parsed media metadata for a file (kind/format/dims/duration/codecs)")],
};

impl ArchiveTypePlugin for NativeMediaPlugin {
    fn spec(&self) -> &HandlerSpec {
        &MEDIA_SPEC
    }

    fn run_command(&self, cmd: &str, args: &[String]) -> anyhow::Result<()> {
        if cmd != "probe" {
            return default_command(self, cmd, args);
        }
        let path = args.first().ok_or_else(|| anyhow::anyhow!("usage: media probe <file>"))?;
        let data = std::fs::read(path)?;
        match extract_media_metadata(path, &data) {
            Some(m) => {
                println!("{:#?}", m);
                Ok(())
            }
            None => anyhow::bail!("media: could not parse '{}'", path),
        }
    }

    fn extract_metadata(&self, path: &str, data: &[u8]) -> Option<ExtensionRow> {
        let m = extract_media_metadata(path, data)?;
        Some(media_meta_to_row(&m))
    }
}

/// Map a [`MediaMeta`] onto the Arrow row. Only present fields are inserted;
/// absent keys yield null cells (the writer appends null when a key is missing).
fn media_meta_to_row(m: &MediaMeta) -> ExtensionRow {
    let mut f: HashMap<String, ExtensionValue> = HashMap::new();

    // String columns
    if let Some(kind) = m.kind {
        f.insert("media_kind".to_string(), ExtensionValue::Str(kind.as_str().to_string()));
    }
    let mut put_str = |k: &str, v: &Option<String>| {
        if let Some(s) = v {
            f.insert(k.to_string(), ExtensionValue::Str(s.clone()));
        }
    };
    put_str("format", &m.format);
    put_str("color", &m.color);
    put_str("audio_codec", &m.audio_codec);
    put_str("video_codec", &m.video_codec);
    put_str("framerate", &m.framerate);
    put_str("title", &m.title);
    put_str("artist", &m.artist);
    put_str("album", &m.album);
    put_str("camera", &m.camera);
    put_str("datetime", &m.datetime);
    put_str("gps", &m.gps);

    // Numeric (UInt32) columns
    let mut put_u32 = |k: &str, v: Option<u32>| {
        if let Some(n) = v {
            f.insert(k.to_string(), ExtensionValue::U32(n));
        }
    };
    put_u32("width", m.width);
    put_u32("height", m.height);
    put_u32("duration_ms", m.duration_ms);
    put_u32("sample_rate", m.sample_rate);
    put_u32("channels", m.channels);
    put_u32("bitrate", m.bitrate);
    put_u32("orientation", m.orientation);

    ExtensionRow { fields: f }
}