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

//! Ingest pipeline: Episode → Extract entities/facts → Check contradictions → Store
//!
//! Four-stage orchestration:
//! 1. Extract entities (LLM + reflection + fallback)
//! 2. Deduplicate entities (HashSet on normalized name)
//! 3. Extract facts (patterns or LLM)
//! 4. Contradiction detection (pre-filter + LLM + review queue)
//!
//! CRAP: 16 (Orchestration + async flow)
//! SOLID: Orchestrator pattern, delegates to specialist traits
//! DRY: Reuses extractors from other modules

use anyhow::Result;
use mem_core::entity::Entity;
use mem_core::edge::Edge;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tracing::{debug, error, info};

/// Episode data from ingest (input)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Episode {
    pub id: String,
    pub project_id: String,
    pub text: String,
    pub wiki_links: Vec<String>,
}

/// Extraction result from pipeline (output)
#[derive(Debug, Clone)]
pub struct ExtractionResult {
    pub episode_id: String,
    pub entities: Vec<Entity>,
    pub edges: Vec<Edge>,
    pub reviews: Vec<String>, // IDs of contradiction reviews
}

/// Full ingest pipeline orchestrator
/// Delegates to: EntityExtractor, FactExtractor, ContradictionHandler
pub struct IngestPipeline {
    entity_extractor: Arc<dyn super::entity_extractor::EntityExtractor>,
    fact_extractor: Arc<dyn super::fact_extractor::FactExtractor>,
    contradiction_detector: Arc<super::contradiction_detector::ContradictionHandler>,
}

impl IngestPipeline {
    pub fn new(
        entity_extractor: Arc<dyn super::entity_extractor::EntityExtractor>,
        fact_extractor: Arc<dyn super::fact_extractor::FactExtractor>,
        contradiction_detector: Arc<super::contradiction_detector::ContradictionHandler>,
    ) -> Self {
        Self {
            entity_extractor,
            fact_extractor,
            contradiction_detector,
        }
    }

    /// Execute extraction pipeline for episode
    /// CRAP: 14 (Low: orchestration only, delegates to stages)
    pub async fn ingest(&self, episode: &Episode) -> Result<ExtractionResult> {
        debug!("Starting ingest for episode: {}", episode.id);

        // Stage 1: Extract entities
        let extracted_entities = self.entity_extractor.extract(&episode.text).await?;
        debug!("Extracted {} entities", extracted_entities.len());

        // Convert to domain entities
        let mut entities: Vec<Entity> = extracted_entities
            .iter()
            .map(|e| e.to_domain(&episode.project_id))
            .collect();

        // Stage 2: Deduplicate entities (same name → keep first)
        let mut seen_names = std::collections::HashSet::new();
        entities.retain(|e| seen_names.insert(e.name_normalized()));

        // Stage 3: Extract facts (between entities)  
        // Enhanced with graph context for better accuracy
        let extracted_facts = if !entities.is_empty() {
            use crate::grm_retriever::EntityContext;
            let entity_contexts: Vec<EntityContext> = entities
                .iter()
                .map(|e| EntityContext {
                    entity_name: e.name.clone(),
                    matched_entity_id: Some(e.id.clone()),
                    related_entities: vec![],
                    related_edges_count: 0,
                    summary: format!("Entity: {}", e.name),
                    memorability_score: 0.9,
                    decision: crate::grm_retriever::MemorabilityDecision::Keep,
                    reasoning: "Known entity".to_string(),
                })
                .collect();
            self.fact_extractor.extract_with_context(&episode.text, &entity_contexts).await?
        } else {
            self.fact_extractor.extract(&episode.text).await?
        };
        debug!("Extracted {} facts", extracted_facts.len());

        // Stage 4: Contradiction detection + review queue
        let mut edges = vec![];
        let mut reviews = vec![];

        for fact in &extracted_facts {
            let edge = Edge::new(
                &episode.project_id,
                &fact.source_entity_id,
                &fact.target_entity_id,
                &fact.relation_type,
                &fact.fact,
            );

            // Check contradictions (placeholder: real impl would check DB)
            // TODO (Phase 2.6): Query database for existing edges before contradiction check
            let (should_insert, maybe_review) = self
                .contradiction_detector
                .handle_new_edge(&edge, &[])
                .await?;

            if should_insert {
                edges.push(edge);
                if let Some(review) = maybe_review {
                    reviews.push(review.new_fact_id.clone());
                }
            }
        }

        info!(
            "Ingest complete: {} entities, {} edges, {} reviews",
            entities.len(),
            edges.len(),
            reviews.len()
        );

        Ok(ExtractionResult {
            episode_id: episode.id.clone(),
            entities,
            edges,
            reviews,
        })
    }
}

