poimen-memory/crates/mem-cli/src/result_compressor.rs

/// Result Compressor: Optimize response size without losing essential information
///
/// Strategies:
/// - Truncate long texts to summary
/// - Extract key sentences
/// - Remove redundant metadata
/// - Compress to multiple formats (JSON, msgpack, CBOR)
/// - Progressive disclosure (compact by default, expand on demand)

use anyhow::Result;
use serde::{Deserialize, Serialize};

/// Compression strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionStrategy {
    /// No compression
    None,
    /// Extract first 100 chars + key sentences
    Summarize,
    /// Remove secondary fields
    Minimal,
    /// Aggressive: ids + scores only
    Ultra,
}

/// Compressed chunk result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressedResult {
    pub id: String,
    pub score: f32,
    pub text: Option<String>,      // Optional if compression=Ultra
    pub category: Option<String>,  // Optional
    pub cache_slot: Option<u32>,   // Optional
}

impl CompressedResult {
    pub fn new(id: &str, score: f32) -> Self {
        Self {
            id: id.to_string(),
            score,
            text: None,
            category: None,
            cache_slot: None,
        }
    }

    pub fn with_text(mut self, text: &str) -> Self {
        self.text = Some(text.to_string());
        self
    }

    pub fn with_category(mut self, category: &str) -> Self {
        self.category = Some(category.to_string());
        self
    }

    pub fn with_cache_slot(mut self, slot: u32) -> Self {
        self.cache_slot = Some(slot);
        self
    }
}

/// Text summarizer
pub struct TextSummarizer {
    max_length: usize,
    sentence_limit: usize,
}

impl TextSummarizer {
    pub fn new(max_length: usize, sentence_limit: usize) -> Self {
        Self {
            max_length,
            sentence_limit,
        }
    }

    /// Extract first N sentences
    pub fn extract_sentences(&self, text: &str, limit: usize) -> String {
        let sentences: Vec<&str> = text
            .split('.')
            .filter(|s| !s.trim().is_empty())
            .take(limit)
            .collect();

        sentences
            .join(". ")
            .trim_end_matches(' ')
            .to_string()
            + if sentences.len() >= limit && !text.ends_with('.') {
                "..."
            } else {
                ""
            }
    }

    /// Truncate to max length with ellipsis
    pub fn truncate(&self, text: &str) -> String {
        if text.len() > self.max_length {
            let truncated = &text[..self.max_length];
            // Find last space to avoid cutting words
            if let Some(pos) = truncated.rfind(' ') {
                format!("{}...", &text[..pos])
            } else {
                format!("{}...", truncated)
            }
        } else {
            text.to_string()
        }
    }

    /// Summarize by extracting key sentences and truncating
    pub fn summarize(&self, text: &str) -> String {
        let key_sentences = self.extract_sentences(text, self.sentence_limit);
        self.truncate(&key_sentences)
    }
}

/// Result compressor
pub struct ResultCompressor {
    summarizer: TextSummarizer,
}

impl ResultCompressor {
    pub fn new(max_text_length: usize, sentence_limit: usize) -> Self {
        Self {
            summarizer: TextSummarizer::new(max_text_length, sentence_limit),
        }
    }

    /// Compress single result
    pub fn compress(
        &self,
        id: &str,
        text: &str,
        score: f32,
        strategy: CompressionStrategy,
    ) -> CompressedResult {
        let mut result = CompressedResult::new(id, score);

        match strategy {
            CompressionStrategy::None => {
                result.text = Some(text.to_string());
            }
            CompressionStrategy::Summarize => {
                result.text = Some(self.summarizer.summarize(text));
            }
            CompressionStrategy::Minimal => {
                result.text = Some(self.summarizer.truncate(text));
            }
            CompressionStrategy::Ultra => {
                result.text = None; // Drop text entirely
            }
        }

        result
    }

    /// Compress multiple results
    pub fn compress_batch(
        &self,
        results: Vec<(String, String, f32)>, // (id, text, score)
        strategy: CompressionStrategy,
    ) -> Vec<CompressedResult> {
        results
            .into_iter()
            .map(|(id, text, score)| self.compress(&id, &text, score, strategy))
            .collect()
    }

    /// Estimate size of compressed results
    pub fn estimate_size(
        &self,
        results: &[CompressedResult],
        include_text: bool,
    ) -> usize {
        let mut size = 0;

        for result in results {
            size += result.id.len() + 4; // id + score (f32)

            if include_text {
                if let Some(text) = &result.text {
                    size += text.len();
                }
            }

            if let Some(category) = &result.category {
                size += category.len();
            }
        }

        size
    }
}

/// Budget-aware compressor (automatically choose compression level)
pub struct BudgetCompressor {
    max_budget_bytes: usize,
    compressor: ResultCompressor,
}

impl BudgetCompressor {
    pub fn new(max_budget_bytes: usize) -> Self {
        Self {
            max_budget_bytes,
            compressor: ResultCompressor::new(500, 3),
        }
    }

    /// Automatically select compression level based on budget
    pub fn select_strategy(&self, estimated_size: usize) -> CompressionStrategy {
        let ratio = estimated_size as f32 / self.max_budget_bytes as f32;

        if ratio < 0.5 {
            CompressionStrategy::None
        } else if ratio < 0.75 {
            CompressionStrategy::Summarize
        } else if ratio < 1.0 {
            CompressionStrategy::Minimal
        } else {
            CompressionStrategy::Ultra
        }
    }

