keep/src/meta_plugin/infer_plugin.rs

use crate::common::PIPESIZE;
use crate::meta_plugin::{
    process_metadata_outputs, register_meta_plugin, BaseMetaPlugin, MetaPlugin, MetaPluginResponse,
    MetaPluginType,
};

#[derive(Debug, Default)]
pub struct InferMetaPlugin {
    buffer: Vec<u8>,
    max_buffer_size: usize,
    is_finalized: bool,
    base: BaseMetaPlugin,
}

impl InferMetaPlugin {
    pub fn new(
        options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
        outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
    ) -> InferMetaPlugin {
        let mut base = BaseMetaPlugin::new();

        if let Some(opts) = options {
            for (key, value) in opts {
                base.options.insert(key, value);
            }
        }

        let max_buffer_size = base
            .options
            .get("max_buffer_size")
            .and_then(|v| v.as_u64())
            .unwrap_or(PIPESIZE as u64) as usize;

        base.outputs.insert(
            "infer_mime_type".to_string(),
            serde_yaml::Value::String("infer_mime_type".to_string()),
        );

        if let Some(outs) = outputs {
            for (key, value) in outs {
                base.outputs.insert(key, value);
            }
        }

        InferMetaPlugin {
            buffer: Vec::new(),
            max_buffer_size,
            is_finalized: false,
            base,
        }
    }
}

impl MetaPlugin for InferMetaPlugin {
    fn meta_type(&self) -> MetaPluginType {
        MetaPluginType::Infer
    }

    fn is_finalized(&self) -> bool {
        self.is_finalized
    }

    fn set_finalized(&mut self, finalized: bool) {
        self.is_finalized = finalized;
    }

    fn set_save_meta(&mut self, save_meta: crate::meta_plugin::SaveMetaFn) {
        self.base.set_save_meta(save_meta);
    }

    fn save_meta(&self, name: &str, value: &str) {
        self.base.save_meta(name, value);
    }

    fn update(&mut self, data: &[u8]) -> MetaPluginResponse {
        if self.is_finalized {
            return MetaPluginResponse {
                metadata: Vec::new(),
                is_finalized: true,
            };
        }

        let remaining = self.max_buffer_size.saturating_sub(self.buffer.len());
        let to_add = &data[..data.len().min(remaining)];
        self.buffer.extend_from_slice(to_add);

        if self.buffer.len() >= self.max_buffer_size {
            let mime_type = infer::get(&self.buffer)
                .map(|kind| kind.mime_type().to_string())
                .unwrap_or_else(|| "application/octet-stream".to_string());

            self.is_finalized = true;

            let metadata = process_metadata_outputs(
                "infer_mime_type",
                serde_yaml::Value::String(mime_type),
                self.base.outputs(),
            )
            .map(|m| vec![m])
            .unwrap_or_default();

            return MetaPluginResponse {
                metadata,
                is_finalized: true,
            };
        }

        MetaPluginResponse {
            metadata: Vec::new(),
            is_finalized: false,
        }
    }

    fn finalize(&mut self) -> MetaPluginResponse {
        if self.is_finalized {
            return MetaPluginResponse {
                metadata: Vec::new(),
                is_finalized: true,
            };
        }

        let mime_type = infer::get(&self.buffer)
            .map(|kind| kind.mime_type().to_string())
            .unwrap_or_else(|| "application/octet-stream".to_string());

        self.is_finalized = true;

        let metadata = process_metadata_outputs(
            "infer_mime_type",
            serde_yaml::Value::String(mime_type),
            self.base.outputs(),
        )
        .map(|m| vec![m])
        .unwrap_or_default();

        MetaPluginResponse {
            metadata,
            is_finalized: true,
        }
    }

    fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
        self.base.outputs()
    }

    fn outputs_mut(
        &mut self,
    ) -> anyhow::Result<&mut std::collections::HashMap<String, serde_yaml::Value>> {
        Ok(self.base.outputs_mut())
    }

    fn default_outputs(&self) -> Vec<String> {
        vec!["infer_mime_type".to_string()]
    }

    fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
        self.base.options()
    }

    fn options_mut(
        &mut self,
    ) -> anyhow::Result<&mut std::collections::HashMap<String, serde_yaml::Value>> {
        Ok(self.base.options_mut())
    }

    fn parallel_safe(&self) -> bool {
        true
    }
}

#[ctor::ctor]
fn register_infer_plugin() {
    register_meta_plugin(MetaPluginType::Infer, |options, outputs| {
        Box::new(InferMetaPlugin::new(options, outputs))
    })
    .expect("Failed to register InferMetaPlugin");
}