leindex 2.0.3

//! Mmap + snapshot persistence for the index builder (extracted from `mod.rs`
//! to keep the parent module under the large-file gate).

use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use tracing::info;

use crate::graph::pdg::ProgramDependenceGraph;
use crate::search::search::{SearchEngine, SearchSnapshot};

use super::fragment;

// ============================================================================
// MMAP EMBEDDING PERSISTENCE (R10)
// ============================================================================

/// Persist all embeddings from the search engine to an mmap-backed binary file.
///
/// After indexing completes, call this to write a `.leindex/embeddings.bin`
/// file that can be memory-mapped for fast read-only access without loading
/// the full embedding matrix into heap memory.
pub(crate) fn persist_embeddings_to_mmap(
    search_engine: &SearchEngine,
    project_path: &Path,
) -> Result<()> {
    let path = crate::search::vector::mmap_embeddings_path(project_path);
    // Skip only when the on-disk mmap already describes the engine's current
    // row population. The old check (mmap-backed + file exists) skipped
    // FOREVER after the first write: an incremental reindex on a
    // mmap-hydrated engine grew the snapshot but never the mmap file, so
    // every later hydration failed the "mmap row count != snapshot
    // indexed_nodes" validation and took the full rebuild path.
    if search_engine.is_mmap_backed() && path.is_file() {
        let engine_rows = search_engine.node_count();
        let file_rows = crate::search::vector::MmapEmbeddingIndex::open(&path)
            .map(|index| index.len())
            .unwrap_or(usize::MAX);
        if file_rows == engine_rows {
            return Ok(());
        }
        info!(
            file_rows,
            engine_rows, "Rewriting mmap embeddings after row-count drift"
        );
    }
    let embeddings = search_engine.collect_embeddings();
    if embeddings.is_empty() {
        return Ok(());
    }
    crate::search::vector::write_mmap_embeddings(&path, &embeddings)
        .map_err(|e| anyhow::anyhow!("Failed to write mmap embeddings: {e}"))?;
    info!(
        count = embeddings.len(),
        path = %path.display(),
        "Persisted embeddings to mmap file"
    );
    Ok(())
}

fn search_snapshot_path(project_path: &Path) -> PathBuf {
    project_path.join(".leindex").join("search_snapshot.bin")
}

/// The (nodes, edges, fingerprint) identity of a PDG as it will be
/// RECONSTRUCTED from the published Pdg layer.
///
/// Snapshot and embedder freshness metadata must describe the persisted
/// graph, not the raw in-memory one: the in-memory graph may legitimately
/// hold several nodes sharing one `node_id` (external/import duplicates) and
/// parallel same-type edges, which the layer format collapses (last-wins,
/// matching the legacy SQL upserts). Identity recorded from the uncollapsed
/// graph can therefore never match a layer-hydrated graph — every cold
/// hydration would take the full TF-IDF rebuild path forever. Deriving the
/// identity through the same collapse the layer round-trip applies makes
/// save/load agreement structural instead of coincidental.
pub(crate) fn persisted_search_identity_from_graph(
    pdg: &crate::graph::pdg::ProgramDependenceGraph,
) -> (usize, usize, String) {
    let mut collapsed = crate::storage::generation::graph_codec::collapse_for_layer(pdg);
    // Mirror hydration exactly: `load_from_storage_inner_at` normalizes
    // external nodes BEFORE fingerprinting, so the identity must describe
    // the same post-normalization graph or the freshness check fails.
    super::normalize_external_nodes(&mut collapsed);
    let fingerprint = pdg_search_fingerprint(&collapsed);
    (collapsed.node_count(), collapsed.edge_count(), fingerprint)
}

