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

/// Query Router: Unified Phase 3+4 pipeline
///
/// Bridges wiki-link graph (Phase 1) with hybrid retrieval (Phase 3)
/// and LLM optimization (Phase 4) into a single query flow.
///
/// Pipeline:
/// 1. Wiki-scope filtering (via WikiLinkGraph)
/// 2. TF-IDF pre-filtering
/// 3. Semantic re-ranking
/// 4. RRF fusion
/// 5. Score thresholding + budget selection + deduplication

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

use mem_ingest::wiki_link::{WikiLinkGraph, WikiLinkParser};
use mem_core::{DocumentScorer, GlobalTfIdfScorer, SemanticScorer};

use crate::hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate};
use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};

/// Query routing configuration
#[derive(Debug, Clone)]
pub struct RouterConfig {
    pub max_wiki_hops: u32,
    pub tfidf_threshold: f32,
    pub prefilter_limit: usize,
    pub score_threshold: f32,
    pub budget_bytes: usize,
    pub dedup_threshold: f32,
    pub rrf_tfidf_weight: f32,
    pub rrf_semantic_weight: f32,
}

impl Default for RouterConfig {
    fn default() -> Self {
        Self {
            max_wiki_hops: 3,
            tfidf_threshold: 0.3,
            prefilter_limit: 50,
            score_threshold: 0.6,
            budget_bytes: 8192,
            dedup_threshold: 0.8,
            rrf_tfidf_weight: 0.4,
            rrf_semantic_weight: 0.6,
        }
    }
}

/// Query routing result with full metrics
#[derive(Debug, Clone)]
pub struct RoutedResult {
    pub selected_chunks: Vec<SelectedChunk>,
    pub route: RetrievalRoute,
    pub wiki_scope_size: usize,
    pub prefilter_size: usize,
    pub metrics: SelectionMetrics,
    pub latency_ms: u64,
}

/// Selected chunk with all scores
#[derive(Debug, Clone)]
pub struct SelectedChunk {
    pub id: String,
    pub text: String,
    pub tfidf_score: f32,
    pub semantic_score: f32,
    pub final_score: f32,
    pub wiki_distance: Option<u32>,
}

/// Query Router: end-to-end Phase 3+4 pipeline
pub struct QueryRouter {
    wiki_filter: WikiScopedFilter,
    retriever: HybridRetriever,
    optimizer: ChunkOptimizer,
    config: RouterConfig,
}

impl QueryRouter {
    pub fn new(
        tfidf_scorer: Arc<GlobalTfIdfScorer>,
        semantic_scorer: Arc<SemanticScorer>,
        config: RouterConfig,
    ) -> Self {
        let wiki_filter = WikiScopedFilter::new(config.max_wiki_hops);
        let retriever = HybridRetriever::new(tfidf_scorer, semantic_scorer);
        let optimizer = ChunkOptimizer::new(
            config.score_threshold,
            config.budget_bytes,
            config.dedup_threshold,
        );

        Self {
            wiki_filter,
            retriever,
            optimizer,
            config,
        }
    }

    /// Execute full query pipeline with wiki-link graph scoping
    pub async fn route_with_wiki_graph(
        &self,
        query: &str,
        wiki_graph: &WikiLinkGraph,
        root_doc: &str,
        all_candidates: Vec<(String, String)>, // (doc_id, text)
    ) -> Result<RoutedResult> {
        let start = std::time::Instant::now();

        // Phase 1: Wiki-scope reduction
        let wiki_reachable = wiki_graph.reachable_docs(root_doc);
        let wiki_scope_size = wiki_reachable.len();

        // Convert wiki-graph to HashMap for WikiScopedFilter
        let graph_map = self.wiki_graph_to_hashmap(wiki_graph, root_doc);

        // Filter candidates by wiki scope
        let scoped_candidates: Vec<_> = all_candidates
            .into_iter()
            .filter(|(doc_id, _)| wiki_reachable.contains(doc_id))
            .collect();

        // Phase 3: Hybrid retrieval
        let route = self.retriever.route_query(query, !wiki_reachable.is_empty(), false);
        let ranked = self.retriever.retrieve(query, scoped_candidates, route.clone()).await?;
        let prefilter_size = ranked.len();

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

        // Phase 4: LLM optimization (threshold + budget + dedup)
        let (selected_opt, metrics) = self.optimizer.optimize(optimizable);

        // Build final result with wiki distances
        let selected_chunks: Vec<SelectedChunk> = selected_opt
            .into_iter()
            .map(|chunk| {
                let wiki_distance = self.calculate_wiki_distance(&chunk.id, root_doc, &graph_map);
                SelectedChunk {
                    id: chunk.id,
                    text: chunk.text,
                    tfidf_score: chunk.score * self.config.rrf_tfidf_weight,
                    semantic_score: chunk.score * self.config.rrf_semantic_weight,
                    final_score: chunk.score,
                    wiki_distance,
                }
            })
            .collect();

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

        Ok(RoutedResult {
            selected_chunks,
            route,
            wiki_scope_size,
            prefilter_size,
            metrics,
            latency_ms,
        })
    }

