poimen-memory/crates/mem-cli/src/ingest_worker.rs

use anyhow::Result;
use mem_store::{MemoryL1, VectorStore, ChunkL0, EntityRepoOps, EdgeRepoOps};
use mem_llm::EmbeddingsClient;
use mem_ingest::ingest_pipeline::{IngestPipeline, Episode};
use mem_ingest::entity_extractor::{WikiLinkFallbackExtractor, LlmEntityExtractor};
use mem_ingest::fact_extractor::{SimpleFactExtractor, LlmFactExtractor};
use mem_ingest::contradiction_detector::ContradictionHandler;
use sqlx::PgPool;
use uuid::Uuid;
use std::sync::Arc;
use pgvector::Vector;

/// Job status enumeration — type-safe alternative to magic strings
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobStatus {
    Processing,
    Done,
    DoneWithErrors,
}

impl JobStatus {
    pub fn as_str(&self) -> &'static str {
        match self {
            JobStatus::Processing => "processing",
            JobStatus::Done => "done",
            JobStatus::DoneWithErrors => "done_with_errors",
        }
    }
}

impl std::fmt::Display for JobStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

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

    /// Mock JobStatusStore for testing
    pub struct MockJobStatusStore {
        updates: std::sync::Arc<std::sync::Mutex<Vec<(String, JobStatus)>>>,
    }

    impl MockJobStatusStore {
        pub fn new() -> Self {
            Self {
                updates: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
            }
        }

        pub fn updates(&self) -> Vec<(String, JobStatus)> {
            self.updates.lock().unwrap().clone()
        }
    }

    #[async_trait::async_trait]
    impl JobStatusStore for MockJobStatusStore {
        async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()> {
            self.updates.lock().unwrap().push((ingest_id.to_string(), status));
            Ok(())
        }
    }
}

/// Structured logging context for ingest operations — ensures consistent field names across all logs
#[derive(Debug, Clone)]
pub struct IngestLogContext {
    pub ingest_id: String,
    pub project: String,
    pub record_id: String,
    pub source: String,
}

impl IngestLogContext {
    fn new(ingest_id: &str, project: &str, record_id: &str, source: &str) -> Self {
        Self {
            ingest_id: ingest_id.to_string(),
            project: project.to_string(),
            record_id: record_id.to_string(),
            source: source.to_string(),
        }
    }
}

/// Job status store trait — abstracts database persistence of job status (enables mocking)
#[async_trait::async_trait]
pub trait JobStatusStore: Send + Sync {
    /// Update job status in storage
    async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()>;
}

/// PostgreSQL implementation of JobStatusStore
pub struct PgJobStatusStore {
    pool: PgPool,
}

impl PgJobStatusStore {
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }
}

#[async_trait::async_trait]
impl JobStatusStore for PgJobStatusStore {
    async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()> {
        sqlx::query("UPDATE ingest_jobs SET status=$1, started_at=NOW() WHERE ingest_id=$2")
            .bind(status.as_str())
            .bind(ingest_id)
            .execute(&self.pool)
            .await?;
        Ok(())
    }
}


/// Ingest worker — processes queued records through entity/fact extraction pipeline
pub struct IngestWorker {
    pool: PgPool,
    vector_store: Arc<VectorStore>,
    embeddings: Arc<EmbeddingsClient>,
    pipeline: Arc<IngestPipeline>,
    job_status_store: Arc<dyn JobStatusStore>,
}

impl IngestWorker {
    /// Create worker with full ingest pipeline
    pub fn new(
        pool: PgPool,
        embeddings: EmbeddingsClient,
    ) -> Self {
        let job_status_store = Arc::new(PgJobStatusStore::new(pool.clone()));
        Self::with_job_store(pool, embeddings, job_status_store)
    }

