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

/// Query Orchestrator: Unified interface combining all phases 1-6
///
/// Orchestrates:
/// - Phase 1: Wiki-link graph traversal
/// - Phase 2: Scoring pipeline
/// - Phase 3: Hybrid retrieval (TF-IDF + semantic + RRF)
/// - Phase 4: LLM optimization (threshold, budget, dedup)
/// - Phase 5: Metadata enhancement (category + intent boost)
/// - Phase 6: Cache alignment (locality + pre-load)

use anyhow::Result;
use std::collections::HashMap;
use std::sync::Arc;
use mem_core::DocumentScorer;

use crate::hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate};
use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
use crate::chunk_metadata::{MetadataExtractor, MetadataBooster, QueryIntent};
use crate::cache_alignment::{KvCacheAligner, CachedChunk, CacheLocalityAnalyzer, RetrievalProfiler};

/// Complete query result with all metadata
#[derive(Debug, Clone)]
pub struct QueryResult {
    pub query: String,
    pub selected_chunks: Vec<OptimizedChunk>,
    pub selection_metrics: SelectionMetrics,
    pub cache_metrics: crate::cache_alignment::CacheMetrics,
    pub profiling: Vec<(String, u64)>, // stage -> duration_ms
    pub total_latency_ms: u64,
    pub query_intent: QueryIntent,
}

/// Chunk with all enrichments
#[derive(Debug, Clone)]
pub struct OptimizedChunk {
    pub id: String,
    pub text: String,
    pub tfidf_score: f32,
    pub semantic_score: f32,
    pub metadata_boost: f32,
    pub final_score: f32,
    pub category: crate::chunk_metadata::ChunkCategory,
    pub cache_distance: u32,
    pub cache_slot: u32,
}

/// Query execution context
pub struct QueryContext {
    pub project: String,
    pub wiki_root_doc: String,
    pub max_wiki_hops: u32,
    pub budget_bytes: usize,
    pub score_threshold: f32,
    pub dedup_threshold: f32,
    pub cache_capacity: usize,
}

impl Default for QueryContext {
    fn default() -> Self {
        Self {
            project: "default".to_string(),
            wiki_root_doc: "index.md".to_string(),
            max_wiki_hops: 3,
            budget_bytes: 8192,
            score_threshold: 0.6,
            dedup_threshold: 0.8,
            cache_capacity: 1000,
        }
    }
}

/// Orchestrator: combines all phases
pub struct QueryOrchestrator {
    retriever: Arc<HybridRetriever>,
    optimizer: Arc<ChunkOptimizer>,
    booster: Arc<MetadataBooster>,
    aligner: Arc<KvCacheAligner>,
    profiler: Arc<RetrievalProfiler>,
}

impl QueryOrchestrator {
    pub fn new(
        tfidf_scorer: Arc<mem_core::GlobalTfIdfScorer>,
        semantic_scorer: Arc<mem_core::SemanticScorer>,
        context: &QueryContext,
    ) -> Self {
        let retriever = Arc::new(HybridRetriever::new(
            tfidf_scorer.clone(),
            semantic_scorer.clone(),
        ));

        let optimizer = Arc::new(ChunkOptimizer::new(
            context.score_threshold,
            context.budget_bytes,
            context.dedup_threshold,
        ));

        let booster = Arc::new(MetadataBooster::new());
        let aligner = Arc::new(KvCacheAligner::new(4096, 100, context.cache_capacity));
        let profiler = Arc::new(RetrievalProfiler::new());

        Self {
            retriever,
            optimizer,
            booster,
            aligner,
            profiler,
        }
    }

    /// End-to-end query execution
    pub async fn execute(
        &self,
        query: &str,
        all_candidates: Vec<(String, String)>, // (doc_id, text)
        context: &QueryContext,
    ) -> Result<QueryResult> {
        let start = std::time::Instant::now();

        // Step 1: Infer query intent (Phase 5)
        let query_intent = MetadataExtractor::infer_query_intent(query);
        self.profiler.record("infer_intent", 1);

        // Step 2: Route retrieval (Phase 3)
        let has_wiki_scope = !context.wiki_root_doc.is_empty();
        let route = self.retriever.route_query(query, has_wiki_scope, false);
        self.profiler.record("route_selection", 2);

        // Step 3: Hybrid retrieval (Phase 3)
        let start_retrieval = std::time::Instant::now();
        let ranked = self
            .retriever
            .retrieve(query, all_candidates, route.clone())
            .await?;
        let retrieval_time = start_retrieval.elapsed().as_millis() as u64;
        self.profiler.record("hybrid_retrieval", retrieval_time);

        // Step 4: Convert to optimizable chunks
        let mut optimizable: Vec<OptimizableChunk> = ranked
            .into_iter()
            .map(|r| {
                let text_len = r.text.len();
                OptimizableChunk {
                    id: r.doc_id,
                    text: r.text,
                    score: r.final_score,
                    confidence: r.semantic_score, // Confidence from semantic
                    size_bytes: text_len,
                }
            })
            .collect();

        // Step 5: Metadata enhancement (Phase 5)
        let start_metadata = std::time::Instant::now();
        for chunk in &mut optimizable {
            let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text);
            let boost = self.booster.calculate_boost(query_intent, &metadata);
            chunk.score = self.booster.apply_boost(chunk.score, boost);
        }
        let metadata_time = start_metadata.elapsed().as_millis() as u64;
        self.profiler.record("metadata_boost", metadata_time);

