poimen-memory/crates/mem-core/src/optimizer/builtin.rs

//! Built-in optimizer and format implementations
//! Bridges existing ContextOptimizer to pluggable system

use async_trait::async_trait;
use super::plugin::{OptimizerPlugin, FormatHandler, OptimizationResult, PluginMetrics};
use super::ContextOptimizer;
use std::sync::{Arc, Mutex};

/// Built-in context optimizer plugin
pub struct BuiltinOptimizer {
    optimizer: Arc<ContextOptimizer>,
    metrics: Arc<Mutex<PluginMetrics>>,
}

impl BuiltinOptimizer {
    pub fn new(optimizer: Arc<ContextOptimizer>) -> Self {
        Self {
            optimizer,
            metrics: Arc::new(Mutex::new(PluginMetrics::default())),
        }
    }
}

#[async_trait]
impl OptimizerPlugin for BuiltinOptimizer {
    fn name(&self) -> &str {
        "builtin-optimizer"
    }

    fn supported_types(&self) -> Vec<&str> {
        vec![
            "text/plain",
            "text/x-log",
            "application/json",
            "text/x-diff",
            "application/x-yaml",
        ]
    }

    async fn optimize(&self, content: &str) -> Result<OptimizationResult, String> {
        let start = std::time::Instant::now();
        let original_len = content.len();

        match self.optimizer.optimize(content) {
            Ok(chunk) => {
                let optimized_len = chunk.compressed.len();
                let ratio = optimized_len as f32 / original_len as f32;

                // Update metrics
                {
                    let mut m = self.metrics.lock().unwrap();
                    m.total_optimizations += 1;
                    m.total_bytes_input += original_len as u64;
                    m.total_bytes_output += optimized_len as u64;
                    m.avg_latency_ms = start.elapsed().as_secs_f32() * 1000.0;
                }

                Ok(OptimizationResult {
                    original: content.to_string(),
                    optimized: chunk.compressed,
                    ratio,
                    plugin: self.name().to_string(),
                    metadata: std::collections::HashMap::new(),
                })
            }
            Err(e) => {
                let mut m = self.metrics.lock().unwrap();
                m.errors += 1;
                Err(format!("Optimization failed: {}", e))
            }
        }
    }

    fn metrics(&self) -> PluginMetrics {
        self.metrics.lock().unwrap().clone()
    }
}

// ============================================================================
// FORMAT HANDLERS
// ============================================================================

/// JSON format handler
pub struct JsonFormatter;

#[async_trait]
impl FormatHandler for JsonFormatter {
    fn name(&self) -> &str {
        "json"
    }

    async fn format(&self, result: &OptimizationResult) -> Result<Vec<u8>, String> {
        serde_json::to_vec(result).map_err(|e| format!("JSON serialization failed: {}", e))
    }

    async fn parse(&self, data: &[u8]) -> Result<OptimizationResult, String> {
        serde_json::from_slice(data)
            .map_err(|e| format!("JSON deserialization failed: {}", e))
    }
}

/// JSONL (newline-delimited JSON) formatter
pub struct JsonlFormatter;

#[async_trait]
impl FormatHandler for JsonlFormatter {
    fn name(&self) -> &str {
        "jsonl"
    }

    async fn format(&self, result: &OptimizationResult) -> Result<Vec<u8>, String> {
        let mut output = serde_json::to_string(result)
            .map_err(|e| format!("JSONL serialization failed: {}", e))?;
        output.push('\n');
        Ok(output.into_bytes())
    }

    async fn parse(&self, data: &[u8]) -> Result<OptimizationResult, String> {
        let line = String::from_utf8(data.to_vec())
            .map_err(|e| format!("UTF-8 decode failed: {}", e))?;
        serde_json::from_str(line.trim())
            .map_err(|e| format!("JSONL deserialization failed: {}", e))
    }
}

/// Raw text format (just optimized content)
pub struct RawFormatter;

#[async_trait]
impl FormatHandler for RawFormatter {
    fn name(&self) -> &str {
        "raw"
    }

    async fn format(&self, result: &OptimizationResult) -> Result<Vec<u8>, String> {
        Ok(result.optimized.as_bytes().to_vec())
    }

    async fn parse(&self, _data: &[u8]) -> Result<OptimizationResult, String> {
        Err("Raw format does not support deserialization".to_string())
    }
}

/// CSV format (for metrics)
pub struct CsvFormatter;

