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

use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Query Context: normalized query + analysis for hybrid search
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct QueryContext {
    // Original query
    pub raw_query: String,
    
    // Normalized (lowercased, trimmed)
    pub normalized_query: String,
    
    // Tokenized terms
    pub tokens: Vec<String>,
    
    // Extracted named entities (year, names, keywords)
    pub entities: HashMap<String, String>,
    
    // Query embedding (to be generated by LLM)
    pub embedding: Option<Vec<f32>>,
    
    // Analysis results
    pub token_count: usize,
    pub has_special_syntax: bool,      // #tag, @mention, "exact phrase"
    pub has_date_filters: bool,        // 2024, "this month"
    pub has_negation: bool,            // -word, NOT phrase
    pub question_type: QuestionType,
    
    // Routing decision
    pub search_strategy: SearchStrategy,
    pub confidence: f32,  // How confident in the routing decision (0.0-1.0)
}

/// Question type classification
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum QuestionType {
    Factual,           // "What is X?" "Define Y"
    Procedural,        // "How do I..." "Steps to..."
    Comparative,       // "Compare X and Y" "Difference between..."
    Troubleshooting,   // "Fix broken..." "Error: ..."
    Navigational,      // "Where is X?" "Find documents about..."
    Open,              // General conversational
}

/// Search strategy (determines which engines to use)
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum SearchStrategy {
    Hybrid,           // Both pgvector + OpenSearch
    SemanticOnly,     // pgvector only (if OpenSearch down)
    LexicalOnly,      // OpenSearch only (if embedding model down)
    LexicalFirst,     // OpenSearch to narrow, then semantic rerank
}

/// RRF (Reciprocal Rank Fusion) configuration
#[derive(Clone, Debug)]
pub struct RRFConfig {
    pub k: f32,                    // Constant (usually 60)
    pub retrieve_k: usize,         // Top-K from each engine (usually 50)
    pub final_k: usize,            // Final top-K to return (usually 10)
}

impl Default for RRFConfig {
    fn default() -> Self {
        Self {
            k: 60.0,
            retrieve_k: 50,
            final_k: 10,
        }
    }
}

/// Query Optimization Engine
pub struct QueryOptimizer {
    enable_entity_extraction: bool,
    enable_question_classification: bool,
}

impl QueryOptimizer {
    pub fn new() -> Self {
        Self {
            enable_entity_extraction: true,
            enable_question_classification: true,
        }
    }

    /// Main entry point: construct query context from user input
    pub async fn optimize_query(&self, raw_query: &str) -> Result<QueryContext> {
        // Stage 1: Normalize
        let normalized = self.normalize_query(raw_query);
        
        // Stage 2: Tokenize
        let tokens = self.tokenize(&normalized);
        
        // Stage 3: Extract entities
        let entities = if self.enable_entity_extraction {
            self.extract_entities(raw_query, &tokens)
        } else {
            HashMap::new()
        };
        
        // Stage 4: Analyze query characteristics
        let token_count = tokens.len();
        let has_special_syntax = self.detect_special_syntax(raw_query);
        let has_date_filters = self.detect_date_filters(&tokens);
        let has_negation = self.detect_negation(&tokens);
        
        // Stage 5: Classify question type
        let question_type = if self.enable_question_classification {
            self.classify_question(raw_query, &tokens)
        } else {
            QuestionType::Open
        };
        
        // Stage 6: Route to search strategy
        let (search_strategy, confidence) = self.route_query(
            token_count,
            has_special_syntax,
            has_date_filters,
            has_negation,
            &question_type,
        );
        
        Ok(QueryContext {
            raw_query: raw_query.to_string(),
            normalized_query: normalized,
            tokens,
            entities,
            embedding: None,
            token_count,
            has_special_syntax,
            has_date_filters,
            has_negation,
            question_type,
            search_strategy,
            confidence,
        })
    }

    /// Stage 1: Normalize query
    fn normalize_query(&self, query: &str) -> String {
        query
            .trim()
            .to_lowercase()
            .replace("  ", " ")  // Remove double spaces
    }

    /// Stage 2: Tokenize
    fn tokenize(&self, query: &str) -> Vec<String> {
        query
            .split_whitespace()
            .map(|s| s.to_string())
            .collect()
    }

    /// Stage 3: Extract entities (years, names, keywords)
    fn extract_entities(&self, raw_query: &str, tokens: &[String]) -> HashMap<String, String> {
        let mut entities = HashMap::new();
        
        for token in tokens {
            // Year detection: YYYY format
            if token.len() == 4 {
                if let Ok(year) = token.parse::<u32>() {
                    if year >= 2000 && year <= 2100 {
                        entities.insert("year".to_string(), token.clone());
                    }
                }
            }
        }
        
        // Detect quoted phrases
        if raw_query.contains('"') {
            let parts: Vec<&str> = raw_query.split('"').collect();
            if parts.len() >= 3 {
                let quoted_phrase = parts[1].to_string();
                entities.insert("exact_phrase".to_string(), quoted_phrase);
            }
        }
        
        entities
    }

