poimen-memory/crates/mem-ingest/src/memorability_gate.rs

//! Memorability Gate: Filter extraction based on graph context
//!
//! Decides whether entities/facts are "worth remembering" by consulting GRM.
//! Configurable thresholds for different decision strategies.
//!
//! CRAP: 12 (Straightforward filtering + thresholds)
//! SOLID: Single responsibility (gate logic), delegates to retriever
//! DRY: Reuses GrmConfig and decision types

use anyhow::Result;
use serde::{Deserialize, Serialize};
use tracing::{debug, info};

use crate::grm_retriever::{
    EntityContext, FactContext, GraphContextRetriever, MemorabilityDecision, GrmConfig, MockGrmRetriever,
};
use mem_core::entity::{Entity, EntityType};
use mem_core::edge::Edge;

/// Entity filtering result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FilteredEntity {
    pub entity: Entity,
    pub context: EntityContext,
    pub filtered: bool,  // true = dropped by GRM gate
    pub reason: String,
}

/// Fact filtering result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FilteredFact {
    pub edge: Edge,
    pub context: FactContext,
    pub filtered: bool,  // true = dropped by GRM gate
    pub reason: String,
    pub requires_review: bool,  // true = queue for human verification
}

/// Memorability Gate
pub struct MemorabilityGate {
    config: GrmConfig,
    retriever: Box<dyn GraphContextRetriever>,
}

impl MemorabilityGate {
    /// Create gate with custom retriever (for testing or custom backends)
    pub fn new(config: GrmConfig, retriever: Box<dyn GraphContextRetriever>) -> Self {
        Self { config, retriever }
    }

    /// Create gate with mock retriever (everything passes)
    pub fn with_mock(config: GrmConfig) -> Self {
        Self {
            config,
            retriever: Box::new(MockGrmRetriever),
        }
    }

    /// Check if GRM gate is enabled
    pub fn is_enabled(&self) -> bool {
        self.config.enabled
    }

    /// Filter entity through GRM gate
    pub async fn filter_entity(&self, entity: &Entity) -> Result<FilteredEntity> {
        if !self.config.enabled {
            debug!("GRM gate disabled, passing entity: {}", entity.name);
            return Ok(FilteredEntity {
                entity: entity.clone(),
                context: EntityContext {
                    entity_name: entity.name.clone(),
                    matched_entity_id: None,
                    related_entities: vec![],
                    related_edges_count: 0,
                    summary: String::new(),
                    memorability_score: 1.0,
                    decision: MemorabilityDecision::Keep,
                    reasoning: "GRM gate disabled".to_string(),
                },
                filtered: false,
                reason: "GRM disabled".to_string(),
            });
        }

        debug!("GRM gate: filtering entity {}", entity.name);
        let context = self.retriever.get_entity_context(&entity.name).await?;

        let (filtered, reason) = match context.decision {
            MemorabilityDecision::Keep => {
                if context.matched_entity_id.is_some() {
                    (true, format!("Existing entity (merge required)"))
                } else {
                    (false, format!("New entity (score: {:.2})", context.memorability_score))
                }
            }
            MemorabilityDecision::Drop => {
                (true, format!("Noise/irrelevant (score: {:.2})", context.memorability_score))
            }
            MemorabilityDecision::ReviewQueue => {
                (false, format!("Low confidence, queued for review (score: {:.2})", context.memorability_score))
            }
            MemorabilityDecision::Merge => {
                (true, format!("Duplicate, requires merge (score: {:.2})", context.memorability_score))
            }
        };

        info!(
            "GRM entity filter: {} → filtered={} ({})",
            entity.name, filtered, reason
        );

        Ok(FilteredEntity {
            entity: entity.clone(),
            context,
            filtered,
            reason,
        })
    }

    /// Filter fact through GRM gate
    pub async fn filter_fact(
        &self,
        edge: &Edge,
        source_name: Option<&str>,
        target_name: Option<&str>,
    ) -> Result<FilteredFact> {
        if !self.config.enabled {
            debug!("GRM gate disabled, passing fact: {}", edge.fact);
            return Ok(FilteredFact {
                edge: edge.clone(),
                context: FactContext {
                    similar_facts_found: 0,
                    contradictory_facts_found: 0,
                    related_entities_coverage: 1.0,
                    memorability_score: 1.0,
                    decision: MemorabilityDecision::Keep,
                    reasoning: "GRM gate disabled".to_string(),
                },
                filtered: false,
                reason: "GRM disabled".to_string(),
                requires_review: false,
            });
        }

        debug!("GRM gate: filtering fact {}", edge.fact);
        let context = self.retriever
            .get_fact_context(
                &edge.source_entity_id,
                &edge.target_entity_id,
                &edge.relation_type,
                &edge.fact,
            )
            .await?;

        let (filtered, requires_review, reason) = match context.decision {
            MemorabilityDecision::Keep => {
                (false, false, format!("Novel fact (score: {:.2})", context.memorability_score))
            }
            MemorabilityDecision::Drop => {
                (true, false, format!("Redundant/noise (score: {:.2})", context.memorability_score))
            }
            MemorabilityDecision::ReviewQueue => {
                (false, true, format!("Low confidence, queued for review (score: {:.2})", context.memorability_score))
            }
            MemorabilityDecision::Merge => {
                (true, false, format!("Duplicate, requires merge (score: {:.2})", context.memorability_score))
            }
        };

        info!(
            "GRM fact filter: {} → {} → filtered={} requires_review={} ({})",
            source_name.unwrap_or("?"),
            target_name.unwrap_or("?"),
            filtered,
            requires_review,
            reason
        );

        Ok(FilteredFact {
            edge: edge.clone(),
            context,
            filtered,
            reason,
            requires_review,
        })
    }

