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

/// Phase 5: Chunk Metadata Index
///
/// Extract and index chunk metadata for improved scoring:
/// 1. Heading extraction (markdown hierarchy)
/// 2. Key term extraction (TF-IDF top terms)
/// 3. Category inference (error|solution|tool|concept)
/// 4. Metadata-based scoring boost
///
/// Benefits:
/// - Better semantic understanding (category context)
/// - Faster ranking (metadata pre-computed)
/// - Query intent matching (match query intent to chunk category)

use anyhow::Result;
use std::collections::{HashMap, HashSet};

/// Chunk category for scoring context
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ChunkCategory {
    Error,       // Problem statement, error trace
    Solution,    // Fix, workaround, resolution
    Tool,        // Command, API, configuration
    Concept,     // Theory, explanation, design pattern
    Reference,   // Documentation, spec, standard
    Unknown,
}

impl ChunkCategory {
    pub fn as_str(&self) -> &str {
        match self {
            ChunkCategory::Error => "error",
            ChunkCategory::Solution => "solution",
            ChunkCategory::Tool => "tool",
            ChunkCategory::Concept => "concept",
            ChunkCategory::Reference => "reference",
            ChunkCategory::Unknown => "unknown",
        }
    }

    pub fn from_str(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "error" => ChunkCategory::Error,
            "solution" => ChunkCategory::Solution,
            "tool" => ChunkCategory::Tool,
            "concept" => ChunkCategory::Concept,
            "reference" => ChunkCategory::Reference,
            _ => ChunkCategory::Unknown,
        }
    }
}

/// Query intent for matching with chunk categories
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum QueryIntent {
    FixError,      // "fix", "debug", "troubleshoot"
    LearnConcept,  // "explain", "understand", "how does"
    UseTool,       // "use", "run", "call", "api"
    FindReference, // "what is", "definition", "spec"
    Unknown,
}

impl QueryIntent {
    /// Match query intent to chunk categories for boost
    pub fn matching_categories(&self) -> Vec<ChunkCategory> {
        match self {
            QueryIntent::FixError => vec![ChunkCategory::Error, ChunkCategory::Solution],
            QueryIntent::LearnConcept => vec![ChunkCategory::Concept, ChunkCategory::Reference],
            QueryIntent::UseTool => vec![ChunkCategory::Tool, ChunkCategory::Solution],
            QueryIntent::FindReference => vec![ChunkCategory::Reference, ChunkCategory::Concept],
            QueryIntent::Unknown => vec![
                ChunkCategory::Error,
                ChunkCategory::Solution,
                ChunkCategory::Tool,
                ChunkCategory::Concept,
            ],
        }
    }
}

/// Extracted chunk metadata
#[derive(Debug, Clone)]
pub struct ChunkMetadata {
    pub chunk_id: String,
    pub heading: Option<String>,      // Highest-level heading
    pub key_terms: Vec<String>,       // Top TF-IDF terms
    pub category: ChunkCategory,
    pub category_confidence: f32,     // 0.0-1.0
}

/// Metadata Extractor
pub struct MetadataExtractor;

impl MetadataExtractor {
    /// Extract heading (first markdown heading)
    pub fn extract_heading(text: &str) -> Option<String> {
        for line in text.lines() {
            if line.starts_with('#') {
                return Some(
                    line
                        .trim_start_matches('#')
                        .trim()
                        .to_string()
                );
            }
        }
        None
    }

    /// Extract top K key terms by word frequency
    pub fn extract_key_terms(text: &str, top_k: usize) -> Vec<String> {
        let mut term_counts: HashMap<String, usize> = HashMap::new();

        // Count word frequencies (case-insensitive, skip common words)
        let stopwords = vec![
            "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with",
            "by", "from", "is", "are", "was", "be", "have", "has", "do", "does", "did",
        ];

        for word in text.split_whitespace() {
            let cleaned = word
                .to_lowercase()
                .chars()
                .filter(|c| c.is_alphanumeric())
                .collect::<String>();

            if !cleaned.is_empty()
                && cleaned.len() > 3
                && !stopwords.contains(&cleaned.as_str())
            {
                *term_counts.entry(cleaned).or_insert(0) += 1;
            }
        }

        // Sort by frequency descending
        let mut terms: Vec<_> = term_counts.into_iter().collect();
        terms.sort_by(|a, b| b.1.cmp(&a.1));

        terms.into_iter().take(top_k).map(|(term, _)| term).collect()
    }

