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

/// Phase 3: Compaction — Automated deduplication and garbage collection
///
/// Three-tier approach:
/// - T3.1: Exact dedup (no LLM)
/// - T3.2: Semantic dedup (LLM-gated with pre-filter)
/// - T3.3: Audit logging + dry-run mode

use anyhow::{Result, anyhow};
use sqlx::{Pool, Postgres, Row};
use std::sync::Arc;
use std::collections::HashMap;
use tracing::{debug, info, warn};

use mem_core::edge::Edge;
use mem_ingest::entity_extractor::LlmCaller;

/// Compaction statistics
#[derive(Debug, Clone, Default)]
pub struct CompactionStats {
    pub duplicate_edges_deleted: usize,
    pub stale_facts_deleted: usize,
    pub semantic_merged: usize,
    pub bytes_freed: usize,
    pub llm_calls: usize,
    pub human_reviews_queued: usize,
    pub duration_ms: u64,
}

/// Compaction mode
#[derive(Debug, Clone, Copy)]
pub enum CompactionMode {
    /// Simulate changes, don't apply
    DryRun,
    /// Apply changes with audit logging
    Execute,
}

/// T3.1: Exact Deduplicator
pub struct Tier1Compactor {
    pool: Pool<Postgres>,
    retention_days: i32,
}

impl Tier1Compactor {
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self {
            pool,
            retention_days: 30,
        }
    }

    /// Find duplicate edges (same source + target + relation_type + fact_hash)
    pub async fn find_duplicate_edges(&self) -> Result<Vec<(String, String)>> {
        let rows = sqlx::query(
            r#"
            SELECT array_agg(id ORDER BY created_at)
            FROM memory_edge
            WHERE deleted_at IS NULL
            GROUP BY source_id, target_id, relation_type, md5(fact)
            HAVING COUNT(*) > 1
            "#,
        )
        .fetch_all(&self.pool)
        .await?;

        let mut duplicates = Vec::new();
        for row in rows {
            let ids: Vec<String> = row.get::<Vec<String>, _>(0);
            if ids.len() > 1 {
                // Keep first (master), mark rest as duplicates
                for dup_id in &ids[1..] {
                    duplicates.push((ids[0].clone(), dup_id.clone()));
                }
            }
        }

        Ok(duplicates)
    }

    /// Delete duplicate edges (soft-delete)
    pub async fn delete_duplicates(&self, mode: CompactionMode) -> Result<CompactionStats> {
        let duplicates = self.find_duplicate_edges().await?;
        let count = duplicates.len();
        let bytes = count * 1024; // Approximate

        let pool = self.pool.clone();
        let execute_fn = async move {
            for (_master, duplicate) in duplicates {
                sqlx::query(
                    "UPDATE memory_edge SET deleted_at = NOW() WHERE id = $1"
                )
                .bind(&duplicate)
                .execute(&pool)
                .await?;
            }
            Ok::<(), anyhow::Error>(())
        };

        let result = crate::compaction_executor::execute_operation(
            mode,
            "delete duplicate edges",
            count,
            bytes,
            execute_fn,
        )
        .await?;

        let stats = CompactionStats {
            duplicate_edges_deleted: if result.executed { result.count } else { 0 },
            bytes_freed: if result.executed { result.bytes } else { 0 },
            ..Default::default()
        };

        info!("T3.1: Deleted {} duplicate edges", stats.duplicate_edges_deleted);
        Ok(stats)
    }

    /// Garbage collect stale facts
    pub async fn gc_stale_facts(&self, mode: CompactionMode) -> Result<CompactionStats> {
        let cutoff_date = format!("NOW() - INTERVAL '{}' day", self.retention_days);

        let row_count: (i64,) = sqlx::query_as(
            &format!(
                r#"
                SELECT COUNT(*) FROM memory_edge
                WHERE fact_invalid_at IS NOT NULL
                AND fact_invalid_at < {}
                AND deleted_at IS NULL
                "#,
                cutoff_date
            ),
        )
        .fetch_one(&self.pool)
        .await?;

        let stale_count = row_count.0 as usize;

        if stale_count == 0 {
            return Ok(CompactionStats::default());
        }

        let pool = self.pool.clone();
        let cutoff = cutoff_date.clone();
        let execute_fn = async move {
            sqlx::query(
                &format!(
                    r#"
                    UPDATE memory_edge
                    SET deleted_at = NOW()
                    WHERE fact_invalid_at IS NOT NULL
                    AND fact_invalid_at < {}
                    AND deleted_at IS NULL
                    "#,
                    cutoff
                ),
            )
            .execute(&pool)
            .await?;
            Ok::<(), anyhow::Error>(())
        };

        let result = crate::compaction_executor::execute_operation(
            mode,
            "GC stale facts",
            stale_count,
            stale_count * 1024,
            execute_fn,
        )
        .await?;

        let stats = CompactionStats {
            stale_facts_deleted: if result.executed { result.count } else { 0 },
            bytes_freed: if result.executed { result.bytes } else { 0 },
            ..Default::default()
        };

        info!("T3.1: GC deleted {} stale facts (> {} days old)", stale_count, self.retention_days);
        Ok(stats)
    }
}