        // Step 6: LLM optimization (Phase 4)
        let start_optimize = std::time::Instant::now();
        let (selected_opt, selection_metrics) = self.optimizer.optimize(optimizable.clone());
        let optimize_time = start_optimize.elapsed().as_millis() as u64;
        self.profiler.record("llm_optimize", optimize_time);

        // Step 7: Cache alignment (Phase 6)
        let start_cache = std::time::Instant::now();
        let cached: Vec<CachedChunk> = selected_opt
            .iter()
            .enumerate()
            .map(|(i, chunk)| CachedChunk {
                chunk_id: chunk.id.clone(),
                text: chunk.text.clone(),
                score: chunk.score,
                cache_distance: 0, // Would be computed from wiki-graph
                access_count: 1,
                last_accessed_slot: i as u32,
            })
            .collect();

        let slots = self.aligner.assign_slots(&cached);
        self.aligner.preload_hot_chunks(
            cached.iter().take(5).map(|c| (c.chunk_id.as_str(), c.text.as_str())).collect()
        )?;
        let cache_time = start_cache.elapsed().as_millis() as u64;
        self.profiler.record("cache_align", cache_time);

        // Step 8: Build optimized chunks with all metadata
        let mut optimized_chunks = Vec::new();
        for (i, chunk) in selected_opt.iter().enumerate() {
            let slot = slots.iter().find(|(id, _)| id == &chunk.id).map(|(_, s)| *s).unwrap_or(0);
            let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text);

            optimized_chunks.push(OptimizedChunk {
                id: chunk.id.clone(),
                text: chunk.text.clone(),
                tfidf_score: chunk.score * 0.4, // Approximate
                semantic_score: chunk.score * 0.6,
                metadata_boost: 0.0, // Already applied
                final_score: chunk.score,
                category: metadata.category,
                cache_distance: 0,
                cache_slot: slot,
            });
        }

        let total_latency = start.elapsed().as_millis() as u64;

        Ok(QueryResult {
            query: query.to_string(),
            selected_chunks: optimized_chunks,
            selection_metrics,
            cache_metrics: self.aligner.get_metrics(),
            profiling: self.profiler.summary(),
            total_latency_ms: total_latency,
            query_intent,
        })
    }
}

/// Memory projection for multi-project queries
pub struct MemoryProjection {
    projects: HashMap<String, Arc<QueryOrchestrator>>,
}

impl MemoryProjection {
    pub fn new() -> Self {
        Self {
            projects: HashMap::new(),
        }
    }

    pub fn register_project(
        &mut self,
        project: &str,
        orchestrator: Arc<QueryOrchestrator>,
    ) {
        self.projects.insert(project.to_string(), orchestrator);
    }

    pub async fn query_project(
        &self,
        project: &str,
        query: &str,
        candidates: Vec<(String, String)>,
        context: &QueryContext,
    ) -> Result<QueryResult> {
        let orchestrator = self
            .projects
            .get(project)
            .ok_or_else(|| anyhow::anyhow!("Project not found: {}", project))?;

        orchestrator.execute(query, candidates, context).await
    }

    pub fn projects(&self) -> Vec<&str> {
        self.projects.keys().map(|s| s.as_str()).collect()
    }
}

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

    #[test]
    fn test_query_context_default() {
        let ctx = QueryContext::default();
        assert_eq!(ctx.project, "default");
        assert_eq!(ctx.budget_bytes, 8192);
        assert_eq!(ctx.max_wiki_hops, 3);
    }

    #[test]
    fn test_memory_projection_register() {
        let mut proj = MemoryProjection::new();
        let vocab = Arc::new(BTreeMap::new());
        let scorer = Arc::new(mem_core::GlobalTfIdfScorer::new(vocab));
        let semantic = Arc::new(mem_core::SemanticScorer::new());

        let orchestrator = Arc::new(QueryOrchestrator::new(scorer, semantic, &QueryContext::default()));
        proj.register_project("test", orchestrator);

        assert!(proj.projects().contains(&"test"));
    }

    #[test]
    fn test_memory_projection_unknown_project() {
        let proj = MemoryProjection::new();
        let candidates = vec![("doc1".to_string(), "content".to_string())];
        let ctx = QueryContext::default();

        let result = tokio::runtime::Runtime::new()
            .unwrap()
            .block_on(proj.query_project("unknown", "test", candidates, &ctx));

        assert!(result.is_err());
    }

    #[test]
    fn test_optimized_chunk_creation() {
        let chunk = OptimizedChunk {
            id: "doc1".to_string(),
            text: "Test content".to_string(),
            tfidf_score: 0.5,
            semantic_score: 0.8,
            metadata_boost: 0.1,
            final_score: 0.9,
            category: crate::chunk_metadata::ChunkCategory::Solution,
            cache_distance: 2,
            cache_slot: 0,
        };

        assert_eq!(chunk.id, "doc1");
        assert_eq!(chunk.final_score, 0.9);
        assert!(chunk.final_score <= 1.0);
    }

    #[test]
    fn test_query_result_structure() {
        let result = QueryResult {
            query: "test".to_string(),
            selected_chunks: vec![],
            selection_metrics: SelectionMetrics {
                selected_count: 0,
                rejected_count: 0,
                total_bytes: 0,
                budget_used_pct: 0.0,
                avg_score: 0.0,
                dedup_removed: 0,
            },
            cache_metrics: crate::cache_alignment::CacheMetrics::new(),
            profiling: vec![],
            total_latency_ms: 100,
            query_intent: QueryIntent::Unknown,
        };

        assert_eq!(result.total_latency_ms, 100);
        assert_eq!(result.selected_chunks.len(), 0);
    }
}