    /// Infer category from text content
    pub fn infer_category(text: &str) -> (ChunkCategory, f32) {
        let lower = text.to_lowercase();

        // Error indicators
        if lower.contains("error") || lower.contains("failed") || lower.contains("crash")
            || lower.contains("bug") || lower.contains("exception")
        {
            return (ChunkCategory::Error, 0.9);
        }

        // Solution indicators
        if lower.contains("fix") || lower.contains("solution") || lower.contains("workaround")
            || lower.contains("resolved") || lower.contains("configure")
        {
            return (ChunkCategory::Solution, 0.85);
        }

        // Tool indicators
        if lower.contains("command") || lower.contains("api") || lower.contains("cli")
            || lower.contains("usage:") || lower.contains("$ ")
        {
            return (ChunkCategory::Tool, 0.8);
        }

        // Concept indicators
        if lower.contains("explain") || lower.contains("concept") || lower.contains("principle")
            || lower.contains("design") || lower.contains("pattern")
        {
            return (ChunkCategory::Concept, 0.8);
        }

        // Reference indicators
        if lower.contains("reference") || lower.contains("documentation") || lower.contains("spec")
            || lower.contains("standard") || lower.contains("definition")
        {
            return (ChunkCategory::Reference, 0.75);
        }

        (ChunkCategory::Unknown, 0.3)
    }

    /// Infer query intent from query text
    pub fn infer_query_intent(query: &str) -> QueryIntent {
        let lower = query.to_lowercase();

        if lower.contains("fix") || lower.contains("debug") || lower.contains("troubleshoot")
            || lower.contains("error")
        {
            QueryIntent::FixError
        } else if lower.contains("explain") || lower.contains("understand")
            || lower.contains("how does") || lower.contains("what is")
        {
            QueryIntent::LearnConcept
        } else if lower.contains("use") || lower.contains("run") || lower.contains("call")
            || lower.contains("api")
        {
            QueryIntent::UseTool
        } else if lower.contains("reference") || lower.contains("definition") || lower.contains("spec")
        {
            QueryIntent::FindReference
        } else {
            QueryIntent::Unknown
        }
    }

    /// Full metadata extraction
    pub fn extract(chunk_id: &str, text: &str) -> ChunkMetadata {
        let (category, confidence) = Self::infer_category(text);

        ChunkMetadata {
            chunk_id: chunk_id.to_string(),
            heading: Self::extract_heading(text),
            key_terms: Self::extract_key_terms(text, 5),
            category,
            category_confidence: confidence,
        }
    }
}

/// Metadata-based Scoring Boost
pub struct MetadataBooster {
    category_boost: HashMap<ChunkCategory, f32>,
}

impl MetadataBooster {
    pub fn new() -> Self {
        let mut category_boost = HashMap::new();
        category_boost.insert(ChunkCategory::Error, 0.1);      // 10% boost
        category_boost.insert(ChunkCategory::Solution, 0.2);   // 20% boost
        category_boost.insert(ChunkCategory::Tool, 0.15);      // 15% boost
        category_boost.insert(ChunkCategory::Concept, 0.1);    // 10% boost
        category_boost.insert(ChunkCategory::Reference, 0.05); // 5% boost
        category_boost.insert(ChunkCategory::Unknown, 0.0);    // No boost

        Self { category_boost }
    }

    /// Calculate boost factor for query intent + chunk category
    pub fn calculate_boost(
        &self,
        query_intent: QueryIntent,
        chunk_metadata: &ChunkMetadata,
    ) -> f32 {
        let matching_categories = query_intent.matching_categories();

        if matching_categories.contains(&chunk_metadata.category) {
            // Match: apply boost
            let base_boost = self
                .category_boost
                .get(&chunk_metadata.category)
                .copied()
                .unwrap_or(0.0);

            // Scale by category confidence
            base_boost * chunk_metadata.category_confidence
        } else {
            0.0 // No boost for mismatched categories
        }
    }

