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

//! Inference Engine (Phase 5.2)
//!
//! Rule-based inference with graph traversal, transitive closure, and
//! confidence propagation through reasoning chains.

use std::collections::{HashMap, HashSet, VecDeque};
use sqlx::PgPool;
use serde::{Deserialize, Serialize};
use tracing::{debug, warn};

/// Inference rule
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceRule {
    /// Rule ID
    pub id: String,
    /// Antecedent predicate (e.g., "depends_on")
    pub antecedent: String,
    /// Medial predicate (optional, for chain rules)
    pub medial: Option<String>,
    /// Consequent predicate (e.g., "related_to")
    pub consequent: String,
    /// Confidence multiplier (0.0-1.0)
    pub confidence_multiplier: f32,
    /// Description
    pub description: String,
}

/// Inferred fact
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct InferredFact {
    /// Source entity ID
    pub source_id: String,
    /// Source entity name
    pub source_name: String,
    /// Target entity ID
    pub target_id: String,
    /// Target entity name
    pub target_name: String,
    /// Inferred relation type
    pub relation_type: String,
    /// Confidence (0.0-1.0)
    pub confidence: f32,
    /// Reasoning chain that led to inference
    pub reasoning_chain: Vec<String>,
    /// Rule IDs applied
    pub rule_ids: Vec<String>,
}

/// Reasoning path
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasoningPath {
    /// Path steps: entity_id → entity_id → ...
    pub path: Vec<String>,
    /// Relations between steps: relation_type → relation_type → ...
    pub relations: Vec<String>,
    /// Accumulated confidence (product of step confidences)
    pub confidence: f32,
    /// Steps in path
    pub step_count: usize,
}

/// Transitive closure result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransitiveClosure {
    /// Starting entity ID
    pub source_id: String,
    /// All reachable entities with relation type and confidence
    pub reachable: Vec<ReachableEntity>,
    /// Total entities reached
    pub entity_count: usize,
    /// Total edges in closure
    pub edge_count: usize,
}

/// Reachable entity info
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReachableEntity {
    /// Entity ID
    pub entity_id: String,
    /// Entity name
    pub entity_name: String,
    /// Relation type from source
    pub relation_type: String,
    /// Combined confidence
    pub confidence: f32,
    /// Hop distance from source
    pub distance: usize,
}

/// Inference Engine
pub struct InferenceEngine {
    pool: PgPool,
    rules: Vec<InferenceRule>,
}

impl InferenceEngine {
    pub fn new(pool: PgPool, rules: Vec<InferenceRule>) -> Self {
        InferenceEngine { pool, rules }
    }

    /// Perform rule-based inference
    ///
    /// Applies inference rules to graph, generating new facts
    pub async fn infer_facts(
        &self,
        project_id: &str,
        entity_id: &str,
        max_hops: usize,
    ) -> Result<Vec<InferredFact>, String> {
        if entity_id.is_empty() || max_hops == 0 {
            return Ok(vec![]);
        }

        let mut inferred = Vec::new();
        let mut visited = HashSet::new();

        // BFS from entity_id applying rules at each step
        let mut queue = VecDeque::new();
        queue.push_back((entity_id.to_string(), 0, 1.0, vec![]));

        while let Some((current_id, depth, confidence, chain)) = queue.pop_front() {
            if depth >= max_hops || visited.contains(&current_id) {
                continue;
            }
            visited.insert(current_id.clone());

            // Get edges from current entity
            let edges = self.fetch_entity_edges(&current_id, project_id).await?;

            for edge in edges {
                // Apply each rule
                for rule in &self.rules {
                    if edge.relation_type == rule.antecedent {
                        let new_confidence = (confidence * rule.confidence_multiplier).min(1.0);

                        if new_confidence > 0.1 {
                            let mut new_chain = chain.clone();
                            new_chain.push(format!("{} --{}→ {}", 
                                current_id, rule.consequent, edge.target_id));

                            inferred.push(InferredFact {
                                source_id: entity_id.to_string(),
                                source_name: "Unknown".to_string(),
                                target_id: edge.target_id.clone(),
                                target_name: edge.target_name.clone(),
                                relation_type: rule.consequent.clone(),
                                confidence: new_confidence,
                                reasoning_chain: new_chain.clone(),
                                rule_ids: vec![rule.id.clone()],
                            });

                            queue.push_back((
                                edge.target_id.clone(),
                                depth + 1,
                                new_confidence,
                                new_chain,
                            ));
                        }
                    }
                }
            }
        }

        // Deduplicate by (source, target, relation)
        let mut deduped: HashMap<(String, String, String), InferredFact> = HashMap::new();
        for fact in inferred {
            let key = (fact.source_id.clone(), fact.target_id.clone(), fact.relation_type.clone());
            deduped.entry(key).or_insert(fact);
        }

        Ok(deduped.into_values().collect())
    }

