poimen-memory/crates/mem-cli/src/handlers/unified_query.rs

//! Unified Query Handler (Phase 4.6)
//!
//! Single endpoint aggregating all search features:
//! - Semantic search (entities, edges, hybrid)
//! - Temporal filtering
//! - Community detection
//! - Path finding
//! - Faceted search

use actix_web::{web, HttpRequest, HttpResponse};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tracing::{debug, error, info};

use crate::http_server::AppState;
use crate::query::{
    SemanticRetriever, EntityResult, EdgeResult, HybridResult,
    CommunityDetector, CommunityDetectionResult,
    PathFinder, PathFindingResult,
    FacetedSearch, AvailableFacets, FacetFilters,
};

/// Unified query request (Phase 4.6)
///
/// Combines all search types and features into single endpoint.
/// Determines behavior via `search_type` parameter.
#[derive(Debug, Deserialize)]
pub struct UnifiedQueryRequest {
    /// Query text (will be embedded)
    pub query: String,

    // Search Type & Mode
    /// "entities" | "edges" | "hybrid" (default: "entities")
    #[serde(default = "default_search_type")]
    pub search_type: String,

    // Entity/Edge Filters
    /// Optional filter by entity type (entity search only)
    pub entity_type: Option<String>,
    /// Optional filter by relation type (edge search only)
    pub relation_type: Option<String>,

    // Scoring
    /// Minimum similarity (0.0-1.0, default 0.5)
    #[serde(default = "default_confidence_floor")]
    pub confidence_floor: f32,
    /// Semantic weight for hybrid (0.0-1.0, default 0.6)
    #[serde(default = "default_semantic_weight")]
    pub semantic_weight: f32,
    /// Lexical weight for hybrid (0.0-1.0, default 0.4)
    #[serde(default = "default_lexical_weight")]
    pub lexical_weight: f32,

    // Pagination
    /// Max results (default 10, max 100)
    #[serde(default = "default_top_k")]
    pub top_k: usize,

    // Temporal Filtering (Phase 4.2)
    /// Earliest event time (ISO 8601)
    pub start_time: Option<chrono::DateTime<chrono::Utc>>,
    /// Latest event time (ISO 8601)
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,

    // Community Detection (Phase 4.3)
    /// Enable community detection
    pub detect_communities: Option<bool>,
    /// Minimum community size (default 3)
    pub min_community_size: Option<usize>,

    // Path Finding (Phase 4.4)
    /// Enable path finding
    pub find_paths: Option<bool>,
    /// Target entity ID for paths
    pub target_entity_id: Option<String>,
    /// Max path depth (default 5, max 10)
    pub max_path_depth: Option<usize>,
    /// K-hop neighborhood size (default 2, max 5)
    pub k_hops: Option<usize>,

    // Faceted Search (Phase 4.5)
    /// Discover available facets
    pub discover_facets: Option<bool>,
    /// Apply facet filters
    pub facet_filters: Option<FacetFilters>,
}

/// Unified response wrapper
///
/// Serializes based on search_type and result content.
#[derive(Debug, Serialize)]
pub struct UnifiedQueryResponse {
    pub query: String,
    pub search_type: String,
    pub results: Vec<Value>,
    pub total_count: usize,
    pub search_time_ms: u128,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub communities: Option<CommunityDetectionResult>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub paths: Option<Vec<PathFindingResult>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub available_facets: Option<AvailableFacets>,
}

fn default_search_type() -> String { "entities".to_string() }
fn default_confidence_floor() -> f32 { 0.5 }
fn default_semantic_weight() -> f32 { 0.6 }
fn default_lexical_weight() -> f32 { 0.4 }
fn default_top_k() -> usize { 10 }

/// POST /memory/query - Unified query endpoint (Phase 4.6)
pub async fn unified_query_handler(
    req: HttpRequest,
    body: web::Json<UnifiedQueryRequest>,
    state: web::Data<AppState>,
) -> HttpResponse {
    let start_time = std::time::Instant::now();

    // 1. Validate JWT + rate limit
    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "query", 500
    ) {
        return response;
    }

    // 2. Validate input
    if let Err(response) = validate_unified_request(&body) {
        return response;
    }

    debug!("Unified query: type={}, query='{}', entity_type={:?}, relation_type={:?}",
           body.search_type, body.query, body.entity_type, body.relation_type);

    // 3. Embed query once (reused for all search types)
    let query_embedding = match state.embeddings.embed_one(&body.query).await {
        Ok(emb) => emb.to_vec(),
        Err(e) => {
            error!("Embedding failed: {}", e);
            return crate::handlers::response_builder::internal_error(
                "Failed to embed query"
            );
        }
    };

    // 4. Route to appropriate search type
    let response = match body.search_type.as_str() {
        "entities" => search_entities(&body, &state, &query_embedding, start_time).await,
        "edges" => search_edges(&body, &state, &query_embedding, start_time).await,
        "hybrid" => search_hybrid(&body, &state, &query_embedding, start_time).await,
        _ => {
            return crate::handlers::response_builder::bad_request(
                "search_type must be 'entities', 'edges', or 'hybrid'"
            );
        }
    };

    response
}