    /// Apply boost to base score
    pub fn apply_boost(&self, base_score: f32, boost: f32) -> f32 {
        (base_score + boost).min(1.0)
    }
}

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

    #[test]
    fn test_extract_heading() {
        let text = "# Debugging Kubernetes Pods\n\nSome content";
        let heading = MetadataExtractor::extract_heading(text);
        assert_eq!(heading, Some("Debugging Kubernetes Pods".to_string()));
    }

    #[test]
    fn test_extract_heading_none() {
        let text = "No heading here\n\nJust content";
        let heading = MetadataExtractor::extract_heading(text);
        assert_eq!(heading, None);
    }

    #[test]
    fn test_extract_key_terms() {
        let text = "kubernetes pod debugging pod kubernetes deployment";
        let terms = MetadataExtractor::extract_key_terms(text, 3);
        assert!(terms.contains(&"kubernetes".to_string()));
        assert!(terms.len() <= 3);
    }

    #[test]
    fn test_infer_category_error() {
        let text = "Pod crash error: exception during startup";
        let (category, _) = MetadataExtractor::infer_category(text);
        assert_eq!(category, ChunkCategory::Error);
    }

    #[test]
    fn test_infer_category_solution() {
        let text = "To fix this issue, configure the pod like this...";
        let (category, _) = MetadataExtractor::infer_category(text);
        assert_eq!(category, ChunkCategory::Solution);
    }

    #[test]
    fn test_infer_category_tool() {
        let text = "Usage: kubectl get pods\n\n$ kubectl apply -f config.yaml";
        let (category, _) = MetadataExtractor::infer_category(text);
        assert_eq!(category, ChunkCategory::Tool);
    }

    #[test]
    fn test_infer_category_concept() {
        let text = "The principle of kuberentes design patterns is...";
        let (category, _) = MetadataExtractor::infer_category(text);
        assert_eq!(category, ChunkCategory::Concept);
    }

    #[test]
    fn test_infer_query_intent_fix_error() {
        let intent = MetadataExtractor::infer_query_intent("How do I fix a pod crash?");
        assert_eq!(intent, QueryIntent::FixError);
    }

    #[test]
    fn test_infer_query_intent_learn() {
        let intent = MetadataExtractor::infer_query_intent("Explain kubernetes concepts");
        assert_eq!(intent, QueryIntent::LearnConcept);
    }

    #[test]
    fn test_infer_query_intent_tool() {
        let intent = MetadataExtractor::infer_query_intent("How to use the kubectl API?");
        assert_eq!(intent, QueryIntent::UseTool);
    }

    #[test]
    fn test_full_metadata_extraction() {
        let text = "# Pod Debugging\n\nError: CrashLoopBackOff. Solution: check logs";
        let metadata = MetadataExtractor::extract("chunk1", text);

        assert_eq!(metadata.chunk_id, "chunk1");
        assert_eq!(metadata.heading, Some("Pod Debugging".to_string()));
        assert!(!metadata.key_terms.is_empty());
        assert!(metadata.category_confidence > 0.0);
    }

    #[test]
    fn test_metadata_booster_matching_category() {
        let booster = MetadataBooster::new();
        let metadata = ChunkMetadata {
            chunk_id: "chunk1".to_string(),
            heading: None,
            key_terms: vec![],
            category: ChunkCategory::Solution,
            category_confidence: 0.9,
        };

        let boost = booster.calculate_boost(QueryIntent::FixError, &metadata);
        assert!(boost > 0.0); // Solution matches FixError intent
    }

    #[test]
    fn test_metadata_booster_mismatched_category() {
        let booster = MetadataBooster::new();
        let metadata = ChunkMetadata {
            chunk_id: "chunk1".to_string(),
            heading: None,
            key_terms: vec![],
            category: ChunkCategory::Reference,
            category_confidence: 0.8,
        };

        let boost = booster.calculate_boost(QueryIntent::FixError, &metadata);
        assert_eq!(boost, 0.0); // Reference doesn't match FixError intent
    }

    #[test]
    fn test_apply_boost_caps_at_1() {
        let booster = MetadataBooster::new();
        let score = booster.apply_boost(0.95, 0.2);
        assert_eq!(score, 1.0); // Capped at 1.0
    }

    #[test]
    fn test_category_to_str() {
        assert_eq!(ChunkCategory::Error.as_str(), "error");
        assert_eq!(ChunkCategory::Solution.as_str(), "solution");
        assert_eq!(ChunkCategory::Unknown.as_str(), "unknown");
    }

    #[test]
    fn test_category_from_str() {
        assert_eq!(ChunkCategory::from_str("error"), ChunkCategory::Error);
        assert_eq!(ChunkCategory::from_str("SOLUTION"), ChunkCategory::Solution);
        assert_eq!(ChunkCategory::from_str("unknown"), ChunkCategory::Unknown);
    }
}