stateset_nsr/knowledge/

query.rs

//! Query Language for Knowledge Base
//!
//! Provides a query interface for retrieving information from the
//! knowledge base using pattern matching and logical conditions.

use super::*;
use std::collections::HashMap;

/// A query against the knowledge base
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Query {
    /// Pattern to match
    pub patterns: Vec<QueryPattern>,
    /// Filter conditions
    pub filters: Vec<QueryFilter>,
    /// Variables to return
    pub select: Vec<String>,
    /// Maximum results
    pub limit: Option<usize>,
    /// Results offset
    pub offset: Option<usize>,
    /// Order by
    pub order_by: Option<OrderBy>,
}

impl Query {
    pub fn new() -> Self {
        Self {
            patterns: Vec::new(),
            filters: Vec::new(),
            select: Vec::new(),
            limit: None,
            offset: None,
            order_by: None,
        }
    }

    /// Add a triple pattern to match
    pub fn pattern(mut self, subject: QueryTerm, predicate: &str, object: QueryTerm) -> Self {
        self.patterns.push(QueryPattern::Triple {
            subject,
            predicate: predicate.to_string(),
            object,
        });
        self
    }

    /// Add a filter condition
    pub fn filter(mut self, filter: QueryFilter) -> Self {
        self.filters.push(filter);
        self
    }

    /// Specify which variables to return
    pub fn select_vars(mut self, vars: Vec<&str>) -> Self {
        self.select = vars.into_iter().map(String::from).collect();
        self
    }

    /// Limit the number of results
    pub fn limit(mut self, limit: usize) -> Self {
        self.limit = Some(limit);
        self
    }

    /// Skip first N results
    pub fn offset(mut self, offset: usize) -> Self {
        self.offset = Some(offset);
        self
    }

    /// Order results
    pub fn order_by(mut self, var: &str, ascending: bool) -> Self {
        self.order_by = Some(OrderBy {
            variable: var.to_string(),
            ascending,
        });
        self
    }
}

impl Default for Query {
    fn default() -> Self {
        Self::new()
    }
}

/// A term in a query pattern
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryTerm {
    /// A variable to bind
    Variable(String),
    /// A specific entity
    Entity(EntityId),
    /// A constant value
    Constant(String),
    /// Any value (wildcard)
    Any,
}

impl QueryTerm {
    pub fn var(name: &str) -> Self {
        QueryTerm::Variable(name.to_string())
    }

    pub fn entity(id: EntityId) -> Self {
        QueryTerm::Entity(id)
    }

    pub fn constant(value: &str) -> Self {
        QueryTerm::Constant(value.to_string())
    }
}

/// A pattern to match in the query
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryPattern {
    /// Match a triple
    Triple {
        subject: QueryTerm,
        predicate: String,
        object: QueryTerm,
    },
    /// Match an entity by type
    EntityType {
        variable: String,
        entity_type: EntityType,
    },
    /// Match an entity by property
    Property {
        entity: QueryTerm,
        property: String,
        value: QueryTerm,
    },
    /// Optional pattern (LEFT JOIN semantics)
    Optional(Box<QueryPattern>),
    /// Union of patterns (OR)
    Union(Vec<QueryPattern>),
    /// Negation (NOT EXISTS)
    NotExists(Box<QueryPattern>),
}

/// Filter condition for query results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryFilter {
    /// Equality comparison
    Equals {
        variable: String,
        value: PropertyValue,
    },
    /// Not equal
    NotEquals {
        variable: String,
        value: PropertyValue,
    },
    /// Greater than (for numeric values)
    GreaterThan { variable: String, value: f64 },
    /// Less than
    LessThan { variable: String, value: f64 },
    /// String contains
    Contains { variable: String, substring: String },
    /// Regex match
    Regex { variable: String, pattern: String },
    /// Confidence threshold
    MinConfidence { variable: String, threshold: f32 },
    /// Has embedding
    HasEmbedding { variable: String },
    /// Similarity threshold (requires embeddings)
    SimilarTo {
        variable: String,
        target: EntityId,
        threshold: f32,
    },
    /// Logical AND of filters
    And(Vec<QueryFilter>),
    /// Logical OR of filters
    Or(Vec<QueryFilter>),
    /// Negation
    Not(Box<QueryFilter>),
}