    /// Stage 4: Detect special syntax (#tag, @mention, "phrases")
    fn detect_special_syntax(&self, query: &str) -> bool {
        query.contains('#') || query.contains('@') || query.contains('"')
    }

    /// Stage 4: Detect date filters
    fn detect_date_filters(&self, tokens: &[String]) -> bool {
        let date_keywords = vec![
            "this", "last", "next",
            "2024", "2025", "2026",
            "january", "february", "march", "april", "may", "june",
            "july", "august", "september", "october", "november", "december",
            "week", "month", "year", "day", "today", "yesterday", "tomorrow",
        ];
        
        tokens.iter().any(|t| date_keywords.contains(&t.as_str()))
    }

    /// Stage 4: Detect negation
    fn detect_negation(&self, tokens: &[String]) -> bool {
        tokens.iter().any(|t| t == "-" || t == "not" || t == "no" || t.starts_with("-"))
    }

    /// Stage 5: Classify question type
    fn classify_question(&self, raw_query: &str, tokens: &[String]) -> QuestionType {
        let query_lower = raw_query.to_lowercase();
        
        // Check first token for question words
        if tokens.is_empty() {
            return QuestionType::Open;
        }
        
        let first_token = &tokens[0];
        
        match first_token.as_str() {
            // Procedural questions
            t if t == "how" => QuestionType::Procedural,
            t if t == "what" => {
                if query_lower.contains("difference") || query_lower.contains("between") {
                    QuestionType::Comparative
                } else {
                    QuestionType::Factual
                }
            }
            // Comparative
            t if t == "compare" || t == "compare" => QuestionType::Comparative,
            // Troubleshooting
            t if t == "fix" || t == "error" || t == "broken" || t == "debug" => {
                QuestionType::Troubleshooting
            }
            // Navigational
            t if t == "where" || t == "find" || t == "show" => QuestionType::Navigational,
            _ => {
                // Heuristics based on content
                if query_lower.contains("how") {
                    QuestionType::Procedural
                } else if query_lower.contains("fix") || query_lower.contains("error") {
                    QuestionType::Troubleshooting
                } else {
                    QuestionType::Open
                }
            }
        }
    }

    /// Stage 6: Route to search strategy
    fn route_query(
        &self,
        token_count: usize,
        has_special_syntax: bool,
        has_date_filters: bool,
        _has_negation: bool,
        question_type: &QuestionType,
    ) -> (SearchStrategy, f32) {
        // Very short queries: lexical better
        if token_count < 3 {
            return (SearchStrategy::LexicalOnly, 0.8);
        }
        
        // Special syntax: preserve exact matches with lexical
        if has_special_syntax {
            if has_date_filters {
                // Special syntax + dates = use lexical to narrow, then semantic
                return (SearchStrategy::LexicalFirst, 0.85);
            } else {
                // Just special syntax = lexical only
                return (SearchStrategy::LexicalOnly, 0.8);
            }
        }
        
        // Date filters present: use cascading (lexical → semantic)
        if has_date_filters {
            return (SearchStrategy::LexicalFirst, 0.9);
        }
        
        // Question type heuristics
        match question_type {
            // Factual questions usually work well with semantic
            QuestionType::Factual => (SearchStrategy::Hybrid, 0.9),
            
            // Procedural questions benefit from both (exact steps + understanding)
            QuestionType::Procedural => (SearchStrategy::Hybrid, 0.95),
            
            // Troubleshooting needs both (exact errors + semantic understanding)
            QuestionType::Troubleshooting => (SearchStrategy::Hybrid, 0.95),
            
            // Comparative: hybrid needed (understanding + multiple docs)
            QuestionType::Comparative => (SearchStrategy::Hybrid, 0.9),
            
            // Navigational: lexical good for finding specific things
            QuestionType::Navigational => (SearchStrategy::LexicalFirst, 0.85),
            
            // Open/general: hybrid default
            QuestionType::Open => (SearchStrategy::Hybrid, 0.8),
        }
    }
}

/// RRF Fusion Engine
pub struct RRFFusion {
    config: RRFConfig,
}

impl RRFFusion {
    pub fn new(config: RRFConfig) -> Self {
        Self { config }
    }

    /// Fuse two ranked lists using Reciprocal Rank Fusion
    pub fn fuse(
        &self,
        semantic_results: Vec<(String, f32)>,  // (id, score)
        lexical_results: Vec<(String, f32)>,
    ) -> Vec<(String, f32)> {
        use std::collections::HashMap;
        
        let mut fused_scores: HashMap<String, f32> = HashMap::new();
        
        // Add semantic ranks with RRF formula: 1 / (k + rank)
        for (rank, (id, _)) in semantic_results.into_iter().enumerate() {
            let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0);
            fused_scores.insert(id, rrf_score);
        }
        