/// Persist search metadata required for fast load_from_storage hydration.
pub(crate) fn persist_search_snapshot(
    search_engine: &SearchEngine,
    project_path: &Path,
    pdg_nodes: usize,
    pdg_edges: usize,
    pdg_fingerprint: String,
) -> Result<()> {
    let mut snapshot = search_engine.search_snapshot(pdg_nodes, pdg_edges, pdg_fingerprint);
    if snapshot.indexed_nodes == 0 {
        return Ok(());
    }

    // Fragment layer (Task 5): persist the fragment embedding matrix + root
    // hash alongside the snapshot so cold-start hydration can rebuild the
    // fragment vector index and validate invariant 8 (root-hash + row count).
    // Task 7: the mmap is written FIRST and the root LAST (the root is the
    // commit marker — a crash mid-persist leaves a mismatched/older root that
    // hydration rejects via `fragment_layer_generation_is_consistent`, so a
    // half-synced fragment tree never serves). The root generation is taken
    // from the sync manifest so `fragment_root.bin` and
    // `fragment_sync_manifest.bin` stay generation-aligned.
    let fragment_embeddings = search_engine.collect_fragment_embeddings();
    persist_fragment_embeddings_to_mmap(project_path, &fragment_embeddings)?;
    if !fragment_embeddings.is_empty() {
        let fragment_ids: Vec<String> = fragment_embeddings
            .iter()
            .map(|(id, _)| id.clone())
            .collect();
        snapshot.fragment_root_hash = Some(fragment::sync::compute_fragment_root_hash_from_ids(
            &fragment_ids,
        ));
        // CRITICAL: the manifest lives in `.leindex/`, so the lookup must use
        // the storage dir (like `load_fragment_root`) — using `project_path`
        // directly reads the wrong path, falls back to generation 0, and would
        // clobber the engine's gen-N root, permanently disabling the layer.
        let generation =
            fragment::sync::load_fragment_sync_manifest(&project_path.join(".leindex"))
                .ok()
                .flatten()
                .map(|manifest| manifest.generation)
                .unwrap_or(0);
        fragment::sync::persist_fragment_root_from_ids(project_path, &fragment_ids, generation)?;
    }

    let path = search_snapshot_path(project_path);
    // Identity sidecar. `SearchSnapshot` carries no timestamps — it is a pure
    // function of (PDG content, indexed node set, token dictionary, fragment
    // layer). The identity is a blake3 of the serialized snapshot itself, so
    // EVERY content-bearing field is covered: the earlier multi-field
    // identity omitted the per-node `token_ids` and the token dictionary,
    // and a comment-only edit inside a function body left every one of those
    // fields unchanged (same ids, byte ranges, counts, fingerprint), so the
    // rewrite was skipped and cold hydration kept serving the pre-edit
    // tokens until an unrelated structural change bumped the fingerprint.
    // When the identity matches what is already on disk, the ~11MB rewrite
    // is byte-equivalent and can be skipped. Serializing in memory to hash
    // is far cheaper than the write + fsync it avoids, and the serialized
    // bytes are reused for the write itself.
    let sidecar_path = project_path
        .join(".leindex")
        .join("search_snapshot.identity");
    let bytes = bincode::serialize(&snapshot).context("Failed to serialize search snapshot")?;
    let identity = format!("v2:{}\n", blake3::hash(&bytes));
    if path.is_file()
        && std::fs::read_to_string(&sidecar_path).ok().as_deref() == Some(identity.as_str())
    {
        tracing::debug!(
            path = %path.display(),
            "Search snapshot unchanged; skipping rewrite"
        );
        return Ok(());
    }
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).with_context(|| {
            format!(
                "Failed to create search snapshot directory: {}",
                parent.display()
            )
        })?;
    }

    std::fs::write(&path, &bytes)
        .with_context(|| format!("Failed to write search snapshot: {}", path.display()))?;
    std::fs::write(&sidecar_path, &identity).with_context(|| {
        format!(
            "Failed to write search snapshot identity sidecar: {}",
            sidecar_path.display()
        )
    })?;
    info!(
        count = snapshot.indexed_nodes,
        path = %path.display(),
        "Persisted search snapshot"
    );
    Ok(())
}