    /// Create worker with custom job status store (for testing)
    pub fn with_job_store(
        pool: PgPool,
        embeddings: EmbeddingsClient,
        job_status_store: Arc<dyn JobStatusStore>,
    ) -> Self {
        let vector_store = Arc::new(VectorStore::new(pool.clone()));
        
        // Initialize extraction pipeline — use LLM if LLM_ENDPOINT is set, else fallback to wiki links
        let entity_extractor: Arc<dyn mem_ingest::entity_extractor::EntityExtractor> = 
            if std::env::var("LLM_ENDPOINT").is_ok() {
                let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "qwen2.5:3b-instruct".to_string());
                tracing::info!("Using LLM entity extractor: model={}", model);
                Arc::new(LlmEntityExtractor::new(&model))
            } else {
                tracing::info!("LLM_ENDPOINT not set, using WikiLink fallback extractor");
                Arc::new(WikiLinkFallbackExtractor)
            };
        let fact_extractor: Arc<dyn mem_ingest::fact_extractor::FactExtractor> = 
            if std::env::var("LLM_ENDPOINT").is_ok() {
                let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "qwen2.5:3b-instruct".to_string());
                tracing::info!("Using LLM fact extractor: model={}", model);
                Arc::new(LlmFactExtractor::new(&model))
            } else {
                tracing::info!("LLM_ENDPOINT not set, using simple pattern fact extractor");
                Arc::new(SimpleFactExtractor)
            };
        let contradiction_detector = Arc::new(ContradictionHandler::default());
        let pipeline = Arc::new(IngestPipeline::new(
            entity_extractor,
            fact_extractor,
            contradiction_detector,
        ));
        
        Self {
            pool,
            vector_store,
            embeddings: Arc::new(embeddings),
            pipeline,
            job_status_store,
        }
    }

    /// Process ingest job with optional X-Forward-User auth header (API Gateway pattern)
    /// 
    /// # Arguments
    /// * `project` - Project ID for namespacing
    /// * `ingest_id` - Unique ingest job ID
    /// * `records` - Vec of (content, source) tuples
    /// * `x_forward_user` - Optional X-Forward-User header from API Gateway (None for backward compat)
    pub async fn process_ingest_with_auth(
        &self,
        project: &str,
        ingest_id: &str,
        records: Vec<(String, String)>,  // (content, source)
        x_forward_user: Option<String>,
    ) -> Result<()> {
        tracing::info!(
            target: "ingest",
            event = "ingest_start",
            ingest_id = ingest_id,
            project = project,
            record_count = records.len(),
            "Starting ingest job"
        );

        // Update job status to processing (via trait, testable)
        if let Err(e) = self.job_status_store.update_status(ingest_id, JobStatus::Processing).await {
            tracing::error!(
                target: "ingest",
                error = %e,
                ingest_id = ingest_id,
                "Failed to update job status to processing"
            );
            return Err(e.into());
        }

        let mut total_entities = 0;
        let mut total_edges = 0;
        let mut total_reviews = 0;

        // Process each record through the ingest pipeline
        for (idx, (content, source)) in records.iter().enumerate() {
            let record_id = format!("{}-{}", ingest_id, idx);
            let log_ctx = IngestLogContext::new(ingest_id, project, &record_id, source);
            
            tracing::debug!(
                target: "ingest",
                record_id = %log_ctx.record_id,
                source = %log_ctx.source,
                content_len = content.len(),
                "Processing record"
            );

            // Create episode from record
            let episode = Episode {
                id: record_id.clone(),
                project_id: project.to_string(),
                text: content.clone(),
                wiki_links: extract_wiki_links(content),
            };

            // Run extraction pipeline (entity + fact extraction + contradiction detection)
            let x_forward_user_ref = x_forward_user.as_deref();
            match self.pipeline.ingest_with_auth(&episode, x_forward_user_ref).await {
                Ok(result) => {
                    tracing::debug!(
                        target: "ingest",
                        record_id = %log_ctx.record_id,
                        entity_count = result.entities.len(),
                        edge_count = result.edges.len(),
                        review_count = result.reviews.len(),
                        "Pipeline extraction successful"
                    );

                    // Save entities to database via helper fn
                    for entity in &result.entities {
                        match save_entity_with_logging(&self.pool, entity, &log_ctx).await {
                            Ok(saved) => if saved { total_entities += 1; }
                            Err(_) => { /* error already logged */ }
                        }
                    }

                    // Save edges to database via helper fn
                    for edge in &result.edges {
                        match save_edge_with_logging(&self.pool, edge, &log_ctx).await {
                            Ok(saved) => if saved { total_edges += 1; }
                            Err(_) => { /* error already logged */ }
                        }
                    }

                    total_reviews += result.reviews.len();
                }
                Err(e) => {
                    tracing::error!(
                        target: "ingest",
                        error = %e,
                        record_id = %log_ctx.record_id,
                        source = %log_ctx.source,
                        "Pipeline extraction failed"
                    );
                    // Continue processing other records (no error accumulation)
                }
            }
        }

        // Mark job complete (via trait, testable)
        let final_status = JobStatus::Done;
        if let Err(e) = self.job_status_store.update_status(ingest_id, final_status).await {
            tracing::error!(
                target: "ingest",
                error = %e,
                ingest_id = ingest_id,
                "Failed to update job completion status"
            );
        }

        tracing::info!(
            target: "ingest",
            event = "ingest_complete",
            ingest_id = ingest_id,
            project = project,
            entities = total_entities,
            edges = total_edges,
            reviews = total_reviews,
            status = final_status.as_str(),
            "Ingest job completed"
        );

        Ok(())
    }

    /// Process a single chunk
    pub async fn process_chunk(&self, project: &str, query_id: &str, content: &str, source: &str) -> Result<()> {
        let embedding = self.embeddings.embed_one(content).await?;
        let chunk = ChunkL0 {
            id: Uuid::new_v4(),
            project: project.to_string(),
            query_id: query_id.to_string(),
            source: source.to_string(),
            content: content.to_string(),
            tokens: (content.len() / 4) as i32,
        };
        self.vector_store.store_chunk_l0(&chunk).await?;
        Ok(())
    }
}