    /// Compress results intelligently to stay within budget
    pub fn compress_to_budget(
        &self,
        results: Vec<(String, String, f32)>,
    ) -> (Vec<CompressedResult>, CompressionStrategy) {
        let estimated = results
            .iter()
            .map(|(_, text, _)| text.len())
            .sum::<usize>();

        let strategy = self.select_strategy(estimated);
        let compressed = self.compressor.compress_batch(results, strategy);

        (compressed, strategy)
    }
}

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

    #[test]
    fn test_compressed_result_builder() {
        let result = CompressedResult::new("doc1", 0.9)
            .with_text("Some text")
            .with_category("solution")
            .with_cache_slot(5);

        assert_eq!(result.id, "doc1");
        assert_eq!(result.score, 0.9);
        assert_eq!(result.text, Some("Some text".to_string()));
        assert_eq!(result.cache_slot, Some(5));
    }

    #[test]
    fn test_text_summarizer_truncate() {
        let summarizer = TextSummarizer::new(20, 3);
        let text = "This is a long text that needs to be truncated";
        let truncated = summarizer.truncate(text);

        assert!(truncated.len() <= 23); // 20 + "..."
        assert!(truncated.ends_with("..."));
    }

    #[test]
    fn test_text_summarizer_extract_sentences() {
        let summarizer = TextSummarizer::new(500, 2);
        let text = "First sentence. Second sentence. Third sentence.";
        let extracted = summarizer.extract_sentences(text, 2);

        assert!(extracted.contains("First sentence"));
        assert!(extracted.contains("Second sentence"));
        assert!(!extracted.contains("Third sentence"));
    }

    #[test]
    fn test_text_summarizer_summarize() {
        let summarizer = TextSummarizer::new(50, 2);
        let text =
            "First sentence. Second sentence. Third sentence with lots of details that continue.";
        let summarized = summarizer.summarize(text);

        assert!(summarized.len() <= 53); // 50 + "..."
        assert!(summarized.contains("First"));
    }

    #[test]
    fn test_result_compressor_none() {
        let compressor = ResultCompressor::new(500, 3);
        let result = compressor.compress("doc1", "test text", 0.9, CompressionStrategy::None);

        assert_eq!(result.text, Some("test text".to_string()));
    }

    #[test]
    fn test_result_compressor_summarize() {
        let compressor = ResultCompressor::new(50, 1);
        let text = "First sentence. Second sentence. Third sentence.";
        let result = compressor.compress("doc1", text, 0.9, CompressionStrategy::Summarize);

        assert!(result.text.is_some());
        if let Some(compressed) = result.text {
            assert!(compressed.len() <= 100);
        }
    }

    #[test]
    fn test_result_compressor_ultra() {
        let compressor = ResultCompressor::new(500, 3);
        let result = compressor.compress("doc1", "test text", 0.9, CompressionStrategy::Ultra);

        assert_eq!(result.text, None);
        assert_eq!(result.id, "doc1");
        assert_eq!(result.score, 0.9);
    }

    #[test]
    fn test_result_compressor_batch() {
        let compressor = ResultCompressor::new(100, 2);
        let results = vec![
            ("doc1".to_string(), "short".to_string(), 0.9),
            ("doc2".to_string(), "another text".to_string(), 0.8),
        ];

        let compressed = compressor.compress_batch(results, CompressionStrategy::Minimal);
        assert_eq!(compressed.len(), 2);
        assert!(compressed[0].text.is_some());
    }

    #[test]
    fn test_estimate_size() {
        let compressor = ResultCompressor::new(500, 3);
        let results = vec![
            CompressedResult::new("doc1", 0.9).with_text("some text"),
            CompressedResult::new("doc2", 0.8).with_text("more text"),
        ];

        let size = compressor.estimate_size(&results, true);
        assert!(size > 0);
    }

    #[test]
    fn test_budget_compressor_select_none() {
        let compressor = BudgetCompressor::new(1000);
        let strategy = compressor.select_strategy(300);
        assert_eq!(strategy, CompressionStrategy::None);
    }

    #[test]
    fn test_budget_compressor_select_summarize() {
        let compressor = BudgetCompressor::new(1000);
        let strategy = compressor.select_strategy(600);
        assert_eq!(strategy, CompressionStrategy::Summarize);
    }

    #[test]
    fn test_budget_compressor_select_ultra() {
        let compressor = BudgetCompressor::new(1000);
        let strategy = compressor.select_strategy(1200);
        assert_eq!(strategy, CompressionStrategy::Ultra);
    }

    #[test]
    fn test_budget_compressor_compress_to_budget() {
        let compressor = BudgetCompressor::new(1000);
        let results = vec![
            ("doc1".to_string(), "short text".to_string(), 0.9),
            ("doc2".to_string(), "more content".to_string(), 0.8),
        ];

        let (compressed, strategy) = compressor.compress_to_budget(results);
        assert!(compressed.len() > 0);
        assert_ne!(strategy, CompressionStrategy::Ultra); // Should not be ultra for small input
    }
}