/// Search entities (with all optional features)
async fn search_entities(
    req: &UnifiedQueryRequest,
    state: &web::Data<AppState>,
    query_embedding: &[f32],
    start_time: std::time::Instant,
) -> HttpResponse {
    let retriever = SemanticRetriever::new(state.pool.clone());

    // Execute entity search
    let results = match retriever.search_entities(
        query_embedding,
        req.top_k,
        req.entity_type.as_deref(),
        req.confidence_floor,
        req.start_time,
        req.end_time,
    ).await {
        Ok(r) => r,
        Err(e) => {
            error!("Entity search failed: {}", e);
            return crate::handlers::response_builder::internal_error(&format!("Search failed: {}", e));
        }
    };

    let count = results.len();
    let elapsed = start_time.elapsed().as_millis();

    // Convert results to JSON
    let results_json: Vec<Value> = results.iter().map(|r| serde_json::to_value(r).unwrap_or(Value::Null)).collect();

    // Optional: Community detection
    let communities = if req.detect_communities.unwrap_or(false) {
        let detector = CommunityDetector::new(state.pool.clone());
        let min_size = req.min_community_size.unwrap_or(3);
        match detector.detect_communities(None, min_size, 0.001).await {
            Ok(result) => Some(result),
            Err(e) => {
                debug!("Community detection failed (non-fatal): {}", e);
                None
            }
        }
    } else {
        None
    };

    // Optional: Path finding
    let paths = if req.find_paths.unwrap_or(false) {
        if let (Some(first_result), Some(target_id)) = (results.first(), &req.target_entity_id) {
            let path_finder = PathFinder::new(state.pool.clone());
            let max_depth = req.max_path_depth.unwrap_or(5);

            match path_finder.shortest_path(&first_result.id, target_id, max_depth).await {
                Ok(Some(path)) => Some(vec![PathFindingResult {
                    source_id: first_result.id.clone(),
                    target_id: target_id.clone(),
                    paths_found: vec![path],
                    path_count: 1,
                    shortest_distance: Some(0),
                    average_distance: 0.0,
                }]),
                _ => None,
            }
        } else {
            None
        }
    } else {
        None
    };

    // Optional: Facet discovery
    let available_facets = if req.discover_facets.unwrap_or(false) {
        let faceted_search = FacetedSearch::new(state.pool.clone());
        match faceted_search.discover_facets("entities", 10).await {
            Ok(facets) => Some(facets),
            Err(e) => {
                debug!("Facet discovery failed (non-fatal): {}", e);
                None
            }
        }
    } else {
        None
    };

    info!("Unified query (entities): {} results in {}ms", count, elapsed);

    let response = UnifiedQueryResponse {
        query: req.query.clone(),
        search_type: "entities".to_string(),
        results: results_json,
        total_count: count,
        search_time_ms: elapsed,
        communities,
        paths,
        available_facets,
    };

    crate::handlers::response_builder::success_response(response)
}

/// Search edges (with temporal and facet filters)
async fn search_edges(
    req: &UnifiedQueryRequest,
    state: &web::Data<AppState>,
    query_embedding: &[f32],
    start_time: std::time::Instant,
) -> HttpResponse {
    let retriever = SemanticRetriever::new(state.pool.clone());

    let results = match retriever.search_edges(
        query_embedding,
        req.top_k,
        req.relation_type.as_deref(),
        req.start_time,
        req.end_time,
    ).await {
        Ok(r) => r,
        Err(e) => {
            error!("Edge search failed: {}", e);
            return crate::handlers::response_builder::internal_error(&format!("Search failed: {}", e));
        }
    };

    let count = results.len();
    let elapsed = start_time.elapsed().as_millis();

    let results_json: Vec<Value> = results.iter().map(|r| serde_json::to_value(r).unwrap_or(Value::Null)).collect();

    // Optional: Facet discovery
    let available_facets = if req.discover_facets.unwrap_or(false) {
        let faceted_search = FacetedSearch::new(state.pool.clone());
        match faceted_search.discover_facets("edges", 10).await {
            Ok(facets) => Some(facets),
            Err(e) => {
                debug!("Facet discovery failed (non-fatal): {}", e);
                None
            }
        }
    } else {
        None
    };

    info!("Unified query (edges): {} results in {}ms", count, elapsed);

    let response = UnifiedQueryResponse {
        query: req.query.clone(),
        search_type: "edges".to_string(),
        results: results_json,
        total_count: count,
        search_time_ms: elapsed,
        communities: None,
        paths: None,
        available_facets,
    };

    crate::handlers::response_builder::success_response(response)
}