#[async_trait]
impl FormatHandler for CsvFormatter {
    fn name(&self) -> &str {
        "csv"
    }

    async fn format(&self, result: &OptimizationResult) -> Result<Vec<u8>, String> {
        let output = format!(
            "{},{},{},{}\n",
            escape_csv(&result.plugin),
            result.original.len(),
            result.optimized.len(),
            format!("{:.2}", result.ratio)
        );
        Ok(output.into_bytes())
    }

    async fn parse(&self, _data: &[u8]) -> Result<OptimizationResult, String> {
        Err("CSV format does not support deserialization".to_string())
    }
}

/// YAML format
pub struct YamlFormatter;

#[async_trait]
impl FormatHandler for YamlFormatter {
    fn name(&self) -> &str {
        "yaml"
    }

    async fn format(&self, result: &OptimizationResult) -> Result<Vec<u8>, String> {
        let yaml = format!(
            "plugin: {}\nratio: {:.2}\noriginal_bytes: {}\noptimized_bytes: {}\n",
            result.plugin,
            result.ratio,
            result.original.len(),
            result.optimized.len()
        );
        Ok(yaml.into_bytes())
    }

    async fn parse(&self, _data: &[u8]) -> Result<OptimizationResult, String> {
        Err("YAML format does not support deserialization".to_string())
    }
}

// ============================================================================
// HELPERS
// ============================================================================

fn escape_csv(s: &str) -> String {
    if s.contains(',') || s.contains('"') || s.contains('\n') {
        format!("\"{}\"", s.replace('"', "\"\""))
    } else {
        s.to_string()
    }
}

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

    #[tokio::test]
    async fn test_builtin_optimizer_logs() {
        let optimizer = Arc::new(ContextOptimizer::new().unwrap());
        let plugin = BuiltinOptimizer::new(optimizer);

        let content = "ERROR: failed\nDEBUG: info\nERROR: error";
        let result = plugin.optimize(content).await.unwrap();

        assert_eq!(result.plugin, "builtin-optimizer");
        assert!(result.ratio <= 1.0);
    }

    #[tokio::test]
    async fn test_json_formatter() {
        let formatter = JsonFormatter;
        let result = OptimizationResult {
            original: "test".to_string(),
            optimized: "tst".to_string(),
            ratio: 0.75,
            plugin: "test".to_string(),
            metadata: Default::default(),
        };

        let formatted = formatter.format(&result).await.unwrap();
        let parsed = formatter.parse(&formatted).await.unwrap();

        assert_eq!(parsed.original, "test");
    }

    #[tokio::test]
    async fn test_jsonl_formatter() {
        let formatter = JsonlFormatter;
        let result = OptimizationResult {
            original: "test".to_string(),
            optimized: "tst".to_string(),
            ratio: 0.75,
            plugin: "test".to_string(),
            metadata: Default::default(),
        };

        let formatted = formatter.format(&result).await.unwrap();
        assert!(String::from_utf8(formatted).unwrap().ends_with('\n'));
    }

    #[tokio::test]
    async fn test_raw_formatter() {
        let formatter = RawFormatter;
        let result = OptimizationResult {
            original: "original content".to_string(),
            optimized: "optimized".to_string(),
            ratio: 0.5,
            plugin: "test".to_string(),
            metadata: Default::default(),
        };

        let formatted = formatter.format(&result).await.unwrap();
        assert_eq!(formatted, b"optimized");
    }

    #[tokio::test]
    async fn test_csv_formatter() {
        let formatter = CsvFormatter;
        let result = OptimizationResult {
            original: "test".to_string(),
            optimized: "tst".to_string(),
            ratio: 0.75,
            plugin: "my-plugin".to_string(),
            metadata: Default::default(),
        };

        let formatted = formatter.format(&result).await.unwrap();
        let csv = String::from_utf8(formatted).unwrap();
        assert!(csv.contains("my-plugin"));
    }

    #[tokio::test]
    async fn test_yaml_formatter() {
        let formatter = YamlFormatter;
        let result = OptimizationResult {
            original: "original".to_string(),
            optimized: "opt".to_string(),
            ratio: 0.33,
            plugin: "test".to_string(),
            metadata: Default::default(),
        };

        let formatted = formatter.format(&result).await.unwrap();
        let yaml = String::from_utf8(formatted).unwrap();
        assert!(yaml.contains("ratio:"));
    }
}