/// Try to load search metadata from an explicit storage directory.
pub(crate) fn try_load_search_snapshot_from_storage(storage_path: &Path) -> Option<SearchSnapshot> {
    let path = storage_path.join("search_snapshot.bin");
    if !path.exists() {
        return None;
    }

    match std::fs::read(&path)
        .with_context(|| format!("Failed to read search snapshot: {}", path.display()))
        .and_then(|bytes| {
            bincode::deserialize::<SearchSnapshot>(&bytes)
                .context("Failed to deserialize search snapshot")
        }) {
        Ok(snapshot) => {
            info!(
                count = snapshot.indexed_nodes,
                path = %path.display(),
                "Loaded search snapshot"
            );
            Some(snapshot)
        }
        Err(e) => {
            tracing::warn!(
                path = %path.display(),
                error = %e,
                "Failed to load search snapshot"
            );
            None
        }
    }
}

struct Blake3FormatWriter<'a>(&'a mut blake3::Hasher);

impl std::fmt::Write for Blake3FormatWriter<'_> {
    fn write_str(&mut self, value: &str) -> std::fmt::Result {
        self.0.update(value.as_bytes());
        Ok(())
    }
}

/// Stable fingerprint of the PDG state that materially affects search
/// hydration. This prevents a snapshot produced for one PDG from being reused
/// after storage changes that happen to preserve node/edge counts.
pub(crate) fn pdg_search_fingerprint(pdg: &ProgramDependenceGraph) -> String {
    use std::fmt::Write as _;

    let mut hasher = blake3::Hasher::new();
    hasher.update(b"leindex-pdg-search-v2");

    // Each record is hashed independently and the digests are sorted, so the
    // result is order-independent by construction — which makes the per-record
    // work (167k BLAKE3 hashes on a mid-size project, ~110 ms serial) safe to
    // spread across cores.
    use rayon::prelude::*;
    let node_ids: Vec<_> = pdg.node_indices().collect();
    let mut nodes: Vec<[u8; 32]> = node_ids
        .par_iter()
        .filter_map(|&node_idx| {
            pdg.get_node(node_idx).map(|node| {
                let mut record = blake3::Hasher::new();
                write!(
                    Blake3FormatWriter(&mut record),
                    "{}\0{:?}\0{}\0{}\0{}\0{}\0{}\0{}",
                    node.id,
                    node.node_type,
                    node.name,
                    node.file_path,
                    node.byte_range.0,
                    node.byte_range.1,
                    node.complexity,
                    node.language
                )
                .expect("hash writer is infallible");
                *record.finalize().as_bytes()
            })
        })
        .collect();
    nodes.par_sort_unstable();
    for node in nodes {
        hasher.update(&node);
    }

    let edge_ids: Vec<_> = pdg.edge_indices().collect();
    let mut edges: Vec<[u8; 32]> = edge_ids
        .par_iter()
        .filter_map(|&edge_idx| {
            let edge = pdg.get_edge(edge_idx)?;
            let (from, to) = pdg.edge_endpoints(edge_idx)?;
            let from = pdg.get_node(from)?;
            let to = pdg.get_node(to)?;
            let mut record = blake3::Hasher::new();
            write!(
                Blake3FormatWriter(&mut record),
                "{}\0{}\0{:?}\0{:?}\0{:?}\0{:?}",
                from.id,
                to.id,
                edge.edge_type,
                edge.metadata.call_count,
                edge.metadata.variable_name,
                edge.metadata.confidence.map(f32::to_bits)
            )
            .expect("hash writer is infallible");
            Some(*record.finalize().as_bytes())
        })
        .collect();
    edges.par_sort_unstable();
    for edge in edges {
        hasher.update(&edge);
    }

    hasher.finalize().to_hex().to_string()
}