/// Hybrid search (semantic + lexical with RRF)
async fn search_hybrid(
    req: &UnifiedQueryRequest,
    state: &web::Data<AppState>,
    query_embedding: &[f32],
    start_time: std::time::Instant,
) -> HttpResponse {
    let retriever = SemanticRetriever::new(state.pool.clone());

    let results = match retriever.hybrid_search(
        query_embedding,
        req.top_k,
        req.semantic_weight,
        req.lexical_weight,
        req.start_time,
        req.end_time,
    ).await {
        Ok(r) => r,
        Err(e) => {
            error!("Hybrid search failed: {}", e);
            return crate::handlers::response_builder::internal_error(&format!("Search failed: {}", e));
        }
    };

    let count = results.len();
    let elapsed = start_time.elapsed().as_millis();

    let results_json: Vec<Value> = results.iter().map(|r| serde_json::to_value(r).unwrap_or(Value::Null)).collect();

    info!("Unified query (hybrid): {} results in {}ms", count, elapsed);

    let response = UnifiedQueryResponse {
        query: req.query.clone(),
        search_type: "hybrid".to_string(),
        results: results_json,
        total_count: count,
        search_time_ms: elapsed,
        communities: None,
        paths: None,
        available_facets: None,
    };

    crate::handlers::response_builder::success_response(response)
}