/// Ordering specification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OrderBy {
    pub variable: String,
    pub ascending: bool,
}

/// Result of a query execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResult {
    /// Bindings for each result row
    pub bindings: Vec<HashMap<String, QueryValue>>,
    /// Total count before limit/offset
    pub total_count: usize,
    /// Execution time in milliseconds
    pub execution_time_ms: u64,
}

impl QueryResult {
    pub fn empty() -> Self {
        Self {
            bindings: Vec::new(),
            total_count: 0,
            execution_time_ms: 0,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.bindings.is_empty()
    }

    pub fn len(&self) -> usize {
        self.bindings.len()
    }
}

/// A value bound to a variable in query results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryValue {
    Entity(Entity),
    EntityId(EntityId),
    Property(PropertyValue),
    Literal(String),
    Number(f64),
    Boolean(bool),
    Null,
}

impl QueryValue {
    pub fn as_entity_id(&self) -> Option<&EntityId> {
        match self {
            QueryValue::EntityId(id) => Some(id),
            QueryValue::Entity(e) => Some(&e.id),
            _ => None,
        }
    }

    pub fn as_entity(&self) -> Option<&Entity> {
        match self {
            QueryValue::Entity(e) => Some(e),
            _ => None,
        }
    }
}

/// Query executor for the knowledge base
pub struct QueryExecutor<'a> {
    kb: &'a KnowledgeBase,
    org_id: Option<String>,
}

impl<'a> QueryExecutor<'a> {
    pub fn new(kb: &'a KnowledgeBase) -> Self {
        Self { kb, org_id: None }
    }

    /// Restrict query execution to a single organization (multi-tenant isolation).
    pub fn with_org_id(mut self, org_id: impl Into<String>) -> Self {
        self.org_id = Some(org_id.into());
        self
    }

    /// Execute a query and return results
    pub fn execute(&self, query: &Query) -> QueryResult {
        let start = std::time::Instant::now();

        // Start with empty bindings
        let mut bindings: Vec<HashMap<String, QueryValue>> = vec![HashMap::new()];

        // Process each pattern
        for pattern in &query.patterns {
            bindings = self.apply_pattern(pattern, bindings);
        }

        // Apply filters
        for filter in &query.filters {
            bindings.retain(|b| self.evaluate_filter(filter, b));
        }

        let total_count = bindings.len();

        // Apply ordering
        if let Some(ref order) = query.order_by {
            bindings.sort_by(|a, b| {
                let av = a.get(&order.variable);
                let bv = b.get(&order.variable);
                let cmp = self.compare_values(av, bv);
                if order.ascending {
                    cmp
                } else {
                    cmp.reverse()
                }
            });
        }

        // Apply offset
        if let Some(offset) = query.offset {
            bindings = bindings.into_iter().skip(offset).collect();
        }

        // Apply limit
        if let Some(limit) = query.limit {
            bindings = bindings.into_iter().take(limit).collect();
        }

        // Project to selected variables
        if !query.select.is_empty() {
            bindings = bindings
                .into_iter()
                .map(|b| {
                    b.into_iter()
                        .filter(|(k, _)| query.select.contains(k))
                        .collect()
                })
                .collect();
        }

        QueryResult {
            bindings,
            total_count,
            execution_time_ms: start.elapsed().as_millis() as u64,
        }
    }