/// Persist neural embeddings to a separate mmap file for fast load_from_storage.
///
/// This stores the ONNX neural embeddings (1024-dim) separately from the
/// TF-IDF embeddings (768-dim) so they can be restored without re-computing
/// ONNX inference for all nodes.
#[cfg(any(feature = "onnx", feature = "remote-embeddings"))]
pub(crate) fn persist_neural_embeddings_to_mmap(
    search_engine: &SearchEngine,
    project_path: &Path,
    identity: Option<&str>,
) -> Result<()> {
    let embeddings = search_engine.collect_neural_embeddings();
    let path = neural_mmap_embeddings_path(project_path);
    if embeddings.is_empty() {
        if path.exists() {
            std::fs::remove_file(&path).map_err(|e| {
                anyhow::anyhow!("Failed to remove stale neural mmap embeddings: {e}")
            })?;
            info!(
                path = %path.display(),
                "Removed stale neural embeddings mmap file"
            );
        }
        return Ok(());
    }
    crate::search::vector::write_mmap_embeddings_with_identity(&path, &embeddings, identity)
        .map_err(|e| anyhow::anyhow!("Failed to write neural mmap embeddings: {e}"))?;
    info!(
        count = embeddings.len(),
        path = %path.display(),
        "Persisted neural embeddings to mmap file"
    );
    Ok(())
}

/// Path for the neural embeddings mmap file.
#[cfg(any(feature = "onnx", feature = "remote-embeddings"))]
pub(crate) fn neural_mmap_embeddings_path(project_path: &Path) -> PathBuf {
    project_path.join(".leindex").join("neural_embeddings.bin")
}

/// Try to load previously persisted neural embeddings from mmap file.
///
/// Returns `None` if the file does not exist or is corrupt.
#[cfg(any(feature = "onnx", feature = "remote-embeddings"))]
pub(crate) fn try_load_neural_mmap_embeddings(
    project_path: &Path,
) -> Option<crate::search::vector::MmapEmbeddingIndex> {
    try_load_neural_mmap_embeddings_from_storage(&project_path.join(".leindex"))
}

/// Try to load neural embeddings from an explicit storage directory.
#[cfg(any(feature = "onnx", feature = "remote-embeddings"))]
pub(crate) fn try_load_neural_mmap_embeddings_from_storage(
    storage_path: &Path,
) -> Option<crate::search::vector::MmapEmbeddingIndex> {
    let path = storage_path.join("neural_embeddings.bin");
    if !path.exists() {
        return None;
    }
    match crate::search::vector::MmapEmbeddingIndex::open(&path) {
        Ok(index) => {
            info!(
                nodes = index.len(),
                dim = index.dimension(),
                path = %path.display(),
                "Loaded neural mmap embedding index"
            );
            Some(index)
        }
        Err(e) => {
            tracing::warn!(
                path = %path.display(),
                error = %e,
                "Failed to load neural mmap embedding index"
            );
            None
        }
    }
}

// ---------------------------------------------------------------------------
// Fragment embeddings mmap twins (fragment-embeddings 1.11.0 Task 4)
// ---------------------------------------------------------------------------
//
// Structural twins of the neural mmap path (`neural_embeddings.bin` →
// `MmapEmbeddingIndex`). Fragment embeddings are content-hash-addressed: the
// IDs written are blake3 content hashes of the enriched fragment text, so a
// row in `fragments_embeddings.bin` maps 1:1 to a `FragmentStore` row key.
// Additive-only (invariant 4): never mutates `embeddings.bin`/
// `neural_embeddings.bin`.
//
// Fragment persistence twins (Task 4) — now live: `persist_search_snapshot`
// persists the fragment mmap + root (Task 5) and `indexing/load.rs` loads the
// mmap for hydration.

/// Path for the fragment embeddings mmap file.
pub(crate) fn fragment_mmap_embeddings_path(project_path: &Path) -> PathBuf {
    project_path
        .join(".leindex")
        .join("fragments_embeddings.bin")
}

