poimen-memory/crates/mem-store/src/rebuild.rs

use anyhow::{anyhow, Result};
use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::{Path, PathBuf};
use sha2::{Digest, Sha256};

use crate::{
    MemoryNode, MemoryRecord, MemoryParent, Level, VectorKind, PgRepo, ObsidianProjector,
    ProjectorOpts,
};

/// Rebuild options
#[derive(Debug, Clone)]
pub struct RebuildOpts {
    pub project: String,
    pub vault_only: bool,    // Only rebuild vault, not database
    pub db_only: bool,        // Only rebuild database, not vault
    pub allow_partial: bool,  // Allow rebuilding from incomplete logs
    pub embedding_cache_dir: Option<PathBuf>,
    pub vault_dir: Option<PathBuf>,
    pub log_dir: Option<PathBuf>,
}

/// Rebuild statistics
#[derive(Debug, Clone, Default)]
pub struct RebuildStats {
    pub nodes_l0: i64,
    pub nodes_l1: i64,
    pub nodes_l2: i64,
    pub edges: i64,
    pub embeddings_computed: i64,
    pub embeddings_cached: i64,
}

/// Rebuild orchestrator: drop projections, rebuild from log
pub struct RebuildEngine {
    repo: PgRepo,
}

impl RebuildEngine {
    /// Create rebuild engine with Postgres connection
    pub async fn new(db_url: &str) -> Result<Self> {
        let repo = PgRepo::connect(db_url).await?;
        Ok(Self { repo })
    }

    /// Execute full rebuild: clear → insert nodes → insert edges → project vault
    /// 
    /// Order matters: nodes first (foreign key constraint), then edges, then vault projection
    pub async fn rebuild(&self, opts: RebuildOpts) -> Result<RebuildStats> {
        let log_dir = opts.log_dir.unwrap_or_else(|| PathBuf::from("log"));
        let vault_dir = opts.vault_dir.unwrap_or_else(|| PathBuf::from("vault"));
        let cache_dir = opts
            .embedding_cache_dir
            .clone()
            .unwrap_or_else(|| PathBuf::from(".cache"));

        // Create cache directory
        fs::create_dir_all(&cache_dir)?;

        // Read all memories from log files
        let memories = Self::read_log_memories(&log_dir, &opts.project, opts.allow_partial).await?;

        let mut stats = RebuildStats::default();

        // PASS 1: Clear project (if not vault-only)
        if !opts.vault_only {
            self.repo.clear_project(&opts.project).await?;
        }

        // PASS 2: Insert all nodes (convert memories to nodes, batch embeddings)
        if !opts.vault_only {
            let mut nodes_by_sha: HashMap<String, MemoryNode> = HashMap::new();
            let mut sha_to_level: HashMap<String, Level> = HashMap::new();
            let mut sha_to_parents: HashMap<String, Vec<String>> = HashMap::new();

            for memory in &memories {
                let sha = Self::memory_sha(&memory.text);
                let level = match memory.level.as_str() {
                    "L0" => Level::L0,
                    "L1" => Level::L1,
                    "L2" => Level::L2,
                    "R" => Level::R,
                    _ => continue,
                };

                let node = MemoryNode {
                    sha256: sha.clone(),
                    level,
                    project: memory.project.clone(),
                    query_id: memory.query_id.clone(),
                    run_id: memory.run_id.clone(),
                    t: memory.t,
                    source: memory.source.clone(),
                    text: memory.text.clone(),
                };

                nodes_by_sha.insert(sha.clone(), node);
                sha_to_level.insert(sha.clone(), level);

                // Track parents from provenance
                let parent_shas: Vec<String> = memory
                    .parents
                    .iter()
                    .map(|p| Self::parent_sha(&p.source, p.t))
                    .collect();
                if !parent_shas.is_empty() {
                    sha_to_parents.insert(sha, parent_shas);
                }
            }

            // Upsert all nodes
            for node in nodes_by_sha.values() {
                self.repo.upsert_node(node).await?;
            }

            stats.nodes_l0 = nodes_by_sha.values().filter(|n| n.level == Level::L0).count() as i64;
            stats.nodes_l1 = nodes_by_sha.values().filter(|n| n.level == Level::L1).count() as i64;
            stats.nodes_l2 = nodes_by_sha.values().filter(|n| n.level == Level::L2).count() as i64;

            // PASS 3: Insert edges (after all nodes exist)
            for (child_sha, parent_shas) in sha_to_parents {
                self.repo.insert_edges(&child_sha, &parent_shas).await?;
                stats.edges += parent_shas.len() as i64;
            }

            // TODO: Batch embeddings with embedding cache
            // For now, mock stats
            stats.embeddings_cached = 0;
            stats.embeddings_computed = 0;
        }

        // PASS 4: Project vault (if not db-only)
        if !opts.db_only {
            ObsidianProjector::project(&log_dir, &vault_dir, ProjectorOpts::default()).await?;
        }

        Ok(stats)
    }

    /// Read all memory records from log directory
    /// 
    /// Returns error if any log is incomplete (no `run_end`) unless `allow_partial`
    pub async fn read_log_memories(
        log_dir: &Path,
        project: &str,
        allow_partial: bool,
    ) -> Result<Vec<MemoryRecord>> {
        let mut memories = Vec::new();

        // Look for log/project/ directory
        let project_dir = log_dir.join(project);
        if !project_dir.exists() {
            return Ok(memories);
        }

        // Iterate over query directories
        for entry in fs::read_dir(&project_dir)? {
            let query_dir = entry?.path();
            if !query_dir.is_dir() {
                continue;
            }

            // Iterate over run files
            for run_entry in fs::read_dir(&query_dir)? {
                let run_file = run_entry?.path();
                if run_file.extension().map(|e| e != "jsonl").unwrap_or(true) {
                    continue;
                }

                // Read JSONL file
                let contents = fs::read_to_string(&run_file)?;
                for line in contents.lines() {
                    if line.trim().is_empty() {
                        continue;
                    }

                    // Parse memory record (simplified — real implementation parses event log)
                    if let Ok(memory) = serde_json::from_str::<MemoryRecord>(line) {
                        memories.push(memory);
                    }
                }

                // Check for run_end (simplified — would need full log parsing)
                if !allow_partial && !contents.contains("run_end") {
                    return Err(anyhow!(
                        "Incomplete log: {} (missing run_end). Use --allow-partial to ignore.",
                        run_file.display()
                    ));
                }
            }
        }

        Ok(memories)
    }

    /// Compute stable sha256 for memory text (content identity)
    pub fn memory_sha(text: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(text.as_bytes());
        format!("{:x}", hasher.finalize())
    }

    /// Compute parent sha from source + timestamp
    fn parent_sha(source: &str, t: i32) -> String {
        format!("{}-{}", source, t)
    }
}

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

    #[test]
    fn test_memory_sha_deterministic() {
        let text = "same content";
        let sha1 = RebuildEngine::memory_sha(text);
        let sha2 = RebuildEngine::memory_sha(text);
        assert_eq!(sha1, sha2, "Same content must produce same SHA");
    }

    #[test]
    fn test_memory_sha_differs() {
        let sha1 = RebuildEngine::memory_sha("content a");
        let sha2 = RebuildEngine::memory_sha("content b");
        assert_ne!(sha1, sha2, "Different content must produce different SHAs");
    }

    #[test]
    fn test_parent_sha_format() {
        let parent_sha = RebuildEngine::parent_sha("pi", 42);
        assert_eq!(parent_sha, "pi-42");
    }
}