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

//! Semantic Search Handler
//!
//! HTTP endpoint for semantic retrieval (vector search).

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

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

/// Request for semantic entity search
#[derive(Debug, Deserialize)]
pub struct SemanticSearchEntityRequest {
    /// Query text (will be embedded)
    pub query: String,
    /// Optional entity type filter
    pub entity_type: Option<String>,
    /// Minimum similarity score (0.0-1.0, default 0.5)
    #[serde(default = "default_confidence_floor")]
    pub confidence_floor: f32,
    /// Maximum number of results (default 10)
    #[serde(default = "default_top_k")]
    pub top_k: usize,
    /// Optional: minimum event_time (ISO 8601)
    pub start_time: Option<chrono::DateTime<chrono::Utc>>,
    /// Optional: maximum event_time (ISO 8601)
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
    /// Optional: include community detection in results
    pub detect_communities: Option<bool>,
    /// Optional: minimum community size (default 3, min 2)
    pub min_community_size: Option<usize>,
    /// Optional: find paths from query result to target entity
    pub find_paths: Option<bool>,
    /// Optional: target entity ID for path finding
    pub target_entity_id: Option<String>,
    /// Optional: maximum hops for path finding (default 5, max 10)
    pub max_path_depth: Option<usize>,
    /// Optional: find k-hop neighborhood around result
    pub k_hops: Option<usize>,
    /// Optional: apply facet filters
    pub facet_filters: Option<FacetFilters>,
    /// Optional: discover available facets
    pub discover_facets: Option<bool>,
}

/// Request for semantic edge search
#[derive(Debug, Deserialize)]
pub struct SemanticSearchEdgeRequest {
    /// Query text (will be embedded)
    pub query: String,
    /// Optional relation type filter
    pub relation_type: Option<String>,
    /// Maximum number of results (default 10)
    #[serde(default = "default_top_k")]
    pub top_k: usize,
    /// Optional: minimum event_time (ISO 8601)
    pub start_time: Option<chrono::DateTime<chrono::Utc>>,
    /// Optional: maximum event_time (ISO 8601)
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
}

/// Request for hybrid search
#[derive(Debug, Deserialize)]
pub struct HybridSearchRequest {
    /// Query text (will be embedded)
    pub query: String,
    /// Weight for semantic score (default 0.6)
    #[serde(default = "default_semantic_weight")]
    pub semantic_weight: f32,
    /// Weight for lexical score (default 0.4)
    #[serde(default = "default_lexical_weight")]
    pub lexical_weight: f32,
    /// Maximum number of results (default 10)
    #[serde(default = "default_top_k")]
    pub top_k: usize,
    /// Optional: minimum event_time (ISO 8601)
    pub start_time: Option<chrono::DateTime<chrono::Utc>>,
    /// Optional: maximum event_time (ISO 8601)
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
}

/// Response for semantic search (with optional community detection, path finding, and facets)
#[derive(Debug, Serialize)]
pub struct SemanticSearchResponse<T> {
    pub query: String,
    pub results: Vec<T>,
    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_confidence_floor() -> f32 { 0.5 }
fn default_top_k() -> usize { 10 }
fn default_semantic_weight() -> f32 { 0.6 }
fn default_lexical_weight() -> f32 { 0.4 }

/// POST /memory/query/semantic/entities - Search entities by semantic similarity
pub async fn search_entities_handler(
    req: HttpRequest,
    body: web::Json<SemanticSearchEntityRequest>,
    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, "semantic_search", 500
    ) {
        return response;
    }

    // 2. Validate input
    if body.query.is_empty() || body.query.len() > 2000 {
        return crate::handlers::response_builder::bad_request(
            "Query must be 1-2000 characters"
        );
    }

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

    // Validate temporal parameters (if provided)
    if let (Some(start), Some(end)) = (body.start_time, body.end_time) {
        if start > end {
            return crate::handlers::response_builder::bad_request(
                "start_time must be <= end_time"
            );
        }
    }

    debug!("Semantic search entities: query='{}', entity_type={:?}, temporal={:?}-{:?}", 
           body.query, body.entity_type, body.start_time, body.end_time);

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

