poimen-memory/crates/mem-cli/src/query/faceted_search.rs

//! Faceted Search Engine
//!
//! Enables multi-dimensional filtering across entities and edges.
//! Supports entity types, relation types, date ranges, confidence levels, and more.

use chrono::{DateTime, Timelike, Utc};
use serde::{Deserialize, Serialize};
use sqlx::{Pool, Postgres};
use std::collections::HashMap;
use tracing::{debug, info};

/// A single facet (filterable dimension)
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum FacetType {
    /// Entity type (e.g., "concept", "person", "technology")
    EntityType,
    /// Relation type (e.g., "depends_on", "related", "inherits")
    RelationType,
    /// Confidence level (e.g., "high", "medium", "low")
    ConfidenceLevel,
    /// Date range (e.g., "today", "this_week", "this_month")
    DateRange,
}

/// A facet value with count
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FacetValue {
    pub name: String,               // e.g., "concept", "high"
    pub count: usize,               // How many results match this value
    pub percentage: f32,            // Percentage of total results (0-100)
}

/// Available facets for a query
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AvailableFacets {
    pub entity_types: Vec<FacetValue>,
    pub relation_types: Vec<FacetValue>,
    pub confidence_levels: Vec<FacetValue>,
    pub date_ranges: Vec<FacetValue>,
    pub total_results: usize,
    pub facet_time_ms: u128,
}

/// Facet filters for a query
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct FacetFilters {
    /// Filter by entity types (OR within facet, AND across facets)
    pub entity_types: Option<Vec<String>>,
    /// Filter by relation types
    pub relation_types: Option<Vec<String>>,
    /// Filter by confidence level ("high"=0.8+, "medium"=0.5-0.8, "low"=<0.5)
    pub confidence_level: Option<String>,
    /// Filter by date range ("today", "week", "month", "year", "all")
    pub date_range: Option<String>,
}

/// Faceted search result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FacetedResult<T> {
    pub results: Vec<T>,
    pub total_count: usize,
    pub available_facets: AvailableFacets,
    pub applied_filters: FacetFilters,
}

/// Faceted Search Engine
pub struct FacetedSearch {
    pub pool: Pool<Postgres>,
}