    fn apply_pattern(
        &self,
        pattern: &QueryPattern,
        bindings: Vec<HashMap<String, QueryValue>>,
    ) -> Vec<HashMap<String, QueryValue>> {
        match pattern {
            QueryPattern::Triple {
                subject,
                predicate,
                object,
            } => {
                let mut new_bindings = Vec::new();

                for binding in bindings {
                    let matching_triples =
                        self.find_matching_triples(subject, predicate, object, &binding);

                    for triple in matching_triples {
                        let mut new_binding = binding.clone();

                        // Bind subject variable
                        if let QueryTerm::Variable(v) = subject {
                            new_binding
                                .insert(v.clone(), QueryValue::EntityId(triple.subject.clone()));
                        }

                        // Bind object variable
                        if let QueryTerm::Variable(v) = object {
                            new_binding
                                .insert(v.clone(), QueryValue::EntityId(triple.object.clone()));
                        }

                        new_bindings.push(new_binding);
                    }
                }

                new_bindings
            }

            QueryPattern::EntityType {
                variable,
                entity_type,
            } => {
                let mut new_bindings = Vec::new();

                for binding in bindings {
                    let entities = if let Some(ref org_id) = self.org_id {
                        self.kb.get_entities_by_type_and_org(entity_type, org_id)
                    } else {
                        self.kb.get_entities_by_type(entity_type)
                    };

                    for entity in entities {
                        let mut new_binding = binding.clone();
                        new_binding.insert(variable.clone(), QueryValue::Entity(entity));
                        new_bindings.push(new_binding);
                    }
                }

                new_bindings
            }

            QueryPattern::Property {
                entity,
                property,
                value,
            } => {
                let mut new_bindings = Vec::new();

                for binding in bindings {
                    // Get the entity
                    let entity_id = match entity {
                        QueryTerm::Variable(v) => {
                            binding.get(v).and_then(|qv| qv.as_entity_id()).cloned()
                        }
                        QueryTerm::Entity(id) => Some(id.clone()),
                        _ => None,
                    };

                    if let Some(id) = entity_id {
                        if let Some(ent) = self.kb.get_entity(&id) {
                            if let Some(ref org_id) = self.org_id {
                                if ent.org_id != *org_id {
                                    continue;
                                }
                            }
                            if let Some(prop_value) = ent.properties.get(property) {
                                let mut new_binding = binding.clone();

                                // Handle value binding
                                match value {
                                    QueryTerm::Variable(v) => {
                                        let qv = self.property_to_query_value(&prop_value);
                                        new_binding.insert(v.clone(), qv);
                                        new_bindings.push(new_binding);
                                    }
                                    QueryTerm::Constant(c) => {
                                        if self.property_matches_constant(&prop_value, c) {
                                            new_bindings.push(new_binding);
                                        }
                                    }
                                    QueryTerm::Any => {
                                        new_bindings.push(new_binding);
                                    }
                                    _ => {}
                                }
                            }
                        }
                    }
                }

                new_bindings
            }

            QueryPattern::Optional(inner) => {
                let mut new_bindings = Vec::new();

                for binding in bindings {
                    let matched = self.apply_pattern(inner, vec![binding.clone()]);

                    if matched.is_empty() {
                        new_bindings.push(binding);
                    } else {
                        new_bindings.extend(matched);
                    }
                }

                new_bindings
            }

            QueryPattern::Union(patterns) => {
                let mut all_bindings = Vec::new();
                for pattern in patterns {
                    let matched = self.apply_pattern(pattern, bindings.clone());
                    all_bindings.extend(matched);
                }
                all_bindings
            }

            QueryPattern::NotExists(inner) => bindings
                .into_iter()
                .filter(|b| {
                    let matched = self.apply_pattern(inner, vec![b.clone()]);
                    matched.is_empty()
                })
                .collect(),
        }
    }