    /// Compute transitive closure for entity
    pub async fn transitive_closure(
        &self,
        entity_id: &str,
        project_id: &str,
        relation_type: Option<&str>,
        max_hops: usize,
    ) -> Result<TransitiveClosure, String> {
        let mut reachable = Vec::new();
        let mut visited: HashMap<String, (f32, usize)> = HashMap::new();

        let mut queue = VecDeque::new();
        queue.push_back((entity_id.to_string(), 1.0, 0));
        visited.insert(entity_id.to_string(), (1.0, 0));

        while let Some((current_id, confidence, distance)) = queue.pop_front() {
            if distance >= max_hops {
                continue;
            }

            let edges = self.fetch_entity_edges(&current_id, project_id).await?;

            for edge in edges {
                // Filter by relation type if specified
                if let Some(rel_type) = relation_type {
                    if edge.relation_type != rel_type {
                        continue;
                    }
                }

                let new_confidence = confidence * 0.95; // Decay confidence per hop

                let target = edge.target_id.clone();
                let entry = visited.entry(target.clone()).or_insert((new_confidence, distance + 1));

                // Keep higher confidence path
                if new_confidence > entry.0 {
                    entry.0 = new_confidence;
                    entry.1 = distance + 1;

                    reachable.push(ReachableEntity {
                        entity_id: target.clone(),
                        entity_name: edge.target_name.clone(),
                        relation_type: edge.relation_type.clone(),
                        confidence: new_confidence,
                        distance: distance + 1,
                    });

                    queue.push_back((target, new_confidence, distance + 1));
                }
            }
        }

        let edge_count = reachable.len();
        let entity_count = visited.len() - 1; // Exclude starting entity

        Ok(TransitiveClosure {
            source_id: entity_id.to_string(),
            reachable,
            entity_count,
            edge_count,
        })
    }

    /// Find all reasoning paths between entities
    pub async fn find_reasoning_paths(
        &self,
        source_id: &str,
        target_id: &str,
        project_id: &str,
        max_hops: usize,
    ) -> Result<Vec<ReasoningPath>, String> {
        let mut paths = Vec::new();
        let mut visited = HashSet::new();

        self.dfs_paths(
            source_id,
            target_id,
            project_id,
            max_hops,
            &mut vec![source_id.to_string()],
            &mut vec![],
            &mut vec![1.0],
            &mut visited,
            &mut paths,
        ).await?;

        Ok(paths)
    }

    /// Check if fact can be inferred from rules
    pub fn check_inference_validity(
        &self,
        antecedent: &str,
        consequent: &str,
    ) -> Option<(String, f32)> {
        for rule in &self.rules {
            if rule.antecedent == antecedent && rule.consequent == consequent {
                return Some((rule.id.clone(), rule.confidence_multiplier));
            }
        }
        None
    }

    /// Get applicable rules for relation type
    pub fn get_applicable_rules(&self, relation_type: &str) -> Vec<&InferenceRule> {
        self.rules.iter().filter(|r| r.antecedent == relation_type).collect()
    }

    // ========== Private Helper Methods ==========

    /// Fetch edges from entity
    async fn fetch_entity_edges(
        &self,
        entity_id: &str,
        project_id: &str,
    ) -> Result<Vec<EdgeInfo>, String> {
        // Stub: would query database
        Ok(vec![])
    }

    /// DFS to find all paths
    async fn dfs_paths(
        &self,
        current: &str,
        target: &str,
        project_id: &str,
        remaining_hops: usize,
        path: &mut Vec<String>,
        relations: &mut Vec<String>,
        confidences: &mut Vec<f32>,
        visited: &mut HashSet<String>,
        results: &mut Vec<ReasoningPath>,
    ) -> Result<(), String> {
        if remaining_hops == 0 {
            return Ok(());
        }

        if current == target && path.len() > 1 {
            let confidence = confidences.iter().product();
            results.push(ReasoningPath {
                path: path.clone(),
                relations: relations.clone(),
                confidence,
                step_count: path.len(),
            });
            return Ok(());
        }

        let edges = self.fetch_entity_edges(current, project_id).await?;

        for edge in edges {
            if !visited.contains(&edge.target_id) {
                visited.insert(edge.target_id.clone());

                path.push(edge.target_id.clone());
                relations.push(edge.relation_type.clone());
                confidences.push(0.9); // Nominal confidence per edge

                self.dfs_paths(
                    &edge.target_id,
                    target,
                    project_id,
                    remaining_hops - 1,
                    path,
                    relations,
                    confidences,
                    visited,
                    results,
                ).await?;

                path.pop();
                relations.pop();
                confidences.pop();
                visited.remove(&edge.target_id);
            }
        }

        Ok(())
    }
}