        // Add lexical ranks (combine if already present)
        for (rank, (id, _)) in lexical_results.into_iter().enumerate() {
            let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0);
            *fused_scores.entry(id).or_insert(0.0) += rrf_score;
        }
        
        // Sort by combined RRF score
        let mut results: Vec<_> = fused_scores.into_iter().collect();
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        
        // Take top-k
        results.truncate(self.config.final_k);
        
        results
    }

    /// Alternative: Weighted Linear Fusion
    pub fn fuse_weighted(
        &self,
        semantic_results: Vec<(String, f32)>,
        lexical_results: Vec<(String, f32)>,
        semantic_weight: f32,
        lexical_weight: f32,
    ) -> Vec<(String, f32)> {
        use std::collections::HashMap;
        
        // Normalize scores to [0.0, 1.0]
        let sem_norm = self.normalize_scores(&semantic_results);
        let lex_norm = self.normalize_scores(&lexical_results);
        
        let sem_map: HashMap<String, f32> = sem_norm.into_iter().collect();
        let lex_map: HashMap<String, f32> = lex_norm.into_iter().collect();
        
        // Merge all IDs
        let mut all_ids = std::collections::HashSet::new();
        all_ids.extend(sem_map.keys().cloned());
        all_ids.extend(lex_map.keys().cloned());
        
        // Calculate weighted scores
        let mut results: Vec<_> = all_ids
            .into_iter()
            .map(|id| {
                let sem_score = sem_map.get(&id).copied().unwrap_or(0.0);
                let lex_score = lex_map.get(&id).copied().unwrap_or(0.0);
                
                let weighted_score = semantic_weight * sem_score + lexical_weight * lex_score;
                (id, weighted_score)
            })
            .collect();
        
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        results.truncate(self.config.final_k);
        
        results
    }

    /// Normalize scores to [0.0, 1.0] range using min-max
    fn normalize_scores(&self, results: &[(String, f32)]) -> Vec<(String, f32)> {
        if results.is_empty() {
            return Vec::new();
        }
        
        let min_score = results.iter().map(|(_, s)| s).fold(f32::INFINITY, |a, &b| a.min(b));
        let max_score = results.iter().map(|(_, s)| s).fold(f32::NEG_INFINITY, |a, &b| a.max(b));
        
        let range = max_score - min_score;
        
        if range < 0.001 {
            // All scores identical
            return results.iter().map(|(id, _)| (id.clone(), 0.5)).collect();
        }
        
        results
            .iter()
            .map(|(id, score)| {
                let normalized = (score - min_score) / range;
                (id.clone(), normalized)
            })
            .collect()
    }
}

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

    #[tokio::test]
    async fn test_query_optimization_procedural() {
        let optimizer = QueryOptimizer::new();
        let ctx = optimizer.optimize_query("How do I fix kubernetes port 8080?").await.unwrap();
        
        assert_eq!(ctx.question_type, QuestionType::Procedural);
        assert_eq!(ctx.search_strategy, SearchStrategy::Hybrid);
        assert!(ctx.confidence >= 0.9);
    }

    #[tokio::test]
    async fn test_query_optimization_short() {
        let optimizer = QueryOptimizer::new();
        let ctx = optimizer.optimize_query("fix port").await.unwrap();
        
        assert_eq!(ctx.token_count, 2);
        assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly);
    }

    #[tokio::test]
    async fn test_query_optimization_special_syntax() {
        let optimizer = QueryOptimizer::new();
        let ctx = optimizer.optimize_query("kubernetes #networking @devops").await.unwrap();
        
        assert!(ctx.has_special_syntax);
        assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly);
    }

    #[test]
    fn test_rrf_fusion() {
        let fusion = RRFFusion::new(RRFConfig::default());
        
        let semantic = vec![
            ("doc1".to_string(), 0.95),
            ("doc2".to_string(), 0.88),
            ("doc3".to_string(), 0.82),
        ];
        
        let lexical = vec![
            ("doc1".to_string(), 8.5),
            ("doc4".to_string(), 7.2),
            ("doc2".to_string(), 6.8),
        ];
        
        let fused = fusion.fuse(semantic, lexical);
        
        // doc1 should be top (in both)
        assert_eq!(fused[0].0, "doc1");
        
        // RRF score: doc1 appears in both lists (rank 1 in each)
        // Score = 1/(60+1) + 1/(60+1) = 2/61 ≈ 0.0328
        assert!(fused[0].1 > 0.03 && fused[0].1 < 0.04, "Expected RRF score ~0.0328, got {}", fused[0].1);
    }

    #[test]
    fn test_weighted_fusion() {
        let fusion = RRFFusion::new(RRFConfig::default());
        
        let semantic = vec![
            ("doc1".to_string(), 0.95),
            ("doc2".to_string(), 0.88),
        ];
        
        let lexical = vec![
            ("doc1".to_string(), 8.5),
            ("doc3".to_string(), 7.2),
        ];
        
        let fused = fusion.fuse_weighted(semantic, lexical, 0.6, 0.4);
        
        // doc1 should rank highest (has both components)
        assert_eq!(fused[0].0, "doc1");
        
        // Score should be normalized and weighted
        // 0.6 * (0.95/0.95) + 0.4 * (8.5/8.5) = 1.0
        assert!((fused[0].1 - 1.0).abs() < 0.01);
    }
}