    /// Execute query without wiki-graph (direct retrieval)
    pub async fn route_direct(
        &self,
        query: &str,
        all_candidates: Vec<(String, String)>,
    ) -> Result<RoutedResult> {
        let start = std::time::Instant::now();

        // Direct retrieval (no wiki scoping)
        let route = RetrievalRoute::Direct;
        let ranked = self.retriever.retrieve(query, all_candidates.clone(), route.clone()).await?;
        let prefilter_size = ranked.len();

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

        // Phase 4: LLM optimization
        let (selected_opt, metrics) = self.optimizer.optimize(optimizable);

        let selected_chunks: Vec<SelectedChunk> = selected_opt
            .into_iter()
            .map(|chunk| SelectedChunk {
                id: chunk.id,
                text: chunk.text,
                tfidf_score: chunk.score * self.config.rrf_tfidf_weight,
                semantic_score: chunk.score * self.config.rrf_semantic_weight,
                final_score: chunk.score,
                wiki_distance: None,
            })
            .collect();

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

        Ok(RoutedResult {
            selected_chunks,
            route,
            wiki_scope_size: all_candidates.len(),
            prefilter_size,
            metrics,
            latency_ms,
        })
    }

    /// Convert WikiLinkGraph to HashMap for distance calculation
    fn wiki_graph_to_hashmap(
        &self,
        wiki_graph: &WikiLinkGraph,
        root_doc: &str,
    ) -> HashMap<String, Vec<String>> {
        let reachable = wiki_graph.reachable_docs(root_doc);
        let mut graph_map = HashMap::new();

        for doc in &reachable {
            let forward = wiki_graph.forward_links(doc);
            graph_map.insert(doc.clone(), forward);
        }

        graph_map
    }

    /// Calculate wiki distance using BFS
    fn calculate_wiki_distance(
        &self,
        doc_id: &str,
        root_doc: &str,
        graph: &HashMap<String, Vec<String>>,
    ) -> Option<u32> {
        if doc_id == root_doc {
            return Some(0);
        }

        let mut visited = std::collections::HashSet::new();
        let mut queue = std::collections::VecDeque::new();

        queue.push_back((root_doc.to_string(), 0u32));
        visited.insert(root_doc.to_string());

        while let Some((current, distance)) = queue.pop_front() {
            if current == doc_id {
                return Some(distance);
            }

            if distance >= self.config.max_wiki_hops {
                continue;
            }

            if let Some(neighbors) = graph.get(&current) {
                for neighbor in neighbors {
                    if !visited.contains(neighbor) {
                        visited.insert(neighbor.clone());
                        queue.push_back((neighbor.clone(), distance + 1));
                    }
                }
            }
        }

        None // Not reachable
    }

    pub fn config(&self) -> &RouterConfig {
        &self.config
    }
}

/// Build wiki-link graph from markdown content
pub struct WikiGraphBuilder;