    fn find_matching_triples(
        &self,
        subject: &QueryTerm,
        predicate: &str,
        object: &QueryTerm,
        binding: &HashMap<String, QueryValue>,
    ) -> Vec<Triple> {
        self.kb
            .triples
            .iter()
            .filter(|t| {
                if let Some(ref org_id) = self.org_id {
                    if t.org_id != *org_id {
                        return false;
                    }
                }
                // Check predicate
                if t.predicate != predicate {
                    return false;
                }

                // Check subject
                let subject_matches = match subject {
                    QueryTerm::Variable(v) => {
                        if let Some(bound) = binding.get(v) {
                            bound.as_entity_id() == Some(&t.subject)
                        } else {
                            true // Unbound, will bind
                        }
                    }
                    QueryTerm::Entity(id) => &t.subject == id,
                    QueryTerm::Constant(c) => t.subject.0 == *c,
                    QueryTerm::Any => true,
                };

                if !subject_matches {
                    return false;
                }

                // Check object
                match object {
                    QueryTerm::Variable(v) => {
                        if let Some(bound) = binding.get(v) {
                            bound.as_entity_id() == Some(&t.object)
                        } else {
                            true // Unbound, will bind
                        }
                    }
                    QueryTerm::Entity(id) => &t.object == id,
                    QueryTerm::Constant(c) => t.object.0 == *c,
                    QueryTerm::Any => true,
                }
            })
            .map(|t| t.clone())
            .collect()
    }

    fn evaluate_filter(&self, filter: &QueryFilter, binding: &HashMap<String, QueryValue>) -> bool {
        match filter {
            QueryFilter::Equals { variable, value } => binding
                .get(variable)
                .is_some_and(|v| self.query_value_equals(v, value)),

            QueryFilter::NotEquals { variable, value } => binding
                .get(variable)
                .map_or(true, |v| !self.query_value_equals(v, value)),

            QueryFilter::GreaterThan { variable, value } => binding
                .get(variable)
                .and_then(|v| self.query_value_as_f64(v))
                .is_some_and(|n| n > *value),

            QueryFilter::LessThan { variable, value } => binding
                .get(variable)
                .and_then(|v| self.query_value_as_f64(v))
                .is_some_and(|n| n < *value),

            QueryFilter::Contains {
                variable,
                substring,
            } => binding
                .get(variable)
                .and_then(|v| self.query_value_as_string(v))
                .is_some_and(|s| s.contains(substring)),

            QueryFilter::Regex { variable, pattern } => {
                if let Ok(re) = regex::Regex::new(pattern) {
                    binding
                        .get(variable)
                        .and_then(|v| self.query_value_as_string(v))
                        .is_some_and(|s| re.is_match(&s))
                } else {
                    false
                }
            }

            QueryFilter::MinConfidence {
                variable,
                threshold,
            } => binding
                .get(variable)
                .and_then(|v| v.as_entity())
                .is_some_and(|e| e.confidence >= *threshold),

            QueryFilter::HasEmbedding { variable } => binding
                .get(variable)
                .and_then(|v| v.as_entity())
                .is_some_and(|e| e.embedding.is_some()),

            QueryFilter::SimilarTo {
                variable,
                target,
                threshold,
            } => {
                if let Some(entity_id) = binding.get(variable).and_then(|v| v.as_entity_id()) {
                    if let Some(ref org_id) = self.org_id {
                        if self
                            .kb
                            .get_entity(target)
                            .is_some_and(|e| e.org_id != *org_id)
                        {
                            return false;
                        }
                    }
                    self.kb
                        .embedding_similarity(entity_id, target)
                        .is_some_and(|sim| sim >= *threshold)
                } else {
                    false
                }
            }

            QueryFilter::And(filters) => filters.iter().all(|f| self.evaluate_filter(f, binding)),

            QueryFilter::Or(filters) => filters.iter().any(|f| self.evaluate_filter(f, binding)),

            QueryFilter::Not(inner) => !self.evaluate_filter(inner, binding),
        }
    }

    fn query_value_equals(&self, qv: &QueryValue, pv: &PropertyValue) -> bool {
        match (qv, pv) {
            (QueryValue::Literal(s), PropertyValue::String(ps)) => s == ps,
            (QueryValue::Number(n), PropertyValue::Float(f)) => (*n - f).abs() < f64::EPSILON,
            (QueryValue::Number(n), PropertyValue::Integer(i)) => {
                (*n - *i as f64).abs() < f64::EPSILON
            }
            (QueryValue::Boolean(b), PropertyValue::Boolean(pb)) => b == pb,
            _ => false,
        }
    }