    /// Batch filter entities
    pub async fn filter_entities(&self, entities: &[Entity]) -> Result<Vec<FilteredEntity>> {
        let mut results = Vec::new();
        for entity in entities {
            results.push(self.filter_entity(entity).await?);
        }
        Ok(results)
    }

    /// Batch filter facts
    pub async fn filter_facts(
        &self,
        edges: &[Edge],
        source_names: Option<&[Option<String>]>,
        target_names: Option<&[Option<String>]>,
    ) -> Result<Vec<FilteredFact>> {
        let mut results = Vec::new();
        for (i, edge) in edges.iter().enumerate() {
            let source = source_names.and_then(|names| names.get(i).and_then(|n| n.as_deref()));
            let target = target_names.and_then(|names| names.get(i).and_then(|n| n.as_deref()));
            results.push(self.filter_fact(edge, source, target).await?);
        }
        Ok(results)
    }

    /// Get statistics about filtering results
    pub fn stats(filtered: &[FilteredEntity]) -> FilterStatistics {
        let total = filtered.len();
        let dropped = filtered.iter().filter(|f| f.filtered).count();
        let kept = total - dropped;
        let avg_score = filtered
            .iter()
            .map(|f| f.context.memorability_score)
            .sum::<f32>() / (total as f32).max(1.0);

        FilterStatistics {
            total,
            kept,
            dropped,
            drop_rate: (dropped as f32 / total as f32).clamp(0.0, 1.0),
            avg_memorability_score: avg_score,
        }
    }
}

/// Filter statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FilterStatistics {
    pub total: usize,
    pub kept: usize,
    pub dropped: usize,
    pub drop_rate: f32,
    pub avg_memorability_score: f32,
}

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

    fn create_test_entity(name: &str) -> Entity {
        Entity::new("poimen", name, EntityType::Person)
    }

    fn create_test_edge(source: &str, target: &str, fact: &str) -> Edge {
        Edge::new("poimen", source, target, "USES", fact)
    }

    #[tokio::test]
    async fn test_gate_disabled() {
        let config = GrmConfig {
            enabled: false,
            ..Default::default()
        };
        let gate = MemorabilityGate::with_mock(config);
        
        let entity = create_test_entity("Rock");
        let result = gate.filter_entity(&entity).await.unwrap();
        
        assert!(!result.filtered);
        assert_eq!(result.reason, "GRM disabled");
    }

    #[tokio::test]
    async fn test_gate_enabled_known_entity() {
        let config = GrmConfig {
            enabled: true,
            entity_memorability_threshold: 0.75,
            ..Default::default()
        };
        let gate = MemorabilityGate::with_mock(config);
        
        let entity = create_test_entity("Rock");
        let result = gate.filter_entity(&entity).await.unwrap();
        
        // With mock retriever, entity "Rock" has high score
        assert_eq!(result.context.decision, MemorabilityDecision::Keep);
    }

    #[tokio::test]
    async fn test_gate_enabled_new_entity() {
        let config = GrmConfig {
            enabled: true,
            entity_memorability_threshold: 0.75,
            ..Default::default()
        };
        let gate = MemorabilityGate::with_mock(config);
        
        let entity = create_test_entity("UnknownPerson");
        let result = gate.filter_entity(&entity).await.unwrap();
        
        // With mock retriever, all entities get KEEP decision
        assert_eq!(result.context.decision, MemorabilityDecision::Keep);
    }

    #[tokio::test]
    async fn test_gate_filter_fact_disabled() {
        let config = GrmConfig {
            enabled: false,
            ..Default::default()
        };
        let gate = MemorabilityGate::with_mock(config);
        
        let edge = create_test_edge("entity-1", "entity-2", "Rock uses Kubernetes");
        let result = gate.filter_fact(&edge, Some("Rock"), Some("Kubernetes")).await.unwrap();
        
        assert!(!result.filtered);
        assert!(!result.requires_review);
    }

    #[tokio::test]
    async fn test_gate_batch_filter_entities() {
        let config = GrmConfig {
            enabled: true,
            ..Default::default()
        };
        let gate = MemorabilityGate::with_mock(config);
        
        let entities = vec![
            create_test_entity("Rock"),
            create_test_entity("Kubernetes"),
            create_test_entity("ArgoCD"),
        ];
        
        let results = gate.filter_entities(&entities).await.unwrap();
        assert_eq!(results.len(), 3);
    }

    #[test]
    fn test_filter_statistics() {
        let filtered = vec![
            FilteredEntity {
                entity: create_test_entity("A"),
                context: EntityContext {
                    entity_name: "A".to_string(),
                    matched_entity_id: None,
                    related_entities: vec![],
                    related_edges_count: 0,
                    summary: String::new(),
                    memorability_score: 0.9,
                    decision: MemorabilityDecision::Keep,
                    reasoning: String::new(),
                },
                filtered: false,
                reason: String::new(),
            },
            FilteredEntity {
                entity: create_test_entity("B"),
                context: EntityContext {
                    entity_name: "B".to_string(),
                    matched_entity_id: None,
                    related_entities: vec![],
                    related_edges_count: 0,
                    summary: String::new(),
                    memorability_score: 0.3,
                    decision: MemorabilityDecision::Drop,
                    reasoning: String::new(),
                },
                filtered: true,
                reason: String::new(),
            },
        ];
        
        let stats = MemorabilityGate::stats(&filtered);
        assert_eq!(stats.total, 2);
        assert_eq!(stats.kept, 1);
        assert_eq!(stats.dropped, 1);
        assert_eq!(stats.drop_rate, 0.5);
    }
}