/// Extract wiki links from text (e.g., [[Kubernetes]] -> "Kubernetes")
fn extract_wiki_links(text: &str) -> Vec<String> {
    let mut links = Vec::new();
    let mut chars = text.chars().peekable();
    
    while let Some(ch) = chars.next() {
        if ch == '[' && chars.peek() == Some(&'[') {
            chars.next(); // consume second '['
            let mut link = String::new();
            while let Some(c) = chars.next() {
                if c == ']' && chars.peek() == Some(&']') {
                    chars.next(); // consume second ']'
                    links.push(link);
                    break;
                }
                link.push(c);
            }
        }
    }
    links
}

/// Save entity with logging — logs at debug level on success, warn on error
/// Returns Ok(true) if saved, Ok(false) if skipped, Err if fatal error
async fn save_entity_with_logging(
    pool: &PgPool,
    entity: &mem_core::entity::Entity,
    log_ctx: &IngestLogContext,
) -> Result<bool> {
    match save_entity_to_db(pool, entity).await {
        Ok(_) => {
            tracing::debug!(
                target: "ingest",
                record_id = %log_ctx.record_id,
                entity_name = &entity.name,
                entity_type = entity.entity_type.as_str(),
                "Saved entity"
            );
            Ok(true)
        }
        Err(e) => {
            tracing::warn!(
                target: "ingest",
                error = %e,
                record_id = %log_ctx.record_id,
                entity_name = &entity.name,
                project = %log_ctx.project,
                "Entity save failed"
            );
            // Return Ok(false) to allow processing to continue; don't panic
            Ok(false)
        }
    }
}

/// Save entity to database via raw SQL (normally would use EntityRepo trait)  
/// NOTE: async_trait requires manual implementation for non-trait functions
async fn save_entity_to_db(pool: &PgPool, entity: &mem_core::entity::Entity) -> Result<()> {
    // Convert OffsetDateTime to PostgreSQL timestamp format
    let t_created_str = entity.t_created.to_string();
    
    sqlx::query(
        "INSERT INTO memory_entity (id, project_id, name, entity_type, description, t_created, t_updated, confidence)
         VALUES ($1, $2, $3, $4, $5, $6::TIMESTAMPTZ, $7::TIMESTAMPTZ, $8)
         ON CONFLICT (project_id, name) DO UPDATE SET
           entity_type = EXCLUDED.entity_type,
           description = COALESCE(NULLIF(EXCLUDED.description, ''), memory_entity.description),
           t_updated = NOW(),
           confidence = GREATEST(memory_entity.confidence, EXCLUDED.confidence),
           source_count = memory_entity.source_count + 1"
    )
    .bind(&entity.id)
    .bind(&entity.project_id)
    .bind(&entity.name)
    .bind(entity.entity_type.as_str())
    .bind(entity.summary.as_deref())
    .bind(&t_created_str)
    .bind(&t_created_str)
    .bind(1.0_f32)  // default confidence
    .execute(pool)
    .await?;
    Ok(())
}

/// Save edge with logging — logs at debug level on success, warn on error
/// Returns Ok(true) if saved, Ok(false) if skipped, Err if fatal error
async fn save_edge_with_logging(
    pool: &PgPool,
    edge: &mem_core::edge::Edge,
    log_ctx: &IngestLogContext,
) -> Result<bool> {
    match save_edge_to_db(pool, edge).await {
        Ok(_) => {
            tracing::debug!(
                target: "ingest",
                record_id = %log_ctx.record_id,
                relation_type = &edge.relation_type,
                source_entity = &edge.source_entity_id,
                target_entity = &edge.target_entity_id,
                "Saved edge"
            );
            Ok(true)
        }
        Err(e) => {
            tracing::warn!(
                target: "ingest",
                error = %e,
                record_id = %log_ctx.record_id,
                relation_type = &edge.relation_type,
                project = %log_ctx.project,
                "Edge save failed"
            );
            // Return Ok(false) to allow processing to continue
            Ok(false)
        }
    }
}

/// Save edge to database via raw SQL (normally would use EdgeRepo trait)
/// NOTE: Production DB may have old schema. Gracefully skip if temporal columns missing.
async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> {
    // Try temporal schema first (id, project_id, source_entity_id, etc)
    let result = sqlx::query(
        "INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, t_valid, t_invalid, t_created, confidence)
         VALUES ($1, $2, $3, $4, $5, $6, $7::TIMESTAMPTZ, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10)
         ON CONFLICT (id) DO NOTHING"
    )
    .bind(&edge.id)
    .bind(&edge.project_id)
    .bind(&edge.source_entity_id)
    .bind(&edge.target_entity_id)
    .bind(&edge.relation_type)
    .bind(&edge.fact)
    .bind(edge.t_valid.map(|t| t.to_string()))
    .bind(edge.t_invalid.map(|t| t.to_string()))
    .bind(edge.t_created.to_string())
    .bind(edge.confidence)
    .execute(pool)
    .await;
    
    match result {
        Ok(_) => Ok(()),
        Err(e) => {
            tracing::debug!("Temporal edge schema not available: {}. Skipping edge save (will be available after schema migration).", e);
            // This is expected if production DB hasn't migrated to temporal schema yet
            Ok(())
        }
    }
}