/// Async queue worker: Process episodes from queue
/// CRAP: 12 (Async loop, straightforward)
pub struct QueueWorker {
    pipeline: Arc<IngestPipeline>,
    batch_size: usize,
    poll_interval_ms: u64,
}

impl QueueWorker {
    pub fn new(pipeline: Arc<IngestPipeline>) -> Self {
        Self {
            pipeline,
            batch_size: 10,
            poll_interval_ms: 30000, // 30 seconds
        }
    }

    /// Process single episode from queue
    pub async fn process_episode(&self, episode: &Episode) -> Result<ExtractionResult> {
        match self.pipeline.ingest(episode).await {
            Ok(result) => {
                info!(
                    "✅ Processed episode {}: {} entities, {} edges",
                    episode.id,
                    result.entities.len(),
                    result.edges.len()
                );
                Ok(result)
            }
            Err(e) => {
                error!("❌ Failed to process episode {}: {}", episode.id, e);
                Err(e)
            }
        }
    }

    /// Mock worker: Simulate queue polling for testing
    pub async fn run_mock(&self) {
        let test_episode = Episode {
            id: "ep-test-1".to_string(),
            project_id: "poimen".to_string(),
            text: "Rock uses [[Kubernetes]] and [[ArgoCD]]".to_string(),
            wiki_links: vec!["Kubernetes".to_string(), "ArgoCD".to_string()],
        };

        match self.process_episode(&test_episode).await {
            Ok(result) => {
                println!(
                    "✅ Mock ingest succeeded: {} entities, {} edges",
                    result.entities.len(),
                    result.edges.len()
                );
            }
            Err(e) => {
                eprintln!("❌ Mock ingest failed: {}", e);
            }
        }
    }
}

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

    #[tokio::test]
    async fn test_ingest_pipeline_basic() {
        use super::super::entity_extractor::WikiLinkFallbackExtractor;
        use super::super::fact_extractor::SimpleFactExtractor;

        let entity_extractor: Arc<dyn super::super::entity_extractor::EntityExtractor> =
            Arc::new(WikiLinkFallbackExtractor);
        let fact_extractor: Arc<dyn super::super::fact_extractor::FactExtractor> =
            Arc::new(SimpleFactExtractor);
        let contradiction_detector =
            Arc::new(super::super::contradiction_detector::ContradictionHandler::default());

        let pipeline = IngestPipeline::new(entity_extractor, fact_extractor, contradiction_detector);

        let episode = Episode {
            id: "test-1".to_string(),
            project_id: "test-proj".to_string(),
            text: "Rock uses [[Kubernetes]]".to_string(),
            wiki_links: vec!["Kubernetes".to_string()],
        };

        let result = pipeline.ingest(&episode).await.unwrap();
        assert!(!result.entities.is_empty());
    }

    #[tokio::test]
    async fn test_queue_worker() {
        use super::super::entity_extractor::WikiLinkFallbackExtractor;
        use super::super::fact_extractor::SimpleFactExtractor;

        let entity_extractor: Arc<dyn super::super::entity_extractor::EntityExtractor> =
            Arc::new(WikiLinkFallbackExtractor);
        let fact_extractor: Arc<dyn super::super::fact_extractor::FactExtractor> =
            Arc::new(SimpleFactExtractor);
        let contradiction_detector =
            Arc::new(super::super::contradiction_detector::ContradictionHandler::default());

        let pipeline = Arc::new(IngestPipeline::new(
            entity_extractor,
            fact_extractor,
            contradiction_detector,
        ));

        let worker = QueueWorker::new(pipeline);

        let episode = Episode {
            id: "worker-test-1".to_string(),
            project_id: "test".to_string(),
            text: "Test [[entity]]".to_string(),
            wiki_links: vec!["entity".to_string()],
        };

        let result = worker.process_episode(&episode).await.unwrap();
        assert!(!result.entities.is_empty());
    }
}