/// Internal edge info
struct EdgeInfo {
    source_id: String,
    target_id: String,
    target_name: String,
    relation_type: String,
}

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

    fn create_test_rules() -> Vec<InferenceRule> {
        vec![
            InferenceRule {
                id: "r1".to_string(),
                antecedent: "depends_on".to_string(),
                medial: None,
                consequent: "related_to".to_string(),
                confidence_multiplier: 0.9,
                description: "Depends implies related".to_string(),
            },
            InferenceRule {
                id: "r2".to_string(),
                antecedent: "uses".to_string(),
                medial: None,
                consequent: "related_to".to_string(),
                confidence_multiplier: 0.85,
                description: "Uses implies related".to_string(),
            },
        ]
    }

    #[test]
    fn test_inference_rule_structure() {
        let rule = InferenceRule {
            id: "r1".to_string(),
            antecedent: "depends_on".to_string(),
            medial: None,
            consequent: "related_to".to_string(),
            confidence_multiplier: 0.9,
            description: "Test rule".to_string(),
        };
        assert_eq!(rule.antecedent, "depends_on");
        assert_eq!(rule.consequent, "related_to");
    }

    #[test]
    fn test_inferred_fact_structure() {
        let fact = InferredFact {
            source_id: "e1".to_string(),
            source_name: "Entity1".to_string(),
            target_id: "e2".to_string(),
            target_name: "Entity2".to_string(),
            relation_type: "related_to".to_string(),
            confidence: 0.81,
            reasoning_chain: vec!["e1 --depends_on→ e2".to_string()],
            rule_ids: vec!["r1".to_string()],
        };
        assert_eq!(fact.confidence, 0.81);
        assert_eq!(fact.reasoning_chain.len(), 1);
    }

    #[test]
    fn test_reasoning_path_structure() {
        let path = ReasoningPath {
            path: vec!["e1".to_string(), "e2".to_string(), "e3".to_string()],
            relations: vec!["depends_on".to_string(), "uses".to_string()],
            confidence: 0.75,
            step_count: 3,
        };
        assert_eq!(path.step_count, 3);
        assert_eq!(path.path.len(), 3);
    }

    #[test]
    fn test_transitive_closure_structure() {
        let closure = TransitiveClosure {
            source_id: "e1".to_string(),
            reachable: vec![],
            entity_count: 0,
            edge_count: 0,
        };
        assert_eq!(closure.entity_count, 0);
    }

    #[test]
    fn test_reachable_entity_structure() {
        let entity = ReachableEntity {
            entity_id: "e2".to_string(),
            entity_name: "Entity2".to_string(),
            relation_type: "related_to".to_string(),
            confidence: 0.85,
            distance: 1,
        };
        assert_eq!(entity.distance, 1);
        assert!(entity.confidence > 0.8);
    }

    #[test]
    fn test_confidence_multiplier() {
        let rule = &create_test_rules()[0];
        let base_confidence = 0.9;
        let result = base_confidence * rule.confidence_multiplier;
        assert!(result < base_confidence);
    }

    #[test]
    fn test_confidence_decay_single_hop() {
        let confidence = 1.0;
        let decay = 0.95;
        let result = confidence * decay;
        assert_eq!(result, 0.95);
    }

    #[test]
    fn test_confidence_decay_two_hops() {
        let confidence = 1.0;
        let decay = 0.95;
        let result = confidence * decay * decay;
        assert!((result - 0.9025).abs() < 0.0001);
    }

    #[test]
    fn test_confidence_chaining() {
        let conf1 = 0.9;
        let conf2 = 0.85;
        let result = conf1 * conf2;
        assert!((result - 0.765).abs() < 0.0001);
    }

    #[test]
    fn test_confidence_bounds() {
        let confidence = 0.95 * 1.1; // Exceed 1.0
        let bounded = confidence.min(1.0);
        assert_eq!(bounded, 1.0);
    }

    #[test]
    fn test_rule_matching() {
        let rules = create_test_rules();
        let rule = rules.iter().find(|r| r.antecedent == "depends_on").unwrap();
        assert_eq!(rule.consequent, "related_to");
    }

    #[test]
    fn test_rule_no_match() {
        let rules = create_test_rules();
        let rule = rules.iter().find(|r| r.antecedent == "nonexistent");
        assert!(rule.is_none());
    }

    #[test]
    fn test_inferred_fact_confidence_calculation() {
        let base = 1.0;
        let multiplier = 0.9;
        let final_conf = (base * multiplier).min(1.0);
        assert_eq!(final_conf, 0.9);
    }

    #[test]
    fn test_reasoning_chain_construction() {
        let chain = vec![
            "e1 --depends_on→ e2".to_string(),
            "e2 --uses→ e3".to_string(),
        ];
        assert_eq!(chain.len(), 2);
    }

    #[test]
    fn test_path_step_count() {
        let path_len = 3;
        let step_count = path_len;
        assert_eq!(step_count, 3);
    }

    #[test]
    fn test_hop_distance_tracking() {
        let mut distance = 0;
        distance += 1; // Hop 1
        distance += 1; // Hop 2
        assert_eq!(distance, 2);
    }

    #[test]
    fn test_max_hops_limit() {
        let max_hops = 5;
        let current_hops = 3;
        assert!(current_hops < max_hops);
    }

    #[test]
    fn test_rule_confidence_multiplier_range() {
        let multipliers = vec![0.5, 0.75, 0.9, 0.95, 1.0];
        for mult in multipliers {
            assert!(mult >= 0.0 && mult <= 1.0);
        }
    }

    #[test]
    fn test_empty_reasoning_paths() {
        let paths: Vec<ReasoningPath> = vec![];
        assert!(paths.is_empty());
    }

    #[test]
    fn test_single_hop_reasoning() {
        let path = vec!["e1".to_string(), "e2".to_string()];
        assert_eq!(path.len(), 2);
    }

    #[test]
    fn test_multi_hop_reasoning() {
        let path = vec![
            "e1".to_string(),
            "e2".to_string(),
            "e3".to_string(),
            "e4".to_string(),
        ];
        assert_eq!(path.len(), 4);
    }

    #[test]
    fn test_relation_chain_length() {
        let relations = vec!["depends_on".to_string(), "uses".to_string()];
        assert_eq!(relations.len(), 2);
    }

    #[test]
    fn test_inference_deduplication() {
        let facts = vec![
            InferredFact {
                source_id: "e1".to_string(),
                source_name: "E1".to_string(),
                target_id: "e2".to_string(),
                target_name: "E2".to_string(),
                relation_type: "related".to_string(),
                confidence: 0.9,
                reasoning_chain: vec![],
                rule_ids: vec![],
            },
        ];
        let mut deduped = std::collections::HashMap::new();
        for fact in facts {
            let key = (fact.source_id.clone(), fact.target_id.clone(), fact.relation_type.clone());
            deduped.insert(key, fact);
        }
        assert_eq!(deduped.len(), 1);
    }

    #[test]
    fn test_transitive_closure_empty() {
        let closure = TransitiveClosure {
            source_id: "e1".to_string(),
            reachable: vec![],
            entity_count: 0,
            edge_count: 0,
        };
        assert_eq!(closure.reachable.len(), 0);
    }

    #[test]
    fn test_transitive_closure_single_hop() {
        let reachable = vec![
            ReachableEntity {
                entity_id: "e2".to_string(),
                entity_name: "E2".to_string(),
                relation_type: "depends_on".to_string(),
                confidence: 0.95,
                distance: 1,
            },
        ];
        assert_eq!(reachable.len(), 1);
        assert_eq!(reachable[0].distance, 1);
    }

    #[test]
    fn test_transitive_closure_multi_hop() {
        let reachable = vec![
            ReachableEntity {
                entity_id: "e2".to_string(),
                entity_name: "E2".to_string(),
                relation_type: "depends_on".to_string(),
                confidence: 0.95,
                distance: 1,
            },
            ReachableEntity {
                entity_id: "e3".to_string(),
                entity_name: "E3".to_string(),
                relation_type: "depends_on".to_string(),
                confidence: 0.90,
                distance: 2,
            },
        ];
        assert_eq!(reachable.len(), 2);
        assert!(reachable[1].confidence < reachable[0].confidence);
    }

    #[test]
    fn test_serialization_inferred_fact() {
        let fact = InferredFact {
            source_id: "e1".to_string(),
            source_name: "E1".to_string(),
            target_id: "e2".to_string(),
            target_name: "E2".to_string(),
            relation_type: "related".to_string(),
            confidence: 0.81,
            reasoning_chain: vec!["e1 --depends_on→ e2".to_string()],
            rule_ids: vec!["r1".to_string()],
        };
        let json = serde_json::to_string(&fact).unwrap();
        assert!(json.contains("0.81"));
    }

    #[test]
    fn test_serialization_reasoning_path() {
        let path = ReasoningPath {
            path: vec!["e1".to_string(), "e2".to_string()],
            relations: vec!["depends_on".to_string()],
            confidence: 0.9,
            step_count: 2,
        };
        let json = serde_json::to_string(&path).unwrap();
        assert!(json.contains("0.9"));
    }
}