    fn query_value_as_f64(&self, qv: &QueryValue) -> Option<f64> {
        match qv {
            QueryValue::Number(n) => Some(*n),
            QueryValue::Property(PropertyValue::Float(f)) => Some(*f),
            QueryValue::Property(PropertyValue::Integer(i)) => Some(*i as f64),
            _ => None,
        }
    }

    fn query_value_as_string(&self, qv: &QueryValue) -> Option<String> {
        match qv {
            QueryValue::Literal(s) => Some(s.clone()),
            QueryValue::Property(PropertyValue::String(s)) => Some(s.clone()),
            QueryValue::Entity(e) => Some(e.name.clone()),
            _ => None,
        }
    }

    fn property_to_query_value(&self, pv: &PropertyValue) -> QueryValue {
        match pv {
            PropertyValue::String(s) => QueryValue::Literal(s.clone()),
            PropertyValue::Integer(i) => QueryValue::Number(*i as f64),
            PropertyValue::Float(f) => QueryValue::Number(*f),
            PropertyValue::Boolean(b) => QueryValue::Boolean(*b),
            PropertyValue::EntityRef(id) => QueryValue::EntityId(id.clone()),
            _ => QueryValue::Null,
        }
    }

    fn property_matches_constant(&self, pv: &PropertyValue, constant: &str) -> bool {
        match pv {
            PropertyValue::String(s) => s == constant,
            PropertyValue::Integer(i) => constant.parse::<i64>() == Ok(*i),
            PropertyValue::Float(f) => constant
                .parse::<f64>()
                .is_ok_and(|c| (*f - c).abs() < f64::EPSILON),
            PropertyValue::Boolean(b) => constant.parse::<bool>() == Ok(*b),
            _ => false,
        }
    }

    fn compare_values(&self, a: Option<&QueryValue>, b: Option<&QueryValue>) -> std::cmp::Ordering {
        use std::cmp::Ordering;

        match (a, b) {
            (None, None) => Ordering::Equal,
            (None, Some(_)) => Ordering::Less,
            (Some(_), None) => Ordering::Greater,
            (Some(av), Some(bv)) => match (av, bv) {
                (QueryValue::Number(an), QueryValue::Number(bn)) => {
                    an.partial_cmp(bn).unwrap_or(Ordering::Equal)
                }
                (QueryValue::Literal(as_), QueryValue::Literal(bs)) => as_.cmp(bs),
                _ => Ordering::Equal,
            },
        }
    }
}

/// Convenience builder for creating queries
#[derive(Default)]
pub struct QueryBuilder {
    query: Query,
}

impl QueryBuilder {
    pub fn new() -> Self {
        Self {
            query: Query::new(),
        }
    }

    /// Match a triple pattern with variables
    pub fn matches(mut self, subject: &str, predicate: &str, object: &str) -> Self {
        let subj = subject
            .strip_prefix('?')
            .map(QueryTerm::var)
            .unwrap_or_else(|| QueryTerm::constant(subject));

        let obj = object
            .strip_prefix('?')
            .map(QueryTerm::var)
            .unwrap_or_else(|| QueryTerm::constant(object));

        self.query.patterns.push(QueryPattern::Triple {
            subject: subj,
            predicate: predicate.to_string(),
            object: obj,
        });
        self
    }

    /// Filter by property value
    pub fn where_eq(mut self, variable: &str, value: PropertyValue) -> Self {
        self.query.filters.push(QueryFilter::Equals {
            variable: variable.to_string(),
            value,
        });
        self
    }

    /// Select specific variables
    pub fn select(mut self, vars: &[&str]) -> Self {
        self.query.select = vars.iter().map(|s| s.to_string()).collect();
        self
    }