/// Persist fragment embeddings to a separate mmap file for fast hydration.
///
/// Mirrors `persist_neural_embeddings_to_mmap`: an empty input removes a stale
/// file (feature-off leaves no orphan artifact); otherwise writes the matrix
/// via `write_mmap_embeddings`. The `embeddings` slice is `(content_hash,
/// Vec<f32>)` pairs collected by Task 5's `SearchEngine::collect_fragment_embeddings`
/// (or its fragment-store fallback); passing them in keeps this task free of a
/// `SearchEngine` dependency that does not exist until Task 5 adds
/// `fragment_vector_index`.
pub(crate) fn persist_fragment_embeddings_to_mmap(
    project_path: &Path,
    embeddings: &[(String, Vec<f32>)],
) -> Result<()> {
    let path = fragment_mmap_embeddings_path(project_path);
    if embeddings.is_empty() {
        if path.exists() {
            std::fs::remove_file(&path).map_err(|e| {
                anyhow::anyhow!("Failed to remove stale fragment mmap embeddings: {e}")
            })?;
            info!(
                path = %path.display(),
                "Removed stale fragment embeddings mmap file"
            );
        }
        return Ok(());
    }
    crate::search::vector::write_mmap_embeddings(&path, embeddings)
        .map_err(|e| anyhow::anyhow!("Failed to write fragment mmap embeddings: {e}"))?;
    info!(
        count = embeddings.len(),
        path = %path.display(),
        "Persisted fragment embeddings to mmap file"
    );
    Ok(())
}

/// Try to load previously persisted fragment embeddings from an explicit
/// storage directory (`.leindex/fragments_embeddings.bin`).
///
/// Returns `None` when the file does not exist or is corrupt (mirrors the
/// neural loader's warn-and-continue behavior).
pub(crate) fn try_load_fragment_mmap_embeddings_from_storage(
    storage_path: &Path,
) -> Option<crate::search::vector::MmapEmbeddingIndex> {
    let path = storage_path.join("fragments_embeddings.bin");
    if !path.exists() {
        return None;
    }
    match crate::search::vector::MmapEmbeddingIndex::open(&path) {
        Ok(index) => {
            info!(
                nodes = index.len(),
                dim = index.dimension(),
                path = %path.display(),
                "Loaded fragment mmap embedding index"
            );
            Some(index)
        }
        Err(e) => {
            tracing::warn!(
                path = %path.display(),
                error = %e,
                "Failed to load fragment mmap embedding index"
            );
            None
        }
    }
}

/// Validate the persisted fragment layer against the snapshot (invariant 8).
///
/// The snapshot records the fragment root hash (filled at persist time by
/// `persist_search_snapshot`); `fragment_root.bin` must carry the SAME hash
/// and a row count matching the fragment mmap. Any mismatch means the fragment
/// artifacts are stale relative to the snapshot — callers must disable the
/// fragment layer (never block the node-level path, invariant 3).
pub(crate) fn fragment_layer_is_valid(
    snapshot_root: Option<&str>,
    fragment_mmap: Option<&crate::search::vector::MmapEmbeddingIndex>,
    storage_path: &Path,
) -> bool {
    let (Some(root), Some(mmap)) = (snapshot_root, fragment_mmap) else {
        return false;
    };
    match fragment::sync::load_fragment_root(storage_path) {
        Ok(Some(state)) => {
            // Task 7: refuse a HALF-SYNCED fragment tree. The manifest + store +
            // root are written together under one bumped generation; a
            // mid-build crash leaves an older root (generation mismatch) which
            // must not serve — the caller falls back to the last complete root
            // (i.e. the fragment layer stays off). Legacy Task 5/6 artifacts
            // with no manifest are accepted.
            state.root_hash == root
                && state.fragment_rows == mmap.len() as u64
                && fragment::sync::fragment_layer_generation_is_consistent(storage_path, &state)
        }
        Ok(None) => false,
        Err(e) => {
            tracing::warn!(
                error = %e,
                "Failed to load fragment root for validation; fragment layer disabled"
            );
            false
        }
    }
}