    // 4. Execute search with temporal filtering
    let retriever = SemanticRetriever::new(state.pool.clone());
    match retriever.search_entities(
        &query_embedding,
        body.top_k,
        body.entity_type.as_deref(),
        body.confidence_floor,
        body.start_time,
        body.end_time,
    ).await {
        Ok(results) => {
            let count = results.len();
            let elapsed = start_time.elapsed().as_millis();

            // 5. Optional: detect communities
            let communities = if body.detect_communities.unwrap_or(false) {
                let detector = CommunityDetector::new(state.pool.clone());
                let min_size = body.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
            };

            // 6. Optional: find paths from first result to target
            let paths = if body.find_paths.unwrap_or(false) {
                if let (Some(first_result), Some(target_id)) = (results.first(), &body.target_entity_id) {
                    let path_finder = PathFinder::new(state.pool.clone());
                    let max_depth = body.max_path_depth.unwrap_or(5);
                    
                    // Find shortest path
                    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
            };

            // 7. Optional: discover available facets
            let available_facets = if body.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!("Semantic entity search completed: {} results in {}ms", count, elapsed);

            let response = SemanticSearchResponse {
                query: body.query.clone(),
                results,
                total_count: count,
                search_time_ms: elapsed,
                communities,
                paths,
                available_facets,
            };

            crate::handlers::response_builder::success_response(response)
        }
        Err(e) => {
            error!("Semantic search failed: {}", e);
            crate::handlers::response_builder::internal_error(
                &format!("Search failed: {}", e)
            )
        }
    }
}

/// POST /memory/query/semantic/edges - Search edges by semantic similarity
pub async fn search_edges_handler(
    req: HttpRequest,
    body: web::Json<SemanticSearchEdgeRequest>,
    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, "semantic_search", 500
    ) {
        return response;
    }

    // 2. Validate input
    if body.query.is_empty() || body.query.len() > 2000 {
        return crate::handlers::response_builder::bad_request(
            "Query must be 1-2000 characters"
        );
    }

    // Validate temporal parameters (if provided)
    if let (Some(start), Some(end)) = (body.start_time, body.end_time) {
        if start > end {
            return crate::handlers::response_builder::bad_request(
                "start_time must be <= end_time"
            );
        }
    }

    debug!("Semantic search edges: query='{}', relation_type={:?}, temporal={:?}-{:?}",
           body.query, body.relation_type, body.start_time, body.end_time);

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

    // 4. Execute search with temporal filtering
    let retriever = SemanticRetriever::new(state.pool.clone());
    match retriever.search_edges(
        &query_embedding,
        body.top_k,
        body.relation_type.as_deref(),
        body.start_time,
        body.end_time,
    ).await {
        Ok(results) => {
            let count = results.len();
            let elapsed = start_time.elapsed().as_millis();
            info!("Semantic edge search completed: {} results in {}ms", count, elapsed);

            let response = SemanticSearchResponse {
                query: body.query.clone(),
                results,
                total_count: count,
                search_time_ms: elapsed,
                communities: None,
                paths: None,
                available_facets: None,
            };

            crate::handlers::response_builder::success_response(response)
        }
        Err(e) => {
            error!("Semantic search failed: {}", e);
            crate::handlers::response_builder::internal_error(
                &format!("Search failed: {}", e)
            )
        }
    }
}