    /// Limit results
    pub fn limit(mut self, n: usize) -> Self {
        self.query.limit = Some(n);
        self
    }

    /// Build the query
    pub fn build(self) -> Query {
        self.query
    }
}

// ============================================================================
// Cached Query Executor
// ============================================================================

use crate::cache::{Cache, CacheConfig};
use std::hash::{Hash, Hasher};
use std::sync::Arc;

/// A cache key for queries based on their serialized representation
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueryCacheKey(String);

impl Hash for QueryCacheKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl QueryCacheKey {
    /// Create a cache key from a query
    pub fn from_query(query: &Query) -> Self {
        Self::from_query_with_org(query, None)
    }

    /// Create a cache key from a query scoped to an organization.
    pub fn from_query_with_org(query: &Query, org_id: Option<&str>) -> Self {
        let key = serde_json::to_string(query).unwrap_or_default();
        let key = match org_id {
            Some(org) => format!("org:{}|{}", org, key),
            None => key,
        };
        QueryCacheKey(key)
    }
}

/// Configuration for the query cache
#[derive(Debug, Clone)]
pub struct QueryCacheConfig {
    /// Maximum number of cached query results
    pub max_entries: usize,
    /// Time-to-live for cached results in seconds
    pub ttl_seconds: u64,
    /// Whether caching is enabled
    pub enabled: bool,
}

impl Default for QueryCacheConfig {
    fn default() -> Self {
        Self {
            max_entries: 1000,
            ttl_seconds: 300, // 5 minutes
            enabled: true,
        }
    }
}

/// A query executor with built-in caching for repeated queries.
///
/// The cached executor wraps a standard `QueryExecutor` and transparently
/// caches query results. Cache keys are computed from the serialized query,
/// so identical queries will hit the cache.
///
/// # Cache Invalidation
///
/// The cache does **not** automatically invalidate when the knowledge base
/// changes. For write-heavy workloads, consider:
/// - Using a shorter TTL
/// - Calling `invalidate_all()` after batch writes
/// - Using the non-cached `QueryExecutor` directly
///
/// # Example
///
/// ```rust,ignore
/// use stateset_nsr::knowledge::query::{CachedQueryExecutor, QueryCacheConfig};
///
/// let cache_config = QueryCacheConfig {
///     max_entries: 5000,
///     ttl_seconds: 600,
///     enabled: true,
/// };
///
/// let executor = CachedQueryExecutor::new(&kb, cache_config);
///
/// // First query - cache miss, executes query
/// let result1 = executor.execute(&query);
///
/// // Second identical query - cache hit, returns cached result
/// let result2 = executor.execute(&query);
///
/// // Check cache statistics
/// let stats = executor.cache_stats();
/// println!("Cache hit rate: {:.2}%", stats.hit_rate * 100.0);
/// ```
pub struct CachedQueryExecutor<'a> {
    executor: QueryExecutor<'a>,
    cache: Arc<Cache<QueryCacheKey, QueryResult>>,
    config: QueryCacheConfig,
    org_id: Option<String>,
}

impl<'a> CachedQueryExecutor<'a> {
    /// Create a new cached query executor
    pub fn new(kb: &'a KnowledgeBase, config: QueryCacheConfig) -> Self {
        let cache = Arc::new(Cache::new(CacheConfig {
            max_entries: config.max_entries,
            ttl_seconds: config.ttl_seconds,
            enable_stats: true,
        }));

        Self {
            executor: QueryExecutor::new(kb),
            cache,
            config,
            org_id: None,
        }
    }

    /// Restrict query execution (and cache keys) to a single organization.
    pub fn with_org_id(mut self, org_id: impl Into<String>) -> Self {
        let org_id = org_id.into();
        self.executor = self.executor.with_org_id(org_id.clone());
        self.org_id = Some(org_id);
        self
    }