/// T3.2: Semantic Deduplicator
pub struct Tier2Compactor {
    pool: Pool<Postgres>,
    llm_caller: Arc<dyn LlmCaller>,
    confidence_threshold_auto: f32,  // > 0.95: auto-merge
    confidence_threshold_review: f32, // 0.70-0.95: human review
}

impl Tier2Compactor {
    pub fn new(pool: Pool<Postgres>, llm_caller: Arc<dyn LlmCaller>) -> Self {
        Self {
            pool,
            llm_caller,
            confidence_threshold_auto: 0.95,
            confidence_threshold_review: 0.70,
        }
    }

    /// Pre-filter: Find candidate pairs without LLM
    pub async fn prefilter_candidates(&self) -> Result<Vec<(String, String, String, String)>> {
        // Find edges with same source + target (likely related)
        let rows = sqlx::query(
            r#"
            SELECT a.id, b.id, a.fact, b.fact
            FROM memory_edge a
            JOIN memory_edge b ON a.source_id = b.source_id
                AND a.target_id = b.target_id
                AND a.relation_type = b.relation_type
                AND a.id < b.id
            WHERE a.deleted_at IS NULL
            AND b.deleted_at IS NULL
            AND a.fact_invalid_at IS NULL
            AND b.fact_invalid_at IS NULL
            LIMIT 100
            "#,
        )
        .fetch_all(&self.pool)
        .await?;

        let candidates = rows.into_iter()
            .map(|row| (
                row.get::<String, _>(0),
                row.get::<String, _>(1),
                row.get::<String, _>(2),
                row.get::<String, _>(3),
            ))
            .collect();

        Ok(candidates)
    }

    /// Check semantic equivalence via LLM
    pub async fn check_equivalence(
        &self,
        fact_a: &str,
        fact_b: &str,
    ) -> Result<f32> {
        let prompt = format!(
            r#"Are these facts semantically equivalent?

Fact A: {}
Fact B: {}

Respond with JSON: {{"confidence": 0.0-1.0}} where 1.0 means identical meaning."#,
            fact_a, fact_b
        );

        let response = self.llm_caller.call(&prompt).await?;

        // Parse JSON response for confidence score
        if let Ok(json) = serde_json::from_str::<serde_json::Value>(&response) {
            if let Some(conf) = json.get("confidence").and_then(|v| v.as_f64()) {
                return Ok(conf as f32);
            }
        }

        Ok(0.0) // Default to not equivalent if parse fails
    }