/// POST /memory/query/hybrid - Hybrid semantic + lexical search
pub async fn hybrid_search_handler(
    req: HttpRequest,
    body: web::Json<HybridSearchRequest>,
    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, "semantic_search", 500
    ) {
        return response;
    }

    // 2. Validate input
    if body.query.is_empty() || body.query.len() > 2000 {
        return crate::handlers::response_builder::bad_request(
            "Query must be 1-2000 characters"
        );
    }

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

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

    debug!("Hybrid search: query='{}', weights=(sem={}, lex={})",
           body.query, body.semantic_weight, body.lexical_weight);

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

    // 4. Execute search with temporal filtering
    let retriever = SemanticRetriever::new(state.pool.clone());
    match retriever.hybrid_search(
        &query_embedding,
        body.top_k,
        body.semantic_weight,
        body.lexical_weight,
        body.start_time,
        body.end_time,
    ).await {
        Ok(results) => {
            let count = results.len();
            let elapsed = start_time.elapsed().as_millis();
            info!("Hybrid search completed: {} results in {}ms", count, elapsed);

            let response = SemanticSearchResponse {
                query: body.query.clone(),
                results,
                total_count: count,
                search_time_ms: elapsed,
                communities: None,
                paths: None,
                available_facets: None,
            };

            crate::handlers::response_builder::success_response(response)
        }
        Err(e) => {
            error!("Hybrid search failed: {}", e);
            crate::handlers::response_builder::internal_error(
                &format!("Search failed: {}", e)
            )
        }
    }
}

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

    #[test]
    fn test_semantic_search_entity_request() {
        let req = SemanticSearchEntityRequest {
            query: "test query".to_string(),
            entity_type: Some("concept".to_string()),
            confidence_floor: 0.5,
            top_k: 10,
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
        };
        assert_eq!(req.query, "test query");
        assert_eq!(req.confidence_floor, 0.5);
    }

    #[test]
    fn test_semantic_search_with_temporal_range() {
        use chrono::{Utc, Duration};
        let now = Utc::now();
        let tomorrow = now + Duration::days(1);

        let req = SemanticSearchEntityRequest {
            query: "test query".to_string(),
            entity_type: None,
            confidence_floor: 0.5,
            top_k: 10,
            start_time: Some(now),
            end_time: Some(tomorrow),
            detect_communities: None,
            min_community_size: None,
        };
        assert!(req.start_time <= req.end_time);
    }

    #[test]
    fn test_semantic_search_with_community_detection() {
        let req = SemanticSearchEntityRequest {
            query: "test query".to_string(),
            entity_type: None,
            confidence_floor: 0.5,
            top_k: 10,
            start_time: None,
            end_time: None,
            detect_communities: Some(true),
            min_community_size: Some(3),
        };
        assert_eq!(req.detect_communities, Some(true));
        assert_eq!(req.min_community_size, Some(3));
    }

    #[test]
    fn test_semantic_search_edge_request() {
        let req = SemanticSearchEdgeRequest {
            query: "test query".to_string(),
            relation_type: Some("related_to".to_string()),
            top_k: 10,
            start_time: None,
            end_time: None,
        };
        assert_eq!(req.query, "test query");
    }

    #[test]
    fn test_hybrid_search_request_defaults() {
        let req = HybridSearchRequest {
            query: "test".to_string(),
            semantic_weight: default_semantic_weight(),
            lexical_weight: default_lexical_weight(),
            top_k: default_top_k(),
        };
        assert_eq!(req.semantic_weight, 0.6);
        assert_eq!(req.lexical_weight, 0.4);
        assert_eq!(req.top_k, 10);
    }

    #[test]
    fn test_semantic_search_response() {
        let response: SemanticSearchResponse<EntityResult> = SemanticSearchResponse {
            query: "test".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.total_count, 0);
    }

    #[test]
    fn test_semantic_search_with_path_finding() {
        let req = SemanticSearchEntityRequest {
            query: "test query".to_string(),
            entity_type: None,
            confidence_floor: 0.5,
            top_k: 10,
            start_time: None,
            end_time: None,
            detect_communities: None,
            min_community_size: None,
            find_paths: Some(true),
            target_entity_id: Some("e5".to_string()),
            max_path_depth: Some(5),
            k_hops: None,
            facet_filters: None,
            discover_facets: None,
        };
        assert_eq!(req.find_paths, Some(true));
        assert_eq!(req.target_entity_id, Some("e5".to_string()));
    }

    #[test]
    fn test_semantic_search_with_facet_discovery() {
        let req = SemanticSearchEntityRequest {
            query: "kubernetes".to_string(),
            entity_type: None,
            confidence_floor: 0.5,
            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,
            facet_filters: None,
            discover_facets: Some(true),
        };
        assert_eq!(req.discover_facets, Some(true));
    }

    #[test]
    fn test_semantic_search_with_facet_filters() {
        let filters = FacetFilters {
            entity_types: Some(vec!["concept".to_string()]),
            relation_types: None,
            confidence_level: Some("high".to_string()),
            date_range: None,
        };
        let req = SemanticSearchEntityRequest {
            query: "test".to_string(),
            entity_type: None,
            confidence_floor: 0.5,
            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,
            facet_filters: Some(filters),
            discover_facets: None,
        };
        assert!(req.facet_filters.is_some());
        assert_eq!(req.facet_filters.unwrap().confidence_level, Some("high".to_string()));
    }
}