    /// Execute a query with caching
    ///
    /// If caching is enabled and a cached result exists for this query,
    /// returns the cached result. Otherwise, executes the query and
    /// caches the result before returning.
    pub fn execute(&self, query: &Query) -> QueryResult {
        if !self.config.enabled {
            return self.executor.execute(query);
        }

        let cache_key = QueryCacheKey::from_query_with_org(query, self.org_id.as_deref());

        // Try to get from cache
        if let Some(cached) = self.cache.get(&cache_key) {
            return cached;
        }

        // Execute query and cache result
        let result = self.executor.execute(query);
        self.cache.insert(cache_key, result.clone());
        result
    }

    /// Execute a query bypassing the cache
    ///
    /// Useful when you need fresh results or when the query is unlikely
    /// to be repeated.
    pub fn execute_uncached(&self, query: &Query) -> QueryResult {
        self.executor.execute(query)
    }

    /// Invalidate all cached results
    ///
    /// Call this after making significant changes to the knowledge base
    /// to ensure queries return fresh results.
    pub fn invalidate_all(&self) {
        self.cache.clear();
    }

    /// Get cache statistics
    pub fn cache_stats(&self) -> crate::cache::CacheStats {
        self.cache.stats()
    }

    /// Check if a query result is cached
    pub fn is_cached(&self, query: &Query) -> bool {
        let cache_key = QueryCacheKey::from_query_with_org(query, self.org_id.as_deref());
        self.cache.contains(&cache_key)
    }