impl FacetedSearch {
    /// Create a new faceted search engine
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }

    /// Discover available facets for a query
    ///
    /// # Arguments
    /// * `search_type` - "entities" or "edges"
    /// * `limit` - Maximum facet values per facet type (default 10, max 50)
    ///
    /// # Returns
    /// AvailableFacets with all discoverable filters
    pub async fn discover_facets(
        &self,
        search_type: &str,
        limit: usize,
    ) -> Result<AvailableFacets, String> {
        let limit = limit.max(5).min(50);
        let start_time = std::time::Instant::now();

        debug!("Discovering facets for {}, limit={}", search_type, limit);

        if search_type == "entities" {
            self.discover_entity_facets(limit).await
        } else if search_type == "edges" {
            self.discover_edge_facets(limit).await
        } else {
            Err(format!("Unknown search type: {}", search_type))
        }
    }

    /// Discover facets for entity searches
    async fn discover_entity_facets(&self, limit: usize) -> Result<AvailableFacets, String> {
        let start_time = std::time::Instant::now();

        // Get entity types
        let entity_types = sqlx::query_as::<_, (String, i64)>(
            "SELECT entity_type, COUNT(*) as cnt 
             FROM memory_entity 
             WHERE deleted_at IS NULL 
             GROUP BY entity_type 
             ORDER BY cnt DESC 
             LIMIT $1"
        )
            .bind(limit as i64)
            .fetch_all(&self.pool)
            .await
            .map_err(|e| format!("Failed to fetch entity types: {}", e))?
            .into_iter()
            .map(|(name, count)| FacetValue {
                name,
                count: count as usize,
                percentage: 0.0,  // Will be set later
            })
            .collect::<Vec<_>>();

        // Get total count
        let total_count: (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM memory_entity WHERE deleted_at IS NULL"
        )
            .fetch_one(&self.pool)
            .await
            .map_err(|e| format!("Failed to get total count: {}", e))?;

        let total = total_count.0 as usize;

        // Calculate percentages
        let entity_types_with_pct: Vec<_> = entity_types
            .into_iter()
            .map(|mut fv| {
                fv.percentage = if total > 0 {
                    (fv.count as f32 / total as f32) * 100.0
                } else {
                    0.0
                };
                fv
            })
            .collect();

        // Confidence levels (fixed)
        let confidence_levels = vec![
            FacetValue {
                name: "high".to_string(),
                count: 0,  // Would need aggregation query
                percentage: 0.0,
            },
            FacetValue {
                name: "medium".to_string(),
                count: 0,
                percentage: 0.0,
            },
            FacetValue {
                name: "low".to_string(),
                count: 0,
                percentage: 0.0,
            },
        ];

        // Date ranges (fixed)
        let date_ranges = vec![
            FacetValue {
                name: "today".to_string(),
                count: 0,
                percentage: 0.0,
            },
            FacetValue {
                name: "this_week".to_string(),
                count: 0,
                percentage: 0.0,
            },
            FacetValue {
                name: "this_month".to_string(),
                count: 0,
                percentage: 0.0,
            },
            FacetValue {
                name: "all_time".to_string(),
                count: 0,
                percentage: 0.0,
            },
        ];

        let elapsed = start_time.elapsed().as_millis();
        info!("Discovered {} entity types in {}ms", entity_types_with_pct.len(), elapsed);

        Ok(AvailableFacets {
            entity_types: entity_types_with_pct,
            relation_types: vec![],  // Empty for entities
            confidence_levels,
            date_ranges,
            total_results: total,
            facet_time_ms: elapsed,
        })
    }

    /// Discover facets for edge searches
    async fn discover_edge_facets(&self, limit: usize) -> Result<AvailableFacets, String> {
        let start_time = std::time::Instant::now();

        // Get relation types
        let relation_types = sqlx::query_as::<_, (String, i64)>(
            "SELECT relation_type, COUNT(*) as cnt 
             FROM memory_edge 
             WHERE fact_invalid_at IS NULL AND deleted_at IS NULL 
             GROUP BY relation_type 
             ORDER BY cnt DESC 
             LIMIT $1"
        )
            .bind(limit as i64)
            .fetch_all(&self.pool)
            .await
            .map_err(|e| format!("Failed to fetch relation types: {}", e))?
            .into_iter()
            .map(|(name, count)| FacetValue {
                name,
                count: count as usize,
                percentage: 0.0,
            })
            .collect::<Vec<_>>();

        // Get total count
        let total_count: (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM memory_edge WHERE fact_invalid_at IS NULL AND deleted_at IS NULL"
        )
            .fetch_one(&self.pool)
            .await
            .map_err(|e| format!("Failed to get total count: {}", e))?;

        let total = total_count.0 as usize;

        // Calculate percentages
        let relation_types_with_pct: Vec<_> = relation_types
            .into_iter()
            .map(|mut fv| {
                fv.percentage = if total > 0 {
                    (fv.count as f32 / total as f32) * 100.0
                } else {
                    0.0
                };
                fv
            })
            .collect();

        // Confidence levels (fixed)
        let confidence_levels = vec![
            FacetValue {
                name: "high".to_string(),
                count: 0,
                percentage: 0.0,
            },
            FacetValue {
                name: "medium".to_string(),
                count: 0,
                percentage: 0.0,
            },
            FacetValue {
                name: "low".to_string(),
                count: 0,
                percentage: 0.0,
            },
        ];

        let elapsed = start_time.elapsed().as_millis();
        info!("Discovered {} relation types in {}ms", relation_types_with_pct.len(), elapsed);

        Ok(AvailableFacets {
            entity_types: vec![],  // Empty for edges
            relation_types: relation_types_with_pct,
            confidence_levels,
            date_ranges: vec![],
            total_results: total,
            facet_time_ms: elapsed,
        })
    }

    /// Apply facet filters to a confidence threshold
    pub fn confidence_floor_from_level(&self, level: Option<&str>) -> f32 {
        match level {
            Some("high") => 0.8,
            Some("medium") => 0.5,
            Some("low") => 0.0,
            _ => 0.0,  // No filter
        }
    }

    /// Convert date range to start/end times
    pub fn date_range_to_times(&self, range: Option<&str>) -> (Option<DateTime<Utc>>, Option<DateTime<Utc>>) {
        let now = Utc::now();

        match range {
            Some("today") => {
                let start = now.with_hour(0).unwrap().with_minute(0).unwrap().with_second(0).unwrap();
                (Some(start), Some(now))
            }
            Some("this_week") => {
                let start = now - chrono::Duration::days(7);
                (Some(start), Some(now))
            }
            Some("this_month") => {
                let start = now - chrono::Duration::days(30);
                (Some(start), Some(now))
            }
            Some("this_year") => {
                let start = now - chrono::Duration::days(365);
                (Some(start), Some(now))
            }
            _ => (None, None),  // No filter
        }
    }

    /// Validate facet filters
    pub fn validate_filters(&self, filters: &FacetFilters) -> Result<(), String> {
        // Validate entity types (non-empty if provided)
        if let Some(types) = &filters.entity_types {
            if types.is_empty() {
                return Err("entity_types cannot be empty if provided".to_string());
            }
            if types.len() > 50 {
                return Err("entity_types cannot exceed 50 items".to_string());
            }
        }

        // Validate relation types
        if let Some(types) = &filters.relation_types {
            if types.is_empty() {
                return Err("relation_types cannot be empty if provided".to_string());
            }
            if types.len() > 50 {
                return Err("relation_types cannot exceed 50 items".to_string());
            }
        }

        // Validate confidence level
        if let Some(level) = &filters.confidence_level {
            if !["high", "medium", "low"].contains(&level.as_str()) {
                return Err("confidence_level must be 'high', 'medium', or 'low'".to_string());
            }
        }

        // Validate date range
        if let Some(range) = &filters.date_range {
            if !["today", "this_week", "this_month", "this_year", "all"].contains(&range.as_str()) {
                return Err("date_range must be 'today', 'this_week', 'this_month', 'this_year', or 'all'".to_string());
            }
        }

        Ok(())
    }
}