/// Validate unified query request
fn validate_unified_request(req: &UnifiedQueryRequest) -> Result<(), HttpResponse> {
    // Query validation
    if req.query.is_empty() || req.query.len() > 2000 {
        return Err(crate::handlers::response_builder::bad_request(
            "Query must be 1-2000 characters"
        ));
    }

    // Search type validation
    if !matches!(req.search_type.as_str(), "entities" | "edges" | "hybrid") {
        return Err(crate::handlers::response_builder::bad_request(
            "search_type must be 'entities', 'edges', or 'hybrid'"
        ));
    }

    // Confidence floor validation
    if req.confidence_floor < 0.0 || req.confidence_floor > 1.0 {
        return Err(crate::handlers::response_builder::bad_request(
            "confidence_floor must be 0.0-1.0"
        ));
    }

    // Semantic/lexical weight validation (hybrid only)
    if req.semantic_weight < 0.0 || req.semantic_weight > 1.0 {
        return Err(crate::handlers::response_builder::bad_request(
            "semantic_weight must be 0.0-1.0"
        ));
    }

    if req.lexical_weight < 0.0 || req.lexical_weight > 1.0 {
        return Err(crate::handlers::response_builder::bad_request(
            "lexical_weight must be 0.0-1.0"
        ));
    }

    // Top K validation
    if req.top_k == 0 || req.top_k > 100 {
        return Err(crate::handlers::response_builder::bad_request(
            "top_k must be 1-100"
        ));
    }

    // Temporal validation
    if let (Some(start), Some(end)) = (req.start_time, req.end_time) {
        if start > end {
            return Err(crate::handlers::response_builder::bad_request(
                "start_time must be <= end_time"
            ));
        }
    }

    // Max path depth validation
    if let Some(depth) = req.max_path_depth {
        if depth == 0 || depth > 10 {
            return Err(crate::handlers::response_builder::bad_request(
                "max_path_depth must be 1-10"
            ));
        }
    }

    // K hops validation
    if let Some(hops) = req.k_hops {
        if hops == 0 || hops > 5 {
            return Err(crate::handlers::response_builder::bad_request(
                "k_hops must be 1-5"
            ));
        }
    }

    // Min community size validation
    if let Some(size) = req.min_community_size {
        if size < 2 || size > 1000 {
            return Err(crate::handlers::response_builder::bad_request(
                "min_community_size must be 2-1000"
            ));
        }
    }

    Ok(())
}

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

    #[test]
    fn test_unified_query_default_search_type() {
        let req = UnifiedQueryRequest {
            query: "test".to_string(),
            search_type: default_search_type(),
            entity_type: None,
            relation_type: None,
            confidence_floor: default_confidence_floor(),
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: default_top_k(),
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert_eq!(req.search_type, "entities");
    }

    #[test]
    fn test_unified_query_entity_search() {
        let req = UnifiedQueryRequest {
            query: "kubernetes".to_string(),
            search_type: "entities".to_string(),
            entity_type: Some("concept".to_string()),
            relation_type: None,
            confidence_floor: 0.7,
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: 20,
            start_time: None,
            end_time: None,
            detect_communities: Some(true),
            min_community_size: Some(3),
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert_eq!(req.search_type, "entities");
        assert_eq!(req.entity_type, Some("concept".to_string()));
        assert_eq!(req.detect_communities, Some(true));
    }

    #[test]
    fn test_unified_query_edge_search() {
        let req = UnifiedQueryRequest {
            query: "depends on".to_string(),
            search_type: "edges".to_string(),
            entity_type: None,
            relation_type: Some("depends_on".to_string()),
            confidence_floor: default_confidence_floor(),
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: 10,
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert_eq!(req.search_type, "edges");
        assert_eq!(req.relation_type, Some("depends_on".to_string()));
    }

    #[test]
    fn test_unified_query_hybrid_search() {
        let req = UnifiedQueryRequest {
            query: "system design".to_string(),
            search_type: "hybrid".to_string(),
            entity_type: None,
            relation_type: None,
            confidence_floor: default_confidence_floor(),
            semantic_weight: 0.7,
            lexical_weight: 0.3,
            top_k: 15,
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert_eq!(req.search_type, "hybrid");
        assert_eq!(req.semantic_weight, 0.7);
        assert_eq!(req.lexical_weight, 0.3);
    }

    #[test]
    fn test_unified_query_with_all_features() {
        let req = UnifiedQueryRequest {
            query: "kubernetes infrastructure".to_string(),
            search_type: "entities".to_string(),
            entity_type: Some("technology".to_string()),
            relation_type: None,
            confidence_floor: 0.7,
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: 20,
            start_time: None,
            end_time: None,
            detect_communities: Some(true),
            min_community_size: Some(5),
            find_paths: Some(true),
            target_entity_id: Some("e_monitoring".to_string()),
            max_path_depth: Some(4),
            k_hops: Some(3),
            discover_facets: Some(true),
            facet_filters: Some(FacetFilters {
                entity_types: Some(vec!["concept".to_string()]),
                relation_types: None,
                confidence_level: Some("high".to_string()),
                date_range: Some("this_month".to_string()),
            }),
        };
        assert_eq!(req.search_type, "entities");
        assert!(req.detect_communities.unwrap_or(false));
        assert!(req.find_paths.unwrap_or(false));
        assert!(req.discover_facets.unwrap_or(false));
    }

    #[test]
    fn test_unified_query_response() {
        let response = UnifiedQueryResponse {
            query: "test".to_string(),
            search_type: "entities".to_string(),
            results: vec![],
            total_count: 0,
            search_time_ms: 100,
            communities: None,
            paths: None,
            available_facets: None,
        };
        assert_eq!(response.query, "test");
        assert_eq!(response.search_type, "entities");
        assert_eq!(response.total_count, 0);
    }

    #[test]
    fn test_validate_unified_query_invalid_query() {
        let req = UnifiedQueryRequest {
            query: "".to_string(),
            search_type: "entities".to_string(),
            entity_type: None,
            relation_type: None,
            confidence_floor: default_confidence_floor(),
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: default_top_k(),
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert!(validate_unified_request(&req).is_err());
    }

    #[test]
    fn test_validate_unified_query_invalid_search_type() {
        let req = UnifiedQueryRequest {
            query: "test".to_string(),
            search_type: "invalid".to_string(),
            entity_type: None,
            relation_type: None,
            confidence_floor: default_confidence_floor(),
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: default_top_k(),
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert!(validate_unified_request(&req).is_err());
    }

    #[test]
    fn test_validate_unified_query_invalid_confidence_floor() {
        let req = UnifiedQueryRequest {
            query: "test".to_string(),
            search_type: "entities".to_string(),
            entity_type: None,
            relation_type: None,
            confidence_floor: 1.5,
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: default_top_k(),
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert!(validate_unified_request(&req).is_err());
    }

    #[test]
    fn test_validate_unified_query_invalid_top_k() {
        let req = UnifiedQueryRequest {
            query: "test".to_string(),
            search_type: "entities".to_string(),
            entity_type: None,
            relation_type: None,
            confidence_floor: default_confidence_floor(),
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: 200,
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert!(validate_unified_request(&req).is_err());
    }

    #[test]
    fn test_validate_unified_query_valid() {
        let req = UnifiedQueryRequest {
            query: "test".to_string(),
            search_type: "entities".to_string(),
            entity_type: None,
            relation_type: None,
            confidence_floor: 0.5,
            semantic_weight: 0.6,
            lexical_weight: 0.4,
            top_k: 20,
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: None,
            target_entity_id: None,
            max_path_depth: None,
            k_hops: None,
            discover_facets: None,
            facet_filters: None,
        };
        assert!(validate_unified_request(&req).is_ok());
    }
}