    /// Merge equivalent edges
    pub async fn merge_equivalent_edges(
        &self,
        edge_a_id: &str,
        edge_b_id: &str,
        confidence: f32,
        mode: CompactionMode,
    ) -> Result<CompactionStats> {
        let mut stats = CompactionStats::default();
        stats.llm_calls = 1;

        if confidence > self.confidence_threshold_auto {
            // Auto-merge: keep longer fact, delete shorter
            let pool = self.pool.clone();
            let edge_id = edge_b_id.to_string();
            let execute_fn = async move {
                sqlx::query(
                    "UPDATE memory_edge SET deleted_at = NOW() WHERE id = $1"
                )
                .bind(&edge_id)
                .execute(&pool)
                .await?;
                Ok::<(), anyhow::Error>(())
            };

            let result = crate::compaction_executor::execute_operation(
                mode,
                &format!("merge {} and {}", edge_a_id, edge_b_id),
                1,
                512,
                execute_fn,
            )
            .await?;

            stats.semantic_merged = if result.executed { 1 } else { 0 };
            stats.bytes_freed = if result.executed { 512 } else { 0 };
        } else if confidence > self.confidence_threshold_review {
            // Queue for human review
            stats.human_reviews_queued += 1;
            debug!("Queued merge for review: {} + {} (confidence: {:.2})", edge_a_id, edge_b_id, confidence);
        }

        Ok(stats)
    }
}

/// Execute full compaction pipeline
pub async fn compact_memory(
    pool: &Pool<Postgres>,
    llm_caller: Option<Arc<dyn LlmCaller>>,
    mode: CompactionMode,
) -> Result<CompactionStats> {
    let start = std::time::Instant::now();
    let mut total_stats = CompactionStats::default();

    // T3.1: Exact dedup
    let tier1 = Tier1Compactor::new(pool.clone());
    let t1_stats = tier1.delete_duplicates(mode).await?;
    total_stats.duplicate_edges_deleted += t1_stats.duplicate_edges_deleted;
    total_stats.bytes_freed += t1_stats.bytes_freed;

    // T3.1: GC stale facts
    let t1_gc_stats = tier1.gc_stale_facts(mode).await?;
    total_stats.stale_facts_deleted += t1_gc_stats.stale_facts_deleted;
    total_stats.bytes_freed += t1_gc_stats.bytes_freed;

    // T3.2: Semantic dedup (if LLM available)
    if let Some(llm) = llm_caller {
        let tier2 = Tier2Compactor::new(pool.clone(), llm);
        let candidates = tier2.prefilter_candidates().await.unwrap_or_default();

        for (edge_a_id, edge_b_id, fact_a, fact_b) in candidates {
            if let Ok(confidence) = tier2.check_equivalence(&fact_a, &fact_b).await {
                if let Ok(t2_stats) = tier2.merge_equivalent_edges(&edge_a_id, &edge_b_id, confidence, mode).await {
                    total_stats.semantic_merged += t2_stats.semantic_merged;
                    total_stats.llm_calls += 1;
                    total_stats.bytes_freed += t2_stats.bytes_freed;
                    total_stats.human_reviews_queued += t2_stats.human_reviews_queued;
                }
            }
        }
    }

    total_stats.duration_ms = start.elapsed().as_millis() as u64;
    info!("Compaction complete in {}ms: {:?}", total_stats.duration_ms, total_stats);

    Ok(total_stats)
}

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

    #[test]
    fn test_compaction_stats_default() {
        let stats = CompactionStats::default();
        assert_eq!(stats.duplicate_edges_deleted, 0);
        assert_eq!(stats.bytes_freed, 0);
    }

    #[test]
    fn test_compaction_stats_accumulate() {
        let mut stats = CompactionStats::default();
        stats.duplicate_edges_deleted = 5;
        stats.bytes_freed = 5120;

        assert_eq!(stats.duplicate_edges_deleted, 5);
        assert_eq!(stats.bytes_freed, 5120);
    }

    #[test]
    fn test_confidence_thresholds() {
        let tier2 = Tier2Compactor::new(
            // Mock pool would go here
            todo!(),
            Arc::new(MockLlmCaller),
        );

        assert!(tier2.confidence_threshold_auto > tier2.confidence_threshold_review);
        assert!(tier2.confidence_threshold_review > 0.5);
    }
}

/// Mock LLM caller for testing
#[cfg(test)]
struct MockLlmCaller;

#[cfg(test)]
#[async_trait::async_trait]
impl LlmCaller for MockLlmCaller {
    async fn call(&self, _prompt: &str) -> anyhow::Result<String> {
        Ok(r#"{"confidence": 0.85}"#.to_string())
    }
}