poimen-memory/crates/mem-cli/src/query/summarizer.rs

//! Result Summarization (Phase 5.4)
//!
//! Abstracting results, extracting key facts, optimizing coherence,
//! and generating length-controlled summaries.

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use tracing::debug;

/// Summarization strategy
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum SummarizationStrategy {
    /// Extractive: Select top-N sentences
    Extractive,
    /// Abstractive: Generate new concise text
    Abstractive,
    /// Hybrid: Extract + rewrite for coherence
    Hybrid,
}

/// Key fact extracted from results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyFact {
    /// Fact content
    pub fact: String,
    /// Importance score (0-1)
    pub importance: f32,
    /// Source entity ID
    pub source_id: String,
    /// Fact type (entity, relation, property)
    pub fact_type: String,
}

/// Summary with metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Summary {
    /// Original content length
    pub original_length: usize,
    /// Summary text
    pub text: String,
    /// Summary length
    pub summary_length: usize,
    /// Compression ratio
    pub compression_ratio: f32,
    /// Key facts in summary
    pub key_facts: Vec<KeyFact>,
    /// Coherence score (0-1)
    pub coherence: f32,
    /// Strategy used
    pub strategy: SummarizationStrategy,
}

/// Coherence metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CoherenceMetrics {
    /// Entity repetition score
    pub entity_coherence: f32,
    /// Sentence flow score
    pub flow_coherence: f32,
    /// Semantic similarity score
    pub semantic_coherence: f32,
}

/// Summarizer engine
pub struct Summarizer;

