clawhdf5/crates/clawhdf5-migrate/src/validate.rs

//! Validate a migration by reading the store back the way an agent would:
//! through `HDF5Memory::open_read_only`, comparing what it loads with the
//! SQLite source, and running a search for a migrated record.

use std::path::Path;

use clawhdf5_agent::{AgentMemory, HDF5Memory, SearchOptions};
use clawhdf5_format::float16::round_to_f16;

use crate::sqlite_reader::SqliteData;
use crate::store_writer::{Migration, US_PER_SEC};

type BoxErr = Box<dyn std::error::Error>;

/// Summary of a migration validation.
#[derive(Debug)]
pub struct ValidationSummary {
    /// Records in the store (including tombstones).
    pub count: usize,
    /// Records in the store that are not deleted.
    pub active: usize,
    pub sessions: usize,
    pub entities: usize,
    pub relations: usize,
    pub embedding_dim: usize,
    pub float16: bool,
    /// Rows whose full content was compared against the source.
    pub rows_checked: u64,
    /// Whether a search for a migrated record found it (`false` when there
    /// was no active migrated record with an embedding to search for).
    pub search_checked: bool,
}

/// Validate the store at `path` against the source rows `migration` wrote.
///
/// Counts and the session / entity / relation rows are always checked in
/// full. Memory records are content-checked on a representative sample, or
/// all of them with `full`. Embeddings must match exactly: the source values
/// themselves in an `f32` store, their [`round_to_f16`] in a `float16` one.
pub fn validate_store(
    path: &Path,
    source: &SqliteData,
    migration: &Migration,
    full: bool,
) -> Result<ValidationSummary, BoxErr> {
    let mut mem = HDF5Memory::open_read_only(path)?;
    let float16 = mem.config().float16;
    let dim = mem.config().embedding_dim;

    // ---- Counts ----
    check_count("record", mem.count(), migration.store_count)?;
    if float16 != migration.float16 {
        return Err(format!(
            "float16 mismatch: store {float16}, expected {}",
            migration.float16
        )
        .into());
    }
    if dim != migration.embedding_dim {
        return Err(format!(
            "embedding_dim mismatch: store has {dim}, expected {}",
            migration.embedding_dim
        )
        .into());
    }
    if !migration.appended_to_existing {
        check_count("record", mem.count(), migration.records.len())?;
        check_count("session", mem.sessions().len(), migration.sessions.len())?;
        check_count(
            "entity",
            mem.knowledge().entities.len(),
            migration.entities.len(),
        )?;
        check_count(
            "relation",
            mem.knowledge().relations.len(),
            migration.relations.len(),
        )?;
    }

    // ---- Memory records (sampled or full) ----
    let mut rows_checked = 0u64;
    let expected_value = |v: f32| if float16 { round_to_f16(v) } else { v };
    for k in sample_indices(migration.records.len(), full) {
        let (idx, src) = migration.records[k];
        let s = &source.chunks[src];
        let c = &mem.cache;
        if idx >= c.len() {
            return Err(
                format!("record {idx} (chunk id {}) is missing from the store", s.id).into(),
            );
        }
        let id = s.id;
        if c.chunks[idx] != s.chunk {
            return Err(format!(
                "record {idx} (chunk id {id}) text mismatch: source {:?}, store {:?}",
                truncate(&s.chunk),
                truncate(&c.chunks[idx])
            )
            .into());
        }
        if c.source_channels[idx] != s.source_channel
            || c.session_ids[idx] != s.session_id
            || c.tags[idx] != s.tags
        {
            return Err(format!("record {idx} (chunk id {id}) string field mismatch").into());
        }
        if c.timestamps[idx].to_bits() != s.timestamp.to_bits() {
            return Err(format!(
                "record {idx} (chunk id {id}) timestamp mismatch: source {}, store {}",
                s.timestamp, c.timestamps[idx]
            )
            .into());
        }
        let deleted = c.tombstones[idx] != 0;
        if deleted != (s.deleted != 0) {
            return Err(format!(
                "record {idx} (chunk id {id}) deleted mismatch: source {}, store {deleted}",
                s.deleted != 0
            )
            .into());
        }
        let got = c.embeddings.get(idx).unwrap_or(&[]);
        if got.len() != s.embedding.len() {
            return Err(format!(
                "record {idx} (chunk id {id}) embedding length mismatch: source {}, store {}",
                s.embedding.len(),
                got.len()
            )
            .into());
        }
        for (j, (&a, &b)) in s.embedding.iter().zip(got).enumerate() {
            let want = expected_value(a);
            if want.to_bits() != b.to_bits() && !(want.is_nan() && b.is_nan()) {
                return Err(format!(
                    "record {idx} (chunk id {id}) embedding[{j}] mismatch: source {a}, \
                     expected {want}, store {b}"
                )
                .into());
            }
        }
        rows_checked += 1;
    }

    // ---- Records tombstoned because their source row was deleted ----
    for &(idx, src) in &migration.deleted_in_store {
        let s = &source.chunks[src];
        let c = &mem.cache;
        if idx >= c.len() || c.chunks[idx] != s.chunk || c.timestamps[idx] != s.timestamp {
            return Err(format!("record {idx} (chunk id {}) mismatch or missing", s.id).into());
        }
        if c.tombstones[idx] == 0 {
            return Err(format!(
                "record {idx} (chunk id {}) is deleted in the source but active in the store",
                s.id
            )
            .into());
        }
        rows_checked += 1;
    }

    // ---- Sessions ----
    let sessions = mem.sessions();
    for &(at, src) in &migration.sessions {
        let s = &source.sessions[src];
        let (Some(e), Some(summary)) = (sessions.entries.get(at), sessions.summaries.get(at))
        else {
            return Err(format!("session {:?} is missing from the store", s.id).into());
        };
        if e.id != s.id
            || e.start_idx != s.start_idx.max(0) as u64
            || e.end_idx != s.end_idx.max(0) as u64
            || e.channel != s.channel
            || *summary != s.summary
            || e.ts != s.timestamp * US_PER_SEC
        {
            return Err(format!("session {:?} mismatch", s.id).into());
        }
        rows_checked += 1;
    }

    // ---- Knowledge graph ----
    let kg = mem.knowledge();
    for &(id, src) in &migration.entities {
        let s = &source.entities[src];
        let Some(e) = kg.get_entity(id) else {
            return Err(format!(
                "entity {:?} (id {}) is missing from the store",
                s.name, s.id
            )
            .into());
        };
        if e.name != s.name || e.entity_type != s.entity_type || e.embedding_idx != s.embedding_idx
        {
            return Err(format!("entity {:?} (id {}) mismatch", s.name, s.id).into());
        }
        rows_checked += 1;
    }
    for &(at, src) in &migration.relations {
        let s = &source.relations[src];
        let r = kg.relations.get(at);
        let ok = r.is_some_and(|r| {
            Some(&r.src) == migration.entity_ids.get(&s.src)
                && Some(&r.tgt) == migration.entity_ids.get(&s.tgt)
                && r.relation == s.relation
                && r.weight == s.weight as f32
                && r.ts == s.timestamp * US_PER_SEC
        });
        if !ok {
            return Err(format!(
                "relation {} -[{}]-> {} mismatch or missing",
                s.src, s.relation, s.tgt
            )
            .into());
        }
        rows_checked += 1;
    }

    // ---- A migrated record must be findable by search ----
    let probe = migration
        .records
        .iter()
        .copied()
        .find(|&(idx, _)| dim > 0 && mem.cache.tombstones[idx] == 0);
    let search_checked = match probe {
        None => false,
        Some((idx, _)) => {
            let query = mem.cache.embeddings[idx].to_vec();
            let text = mem.cache.chunks[idx].clone();
            let hits = mem.search(&query, &text, &SearchOptions::new(10));
            // A record with the same text is as good a hit: the source may
            // hold duplicates, and they tie.
            if !hits.iter().any(|h| h.index == idx || h.chunk == text) {
                return Err(format!(
                    "search for migrated record {idx} ({:?}) did not return it",
                    truncate(&text)
                )
                .into());
            }
            true
        }
    };

    Ok(ValidationSummary {
        count: mem.count(),
        active: mem.count_active(),
        sessions: mem.sessions().len(),
        entities: mem.knowledge().entities.len(),
        relations: mem.knowledge().relations.len(),
        embedding_dim: dim,
        float16,
        rows_checked,
        search_checked,
    })
}