    /// Get the number of cached query results
    pub fn cached_count(&self) -> usize {
        self.cache.len()
    }
}

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

    #[test]
    fn query_executor_respects_org_id_filter() {
        let kb = KnowledgeBase::new();

        let paris_a = kb.add_entity(Entity::with_org_id(
            "Paris",
            EntityType::Instance,
            "org_a".to_string(),
        ));
        let france_a = kb.add_entity(Entity::with_org_id(
            "France",
            EntityType::Instance,
            "org_a".to_string(),
        ));
        let paris_b = kb.add_entity(Entity::with_org_id(
            "Paris",
            EntityType::Instance,
            "org_b".to_string(),
        ));
        let france_b = kb.add_entity(Entity::with_org_id(
            "France",
            EntityType::Instance,
            "org_b".to_string(),
        ));

        kb.add_triple(Triple::with_org_id(
            paris_a,
            "capital_of",
            france_a,
            "org_a".to_string(),
        ))
        .unwrap();
        kb.add_triple(Triple::with_org_id(
            paris_b,
            "capital_of",
            france_b,
            "org_b".to_string(),
        ))
        .unwrap();

        let query = QueryBuilder::new()
            .matches("?city", "capital_of", "?country")
            .build();

        let result_a = QueryExecutor::new(&kb).with_org_id("org_a").execute(&query);
        assert_eq!(result_a.bindings.len(), 1);

        let result_b = QueryExecutor::new(&kb).with_org_id("org_b").execute(&query);
        assert_eq!(result_b.bindings.len(), 1);
    }

    #[test]
    fn query_cache_key_includes_org_id() {
        let query = QueryBuilder::new()
            .matches("?city", "capital_of", "?country")
            .build();

        let key_a = QueryCacheKey::from_query_with_org(&query, Some("org_a"));
        let key_b = QueryCacheKey::from_query_with_org(&query, Some("org_b"));
        assert_ne!(key_a, key_b);
    }

    #[test]
    fn test_query_execution() {
        let kb = KnowledgeBase::new();

        let paris = kb.add_entity(
            Entity::new("Paris", EntityType::Instance)
                .with_property("population", PropertyValue::Integer(2_161_000)),
        );
        let france = kb.add_entity(Entity::new("France", EntityType::Instance));
        let berlin = kb.add_entity(
            Entity::new("Berlin", EntityType::Instance)
                .with_property("population", PropertyValue::Integer(3_645_000)),
        );
        let germany = kb.add_entity(Entity::new("Germany", EntityType::Instance));

        kb.add_triple(Triple::new(paris.clone(), "capital_of", france.clone()))
            .unwrap();
        kb.add_triple(Triple::new(berlin.clone(), "capital_of", germany.clone()))
            .unwrap();

        let executor = QueryExecutor::new(&kb);

        // Query: Find all capitals
        let query = QueryBuilder::new()
            .matches("?city", "capital_of", "?country")
            .build();

        let result = executor.execute(&query);
        assert_eq!(result.bindings.len(), 2);
    }

    #[test]
    fn optional_pattern_preserves_base_binding_when_no_match() {
        let kb = KnowledgeBase::new();

        let alice = kb.add_entity(Entity::new("Alice", EntityType::Instance));
        let bob = kb.add_entity(Entity::new("Bob", EntityType::Instance));

        kb.add_triple(Triple::new(alice.clone(), "friend_of", bob.clone()))
            .unwrap();

        let query = Query {
            patterns: vec![
                QueryPattern::Triple {
                    subject: QueryTerm::Variable("s".into()),
                    predicate: "friend_of".to_string(),
                    object: QueryTerm::Variable("o".into()),
                },
                QueryPattern::Optional(Box::new(QueryPattern::Triple {
                    subject: QueryTerm::Variable("s".into()),
                    predicate: "likes".to_string(),
                    object: QueryTerm::Variable("x".into()),
                })),
            ],
            filters: Vec::new(),
            select: vec!["s".to_string(), "o".to_string(), "x".to_string()],
            limit: None,
            offset: None,
            order_by: None,
        };

        let executor = QueryExecutor::new(&kb);
        let result = executor.execute(&query);

        assert_eq!(result.bindings.len(), 1);
        let binding = &result.bindings[0];
        assert!(binding.contains_key("s"));
        assert!(binding.contains_key("o"));
        assert!(!binding.contains_key("x")); // optional pattern had no match, but base binding remained
    }

    #[test]
    fn test_cached_query_executor() {
        let kb = KnowledgeBase::new();

        let paris = kb.add_entity(Entity::new("Paris", EntityType::Instance));
        let france = kb.add_entity(Entity::new("France", EntityType::Instance));

        kb.add_triple(Triple::new(paris.clone(), "capital_of", france.clone()))
            .unwrap();

        let config = QueryCacheConfig::default();
        let cached_executor = CachedQueryExecutor::new(&kb, config);

        let query = QueryBuilder::new()
            .matches("?city", "capital_of", "?country")
            .build();

        // First execution - cache miss
        let result1 = cached_executor.execute(&query);
        assert_eq!(result1.bindings.len(), 1);

        let stats1 = cached_executor.cache_stats();
        assert_eq!(stats1.misses, 1);
        assert_eq!(stats1.hits, 0);

        // Second execution - cache hit
        let result2 = cached_executor.execute(&query);
        assert_eq!(result2.bindings.len(), 1);

        let stats2 = cached_executor.cache_stats();
        assert_eq!(stats2.misses, 1);
        assert_eq!(stats2.hits, 1);
        assert!(stats2.hit_rate > 0.4); // Should be ~50%

        // Verify cache contains the query
        assert!(cached_executor.is_cached(&query));
        assert_eq!(cached_executor.cached_count(), 1);
    }

    #[test]
    fn test_cached_executor_invalidation() {
        let kb = KnowledgeBase::new();

        let alice = kb.add_entity(Entity::new("Alice", EntityType::Instance));
        let bob = kb.add_entity(Entity::new("Bob", EntityType::Instance));

        kb.add_triple(Triple::new(alice.clone(), "knows", bob.clone()))
            .unwrap();

        let cached_executor = CachedQueryExecutor::new(&kb, QueryCacheConfig::default());

        let query = QueryBuilder::new()
            .matches("?person", "knows", "?other")
            .build();

        // Execute and cache
        cached_executor.execute(&query);
        assert_eq!(cached_executor.cached_count(), 1);

        // Invalidate cache
        cached_executor.invalidate_all();
        assert_eq!(cached_executor.cached_count(), 0);
        assert!(!cached_executor.is_cached(&query));
    }
}