/// Entity detection helper (DRY)
fn is_capitalized_entity(word: &str, min_len: usize) -> bool {
    word.chars().next().map_or(false, |c| c.is_uppercase()) && word.len() >= min_len
}

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

    /// Generate summary from results
    pub fn summarize(
        &self,
        content: &str,
        max_length: usize,
        strategy: SummarizationStrategy,
    ) -> Result<Summary, String> {
        if content.is_empty() {
            return Err("Content cannot be empty".to_string());
        }

        if max_length < 50 {
            return Err("Summary length must be at least 50 characters".to_string());
        }

        let original_length = content.len();
        debug!("Summarizing {} chars to ~{} chars", original_length, max_length);

        let summary_text = match strategy {
            SummarizationStrategy::Extractive => {
                self.extractive_summarize(content, max_length)?
            }
            SummarizationStrategy::Abstractive => {
                self.abstractive_summarize(content, max_length)?
            }
            SummarizationStrategy::Hybrid => {
                self.hybrid_summarize(content, max_length)?
            }
        };

        let summary_length = summary_text.len();
        let compression_ratio = summary_length as f32 / original_length as f32;

        let key_facts = self.extract_key_facts(content, &summary_text);
        let coherence = self.compute_coherence(&summary_text);

        Ok(Summary {
            original_length,
            text: summary_text,
            summary_length,
            compression_ratio,
            key_facts,
            coherence,
            strategy,
        })
    }

    /// Extractive summarization: select top sentences
    fn extractive_summarize(&self, content: &str, max_length: usize) -> Result<String, String> {
        let sentences = self.split_sentences(content);

        if sentences.is_empty() {
            return Ok(content.to_string());
        }

        // Score sentences
        let mut scored: Vec<(usize, &str, f32)> = sentences
            .iter()
            .enumerate()
            .map(|(idx, sent)| (idx, *sent, self.score_sentence(sent, content)))
            .collect();

        // Sort by score descending
        scored.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));

        // Select top sentences by score
        let mut selected = Vec::new();
        let mut current_length = 0;

        for (idx, sent, _score) in scored {
            if current_length + sent.len() + 1 > max_length && !selected.is_empty() {
                break;
            }
            selected.push((idx, sent));
            current_length += sent.len() + 1;
        }

        // Preserve original order
        selected.sort_by_key(|a| a.0);
        let result = selected.into_iter().map(|a| a.1).collect::<Vec<_>>().join(" ");

        Ok(result)
    }

    /// Abstractive summarization: rewrite content
    fn abstractive_summarize(&self, content: &str, max_length: usize) -> Result<String, String> {
        // Stub: Real implementation would use LLM or neural abstractive model
        // For now, use aggressive extractive + rewriting heuristics

        let sentences = self.split_sentences(content);
        let key_phrases = self.extract_phrases(&sentences);

        let mut result = String::new();
        for phrase in key_phrases.iter().take(3) {
            if result.len() + phrase.len() + 2 > max_length {
                break;
            }
            if !result.is_empty() {
                result.push_str(". ");
            }
            result.push_str(phrase);
        }

        if result.is_empty() {
            result = self.extractive_summarize(content, max_length)?;
        }

        Ok(result)
    }

    /// Hybrid: extract + rewrite for coherence
    fn hybrid_summarize(&self, content: &str, max_length: usize) -> Result<String, String> {
        // Start with extractive
        let extracted = self.extractive_summarize(content, max_length)?;

        // Rewrite for coherence
        let rewritten = self.improve_coherence(&extracted);

        Ok(rewritten)
    }

    /// Extract key facts from content
    fn extract_key_facts(&self, _original: &str, summary: &str) -> Vec<KeyFact> {
        let mut facts = Vec::new();

        // Extract capitalized entities (simple heuristic)
        let words: Vec<&str> = summary.split_whitespace().collect();
        let mut entity_scores: HashMap<String, f32> = HashMap::new();

        for (idx, window) in words.windows(2).enumerate() {
            if is_capitalized_entity(window[0], 2) {
                let entity = window[0].to_string();
                let score = (idx as f32 / words.len() as f32).max(0.5); // Recency + presence
                entity_scores
                    .entry(entity.clone())
                    .and_modify(|s| *s = (*s + score) / 2.0)
                    .or_insert(score);
            }
        }

        // Convert to KeyFacts
        for (entity, score) in entity_scores {
            facts.push(KeyFact {
                fact: entity.clone(),
                importance: score.min(1.0),
                source_id: format!("entity_{}", entity.to_lowercase()),
                fact_type: "entity".to_string(),
            });
        }

        // Sort by importance
        facts.sort_by(|a, b| b.importance.partial_cmp(&a.importance).unwrap_or(std::cmp::Ordering::Equal));

        facts.into_iter().take(5).collect()
    }

    /// Compute coherence metrics
    fn compute_coherence(&self, text: &str) -> f32 {
        let metrics = self.compute_coherence_metrics(text);

        // Average of all metrics
        (metrics.entity_coherence + metrics.flow_coherence + metrics.semantic_coherence) / 3.0
    }

    /// Score sentence for importance
    fn score_sentence(&self, sentence: &str, document: &str) -> f32 {
        let words: Vec<&str> = sentence.split_whitespace().collect();
        let unique_words: HashSet<_> = words.iter().cloned().collect();

        // TF-IDF-like scoring
        let mut score = 0.0;

        for word in &unique_words {
            let tf = words.iter().filter(|w| *w == word).count() as f32 / words.len() as f32;
            let doc_freq = document.split_whitespace().filter(|w| w == word).count() as f32;
            let idf = (document.len() as f32 / doc_freq.max(1.0)).log2();

            score += tf * idf;
        }

        // Boost for position (earlier sentences more important)
        score = score * 0.9 + 0.1;

        score.min(1.0)
    }

    /// Split text into sentences
    fn split_sentences(&self, text: &str) -> Vec<&str> {
        text.split('.').map(|s| s.trim()).filter(|s| !s.is_empty()).collect()
    }

    /// Extract key phrases from sentences
    fn extract_phrases(&self, sentences: &[&str]) -> Vec<String> {
        let mut phrases = Vec::new();

        for sentence in sentences {
            let words: Vec<&str> = sentence.split_whitespace().collect();

            // Extract noun phrases (capitalized sequences)
            let mut phrase = String::new();
            for word in words {
                if is_capitalized_entity(word, 1) {
                    if !phrase.is_empty() {
                        phrase.push(' ');
                    }
                    phrase.push_str(word);
                } else if !phrase.is_empty() {
                    phrases.push(phrase.clone());
                    phrase.clear();
                }
            }

            if !phrase.is_empty() {
                phrases.push(phrase);
            }
        }

        phrases
    }

    /// Improve coherence by rewriting
    fn improve_coherence(&self, text: &str) -> String {
        // Simple heuristic: add connectors between sentences
        let sentences = self.split_sentences(text);

        let mut result = String::new();
        for (idx, sent) in sentences.iter().enumerate() {
            if idx > 0 {
                // Add transition word
                let transitions = vec!["Furthermore, ", "Moreover, ", "Additionally, ", "However, "];
                let transition = transitions[idx % transitions.len()];
                result.push_str(transition);
            }

            result.push_str(sent);
            if !sent.ends_with('.') {
                result.push('.');
            }
            result.push(' ');
        }

        result.trim().to_string()
    }

    /// Compute coherence metrics
    fn compute_coherence_metrics(&self, text: &str) -> CoherenceMetrics {
        let sentences = self.split_sentences(text);

        // Entity coherence: how well entities flow
        let entity_coherence = if sentences.len() > 1 {
            let mut coherence = 0.0;
            for window in sentences.windows(2) {
                let entities1 = self.extract_entities(window[0]);
                let entities2 = self.extract_entities(window[1]);

                let overlap = entities1
                    .iter()
                    .filter(|e| entities2.contains(e))
                    .count();
                coherence += overlap as f32 / (entities1.len().max(entities2.len()) as f32).max(1.0);
            }
            (coherence / (sentences.len() - 1) as f32).min(1.0)
        } else {
            0.8
        };

        // Flow coherence: sentence length variation
        let lengths: Vec<usize> = sentences.iter().map(|s| s.len()).collect();
        let avg_len = lengths.iter().sum::<usize>() as f32 / lengths.len() as f32;
        let variance = lengths
            .iter()
            .map(|l| (*l as f32 - avg_len).powi(2))
            .sum::<f32>()
            / lengths.len() as f32;
        let flow_coherence = (1.0 / (1.0 + variance / 1000.0)).min(1.0);

        // Semantic coherence: vocabulary richness
        let words: Vec<&str> = text.split_whitespace().collect();
        let unique_words: HashSet<_> = words.iter().cloned().collect();
        let semantic_coherence = (unique_words.len() as f32 / words.len() as f32).min(1.0);

        CoherenceMetrics {
            entity_coherence,
            flow_coherence,
            semantic_coherence,
        }
    }

    /// Extract entities from text (DRY: uses is_capitalized_entity)
    fn extract_entities(&self, text: &str) -> HashSet<String> {
        let mut entities = HashSet::new();
        let words: Vec<&str> = text.split_whitespace().collect();

        for word in words {
            if is_capitalized_entity(word, 2) {
                entities.insert(word.to_lowercase());
            }
        }

        entities
    }
}

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

    fn sample_content() -> &'static str {
        "Kubernetes is a container orchestration platform. Docker is used for containerization. \
         Kubernetes manages Docker containers at scale. Microservices are the primary use case. \
         Load balancing and auto-scaling are key features."
    }

    #[test]
    fn test_summarizer_creation() {
        let summarizer = Summarizer::new();
        assert_eq!(std::mem::size_of_val(&summarizer), 0); // Zero-sized type
    }

    #[test]
    fn test_extractive_summarize() {
        let summarizer = Summarizer::new();
        let result = summarizer.extractive_summarize(sample_content(), 100);
        assert!(result.is_ok());
        assert!(result.unwrap().len() <= 150); // Allow some overflow
    }

    #[test]
    fn test_abstractive_summarize() {
        let summarizer = Summarizer::new();
        let result = summarizer.abstractive_summarize(sample_content(), 100);
        assert!(result.is_ok());
        assert!(!result.unwrap().is_empty());
    }

    #[test]
    fn test_hybrid_summarize() {
        let summarizer = Summarizer::new();
        let result = summarizer.hybrid_summarize(sample_content(), 100);
        assert!(result.is_ok());
        assert!(!result.unwrap().is_empty());
    }

    #[test]
    fn test_summarize_extractive() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 100, SummarizationStrategy::Extractive);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert!(summary.compression_ratio < 1.0);
    }

    #[test]
    fn test_summarize_abstractive() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 100, SummarizationStrategy::Abstractive);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert!(!summary.text.is_empty());
    }

    #[test]
    fn test_summarize_hybrid() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 100, SummarizationStrategy::Hybrid);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert!(summary.strategy == SummarizationStrategy::Hybrid);
    }

    #[test]
    fn test_summary_compression_ratio() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 100, SummarizationStrategy::Extractive);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert!(summary.compression_ratio < 1.0);
    }

    #[test]
    fn test_summary_key_facts() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 200, SummarizationStrategy::Extractive);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert!(!summary.key_facts.is_empty());
    }

    #[test]
    fn test_summary_coherence() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 200, SummarizationStrategy::Hybrid);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert!(summary.coherence >= 0.0 && summary.coherence <= 1.0);
    }

    #[test]
    fn test_split_sentences() {
        let summarizer = Summarizer::new();
        let sentences = summarizer.split_sentences(sample_content());
        assert!(sentences.len() > 1);
    }

    #[test]
    fn test_score_sentence() {
        let summarizer = Summarizer::new();
        let score = summarizer.score_sentence("Kubernetes is important", sample_content());
        assert!(score >= 0.0 && score <= 1.0);
    }

    #[test]
    fn test_extract_key_facts() {
        let summarizer = Summarizer::new();
        let facts = summarizer.extract_key_facts(sample_content(), sample_content());
        assert!(!facts.is_empty());
    }

    #[test]
    fn test_compute_coherence() {
        let summarizer = Summarizer::new();
        let coherence = summarizer.compute_coherence(sample_content());
        assert!(coherence >= 0.0 && coherence <= 1.0);
    }

    #[test]
    fn test_compute_coherence_metrics() {
        let summarizer = Summarizer::new();
        let metrics = summarizer.compute_coherence_metrics(sample_content());
        assert!(metrics.entity_coherence >= 0.0 && metrics.entity_coherence <= 1.0);
        assert!(metrics.flow_coherence >= 0.0 && metrics.flow_coherence <= 1.0);
        assert!(metrics.semantic_coherence >= 0.0 && metrics.semantic_coherence <= 1.0);
    }

    #[test]
    fn test_improve_coherence() {
        let summarizer = Summarizer::new();
        let improved = summarizer.improve_coherence("Sentence one. Sentence two.");
        assert!(improved.contains("Furthermore") || improved.contains("Moreover"));
    }

    #[test]
    fn test_extract_entities() {
        let summarizer = Summarizer::new();
        let entities = summarizer.extract_entities("Kubernetes and Docker are tools");
        assert!(entities.contains("kubernetes"));
        assert!(entities.contains("docker"));
    }

    #[test]
    fn test_extract_phrases() {
        let summarizer = Summarizer::new();
        let sentences = vec!["Kubernetes is a platform", "Docker is a tool"];
        let phrases = summarizer.extract_phrases(&sentences);
        assert!(!phrases.is_empty());
    }

    #[test]
    fn test_summarize_empty_content() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize("", 100, SummarizationStrategy::Extractive);
        assert!(result.is_err());
    }

    #[test]
    fn test_summarize_too_short_max_length() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 10, SummarizationStrategy::Extractive);
        assert!(result.is_err());
    }

    #[test]
    fn test_summary_original_length() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 100, SummarizationStrategy::Extractive);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert_eq!(summary.original_length, sample_content().len());
    }

    #[test]
    fn test_summary_strategy_tracked() {
        let summarizer = Summarizer::new();
        let result = summarizer.summarize(sample_content(), 100, SummarizationStrategy::Extractive);
        assert!(result.is_ok());
        let summary = result.unwrap();
        assert_eq!(summary.strategy, SummarizationStrategy::Extractive);
    }

    #[test]
    fn test_key_fact_structure() {
        let fact = KeyFact {
            fact: "Kubernetes".to_string(),
            importance: 0.9,
            source_id: "entity_kubernetes".to_string(),
            fact_type: "entity".to_string(),
        };
        assert_eq!(fact.importance, 0.9);
    }

    #[test]
    fn test_coherence_metrics_structure() {
        let metrics = CoherenceMetrics {
            entity_coherence: 0.8,
            flow_coherence: 0.9,
            semantic_coherence: 0.7,
        };
        assert!(metrics.entity_coherence > 0.7);
    }

    #[test]
    fn test_summary_structure() {
        let summary = Summary {
            original_length: 100,
            text: "Summary".to_string(),
            summary_length: 7,
            compression_ratio: 0.07,
            key_facts: vec![],
            coherence: 0.8,
            strategy: SummarizationStrategy::Extractive,
        };
        assert!(summary.compression_ratio < 1.0);
    }

    #[test]
    fn test_summarization_strategies() {
        let strategies = vec![
            SummarizationStrategy::Extractive,
            SummarizationStrategy::Abstractive,
            SummarizationStrategy::Hybrid,
        ];
        assert_eq!(strategies.len(), 3);
    }

    #[test]
    fn test_sentence_scoring_consistency() {
        let summarizer = Summarizer::new();
        let score1 = summarizer.score_sentence("Kubernetes", sample_content());
        let score2 = summarizer.score_sentence("Kubernetes", sample_content());
        assert_eq!(score1, score2);
    }

    #[test]
    fn test_long_content_summarization() {
        let summarizer = Summarizer::new();
        let long_content = sample_content().repeat(10);
        let result = summarizer.summarize(&long_content, 200, SummarizationStrategy::Extractive);
        assert!(result.is_ok());
    }

    #[test]
    fn test_short_content_summarization() {
        let summarizer = Summarizer::new();
        let short = "Kubernetes is great.";
        let result = summarizer.summarize(short, 50, SummarizationStrategy::Extractive);
        assert!(result.is_ok());
    }
}