fn check_count(kind: &str, got: usize, expected: usize) -> Result<(), BoxErr> {
    if got != expected {
        return Err(format!("{kind} count mismatch: store has {got}, expected {expected}").into());
    }
    Ok(())
}

fn truncate(s: &str) -> String {
    if s.len() <= 40 {
        s.to_string()
    } else {
        let cut = s.char_indices().nth(40).map(|(i, _)| i).unwrap_or(s.len());
        format!("{}…", &s[..cut])
    }
}

/// Indices of records to content-check. Full = all; otherwise a spread of
/// representative rows (first/last and evenly-spaced interior samples).
fn sample_indices(n: usize, full: bool) -> Vec<usize> {
    if n == 0 {
        return Vec::new();
    }
    if full || n <= 16 {
        return (0..n).collect();
    }
    let mut idx: Vec<usize> = (0..16).map(|k| k * (n - 1) / 15).collect();
    idx.dedup();
    idx
}

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

    #[test]
    fn truncate_short_string_unchanged() {
        assert_eq!(truncate("hello"), "hello");
    }

    /// A multi-byte character straddling byte offset 40 must not panic a
    /// byte-index slice — this is arbitrary UTF-8 chunk text from an
    /// untrusted source database, not test-only input.
    #[test]
    fn truncate_multibyte_char_at_boundary_does_not_panic() {
        // 39 ASCII bytes then a 4-byte emoji straddling the byte-40 cut point.
        let s = format!("{}{}", "a".repeat(39), "😀".repeat(5));
        let result = truncate(&s);
        assert!(result.ends_with('…'));
        assert!(result.chars().count() < s.chars().count());
    }

    #[test]
    fn truncate_exactly_at_limit_unchanged() {
        let s = "a".repeat(40);
        assert_eq!(truncate(&s), s);
    }
}