impl WikiGraphBuilder {
    /// Build graph from list of (doc_id, content) pairs
    pub fn build_from_docs(
        project: &str,
        docs: Vec<(&str, &str)>,
    ) -> Result<WikiLinkGraph> {
        let mut graph = WikiLinkGraph::new(project);

        for (doc_id, content) in docs {
            let links = WikiLinkParser::parse_links(content)?;
            for target in links {
                graph.add_link(doc_id, &target);
            }
        }

        Ok(graph)
    }
}

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

    fn create_test_router() -> QueryRouter {
        let vocab = Arc::new(BTreeMap::new());
        let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab));
        let semantic = Arc::new(SemanticScorer::new());

        QueryRouter::new(tfidf, semantic, RouterConfig::default())
    }

    fn create_test_wiki_graph() -> WikiLinkGraph {
        let mut graph = WikiLinkGraph::new("test");
        graph.add_link("index.md", "tools/kubectl.md");
        graph.add_link("tools/kubectl.md", "debugging/pod-crashes.md");
        graph.add_link("debugging/pod-crashes.md", "solutions/restart-pod.md");
        graph
    }

    #[test]
    fn test_router_config_default() {
        let config = RouterConfig::default();
        assert_eq!(config.max_wiki_hops, 3);
        assert_eq!(config.score_threshold, 0.6);
        assert_eq!(config.budget_bytes, 8192);
    }

    #[test]
    fn test_wiki_graph_to_hashmap() {
        let router = create_test_router();
        let graph = create_test_wiki_graph();

        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");

        assert!(hashmap.contains_key("index.md"));
        assert!(hashmap.contains_key("tools/kubectl.md"));
        assert!(hashmap.contains_key("debugging/pod-crashes.md"));
    }

    #[test]
    fn test_calculate_wiki_distance_root() {
        let router = create_test_router();
        let graph = create_test_wiki_graph();
        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");

        let distance = router.calculate_wiki_distance("index.md", "index.md", &hashmap);
        assert_eq!(distance, Some(0));
    }

    #[test]
    fn test_calculate_wiki_distance_direct_child() {
        let router = create_test_router();
        let graph = create_test_wiki_graph();
        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");

        let distance = router.calculate_wiki_distance("tools/kubectl.md", "index.md", &hashmap);
        assert_eq!(distance, Some(1));
    }

    #[test]
    fn test_calculate_wiki_distance_grandchild() {
        let router = create_test_router();
        let graph = create_test_wiki_graph();
        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");

        let distance = router.calculate_wiki_distance("debugging/pod-crashes.md", "index.md", &hashmap);
        assert_eq!(distance, Some(2));
    }

    #[test]
    fn test_calculate_wiki_distance_unreachable() {
        let router = create_test_router();
        let graph = create_test_wiki_graph();
        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");

        let distance = router.calculate_wiki_distance("unknown.md", "index.md", &hashmap);
        assert_eq!(distance, None);
    }

    #[tokio::test]
    async fn test_route_direct() {
        let router = create_test_router();
        let candidates = vec![
            ("doc1".to_string(), "kubernetes pod debugging".to_string()),
            ("doc2".to_string(), "docker container deployment".to_string()),
        ];

        let result = router.route_direct("kubernetes", candidates).await.unwrap();

        assert_eq!(result.route, RetrievalRoute::Direct);
        assert!(result.latency_ms >= 0);
    }

    #[tokio::test]
    async fn test_route_with_wiki_graph() {
        let router = create_test_router();
        let graph = create_test_wiki_graph();

        let candidates = vec![
            ("index.md".to_string(), "main index".to_string()),
            ("tools/kubectl.md".to_string(), "kubectl tool".to_string()),
            ("debugging/pod-crashes.md".to_string(), "debugging content".to_string()),
            ("unrelated.md".to_string(), "not in graph".to_string()),
        ];

        let result = router
            .route_with_wiki_graph("kubectl", &graph, "index.md", candidates)
            .await
            .unwrap();

        // Should filter out "unrelated.md" (not reachable from index.md)
        assert!(result.wiki_scope_size <= 4);
        assert_eq!(result.route, RetrievalRoute::WikiScoped);
    }

    #[test]
    fn test_wiki_graph_builder() {
        let docs = vec![
            ("index.md", "# Index\nSee [[tools/kubectl.md]] for tools."),
            ("tools/kubectl.md", "# Kubectl\nSee [[debugging.md]] for debugging."),
        ];

        let graph = WikiGraphBuilder::build_from_docs("test", docs).unwrap();

        let reachable = graph.reachable_docs("index.md");
        assert!(reachable.contains("index.md"));
        assert!(reachable.contains("tools/kubectl.md"));
        assert!(reachable.contains("debugging.md"));
    }

    #[test]
    fn test_selected_chunk_structure() {
        let chunk = SelectedChunk {
            id: "doc1".to_string(),
            text: "content".to_string(),
            tfidf_score: 0.4,
            semantic_score: 0.6,
            final_score: 0.9,
            wiki_distance: Some(1),
        };

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

    #[test]
    fn test_routed_result_structure() {
        let result = RoutedResult {
            selected_chunks: vec![],
            route: RetrievalRoute::WikiScoped,
            wiki_scope_size: 10,
            prefilter_size: 5,
            metrics: SelectionMetrics {
                selected_count: 3,
                rejected_count: 2,
                total_bytes: 1000,
                budget_used_pct: 12.5,
                avg_score: 0.8,
                dedup_removed: 0,
            },
            latency_ms: 50,
        };

        assert_eq!(result.wiki_scope_size, 10);
        assert_eq!(result.prefilter_size, 5);
        assert_eq!(result.metrics.selected_count, 3);
    }
}