clawhdf5/crates/clawhdf5-agent/src/schema.rs

//! HDF5 schema creation and validation.
//!
//! Handles building the HDF5 file structure from in-memory caches and
//! reading/validating existing files.

use clawhdf5::AttrValue;
use clawhdf5::FillTime;
use clawhdf5_format::datatype::{CharacterSet, Datatype, StringPadding};

use crate::MemoryConfig;
use crate::MemoryError;
use crate::cache::MemoryCache;
use crate::knowledge::KnowledgeCache;
use crate::session::SessionCache;
use crate::wal::WalMark;

pub const SCHEMA_VERSION: &str = "1.0";
pub const ZEROCLAW_VERSION: &str = "0.8.0";

/// `/meta` attributes holding the [`WalMark`] of the WAL prefix already folded
/// into this file. Absent on files written before the mark existed, and when
/// the checkpoint was taken with an empty WAL.
const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len";
const WAL_APPLIED_CRC_ATTR: &str = "wal_applied_crc";
const ANN_GENERATION_ATTR: &str = "ann_generation";

/// Build a complete HDF5 file from the in-memory state.
pub fn build_hdf5_file(
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
) -> Result<Vec<u8>, MemoryError> {
    build_hdf5_file_with_mark(config, cache, sessions, knowledge, None)
}

/// [`build_hdf5_file`], recording which WAL prefix this state already
/// contains (see [`WalMark`]) so a crash before the WAL is truncated doesn't
/// replay those entries a second time.
pub fn build_hdf5_file_with_mark(
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
    wal_applied: Option<WalMark>,
) -> Result<Vec<u8>, MemoryError> {
    let meta = CheckpointMeta {
        wal_applied,
        ann_generation: None,
    };
    build_hdf5_file_with_meta(config, cache, sessions, knowledge, &meta)
}

/// Bookkeeping a checkpoint records in `/meta` beside the store's contents.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CheckpointMeta {
    /// The WAL prefix this checkpoint already contains; see [`WalMark`].
    pub wal_applied: Option<WalMark>,
    /// Identifies the vector-index sidecar (`<store>.h5.ann`) written with this
    /// checkpoint. A sidecar is loaded only if it carries the same value, so
    /// one left over from another checkpoint can never be attached to records
    /// it wasn't built from.
    pub ann_generation: Option<u64>,
}

/// [`build_hdf5_file`] with checkpoint bookkeeping.
pub fn build_hdf5_file_with_meta(
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
    checkpoint: &CheckpointMeta,
) -> Result<Vec<u8>, MemoryError> {
    let wal_applied = checkpoint.wal_applied;
    let mut builder = clawhdf5::FileBuilder::new();

    // /meta group with schema attributes
    let mut meta = builder.create_group("meta");
    meta.set_attr("schema_version", AttrValue::String(SCHEMA_VERSION.into()));
    meta.set_attr("created_at", AttrValue::String(config.created_at.clone()));
    meta.set_attr("agent_id", AttrValue::String(config.agent_id.clone()));
    meta.set_attr("embedder", AttrValue::String(config.embedder.clone()));
    meta.set_attr("embedding_dim", AttrValue::I64(config.embedding_dim as i64));
    meta.set_attr("chunk_size", AttrValue::I64(config.chunk_size as i64));
    meta.set_attr("overlap", AttrValue::I64(config.overlap as i64));
    // Behavioural settings. These used to live only in memory, so reopening a
    // store silently reset them to defaults — e.g. a compressed store was
    // rewritten uncompressed by the first checkpoint after a reopen. Loaders
    // treat each one as optional so older files keep opening.
    meta.set_attr("float16", AttrValue::I64(config.float16.into()));
    meta.set_attr("compression", AttrValue::I64(config.compression.into()));
    meta.set_attr(
        "compression_level",
        AttrValue::I64(config.compression_level.into()),
    );
    meta.set_attr(
        "compact_threshold",
        AttrValue::F64(config.compact_threshold.into()),
    );
    meta.set_attr("hebbian_boost", AttrValue::F64(config.hebbian_boost.into()));
    meta.set_attr("decay_factor", AttrValue::F64(config.decay_factor.into()));
    meta.set_attr("wal_enabled", AttrValue::I64(config.wal_enabled.into()));
    meta.set_attr(
        "wal_max_entries",
        AttrValue::I64(config.wal_max_entries as i64),
    );
    meta.set_attr(
        "edgehdf5_version",
        AttrValue::String(ZEROCLAW_VERSION.into()),
    );
    if let Some(mark) = wal_applied.filter(|m| m.len > 0) {
        meta.set_attr(WAL_APPLIED_LEN_ATTR, AttrValue::I64(mark.len as i64));
        meta.set_attr(WAL_APPLIED_CRC_ATTR, AttrValue::I64(i64::from(mark.crc)));
    }
    if let Some(generation) = checkpoint.ann_generation {
        // Stored as the i64 with the same bits; attributes have no u64 scalar
        // round trip through every reader.
        meta.set_attr(ANN_GENERATION_ATTR, AttrValue::I64(generation as i64));
    }
    // Need at least one dataset in the group for it to be a proper group
    meta.create_dataset("_marker").with_u8_data(&[1]).compact();
    let finished_meta = meta.finish();
    builder.add_group(finished_meta);

    // /memory group
    build_memory_group(&mut builder, config, cache)?;

    // /sessions group
    build_sessions_group(&mut builder, sessions)?;

    // /knowledge_graph group
    build_knowledge_group(&mut builder, knowledge)?;

    builder
        .finish()
        .map_err(|e| MemoryError::Hdf5(e.to_string()))
}

fn build_memory_group(
    builder: &mut clawhdf5::FileBuilder,
    config: &MemoryConfig,
    cache: &MemoryCache,
) -> Result<(), MemoryError> {
    let mut group = builder.create_group("memory");

    // chunks: fixed-length string array
    write_string_dataset(&mut group, "chunks", &cache.chunks);

    // embeddings: f32 [N x D]
    let n = cache.embeddings.len() as u64;
    let d = cache.embedding_dim as u64;
    let flat = cache.flat_embeddings();
    {
        let ds = group
            .create_dataset("embeddings")
            .with_f32_data(flat)
            .with_shape(&[n, d]);

        // Chunk size tuning: target ~256KB per chunk for optimal I/O
        if n > 0 && d > 0 {
            let target_chunk_bytes: u64 = 256 * 1024;
            let rows_per_chunk = (target_chunk_bytes / (d * 4)).max(1).min(n);
            ds.with_chunks(&[rows_per_chunk, d]);

            // Compression. Shuffle is applied automatically (auto-shuffle
            // pre-filter). Zstd is faster than deflate at the same ratio but
            // pulls in libzstd, so it is opt-in via the `zstd` feature; the
            // default build uses deflate, which is always available. (This
            // used to call `with_zstd` unconditionally, so without the
            // feature every checkpoint of a compressed store failed with
            // "unsupported filter: 32015".) Both are standard HDF5 filters;
            // reading a zstd-compressed store needs a zstd-enabled build.
            if config.compression {
                #[cfg(feature = "zstd")]
                {
                    let level = if config.compression_level > 0 {
                        config.compression_level.min(22)
                    } else {
                        3 // fast + good ratio for f32 embeddings
                    };
                    ds.with_zstd(level);
                }
                #[cfg(not(feature = "zstd"))]
                {
                    let level = if config.compression_level > 0 {
                        config.compression_level.min(9)
                    } else {
                        4
                    };
                    ds.with_deflate(level);
                }
            }
        }

        // Skip fill-value initialization — embeddings are fully written
        ds.fill_time(FillTime::Never);
        // Page-aligned for sequential scans
        ds.align(4096);
    }

    // source_channel: fixed-length string array
    write_string_dataset(&mut group, "source_channel", &cache.source_channels);

    // timestamps: f64 array
    group
        .create_dataset("timestamps")
        .with_f64_data(&cache.timestamps)
        .fill_time(FillTime::Never);

    // session_ids: fixed-length string array (auto-compressed when large)
    write_string_dataset(&mut group, "session_ids", &cache.session_ids);

    // tags: fixed-length string array (auto-compressed when large)
    write_string_dataset(&mut group, "tags", &cache.tags);

    // tombstones: u8 array — use compact if small
    {
        let ds = group
            .create_dataset("tombstones")
            .with_u8_data(&cache.tombstones);
        if cache.tombstones.len() <= 65536 {
            ds.compact();
        }
        ds.fill_time(FillTime::Never);
    }

    // norms: f32 array (pre-computed L2 norms)
    group
        .create_dataset("norms")
        .with_f32_data(&cache.norms)
        .fill_time(FillTime::Never);

    // activation_weights: f32 array (Hebbian activation weights)
    group
        .create_dataset("activation_weights")
        .with_f32_data(&cache.activation_weights)
        .fill_time(FillTime::Never);

    let finished = group.finish();
    builder.add_group(finished);
    Ok(())
}

fn build_sessions_group(
    builder: &mut clawhdf5::FileBuilder,
    sessions: &SessionCache,
) -> Result<(), MemoryError> {
    let mut group = builder.create_group("sessions");

    let ids: Vec<String> = sessions.entries.iter().map(|e| e.id.clone()).collect();
    write_string_dataset(&mut group, "ids", &ids);

    let start_idxs: Vec<i64> = sessions
        .entries
        .iter()
        .map(|e| e.start_idx as i64)
        .collect();
    group
        .create_dataset("start_idxs")
        .with_i64_data(&start_idxs);

    let end_idxs: Vec<i64> = sessions.entries.iter().map(|e| e.end_idx as i64).collect();
    group.create_dataset("end_idxs").with_i64_data(&end_idxs);

    let channels: Vec<String> = sessions.entries.iter().map(|e| e.channel.clone()).collect();
    write_string_dataset(&mut group, "channels", &channels);

    let timestamps: Vec<f64> = sessions.entries.iter().map(|e| e.ts).collect();
    group
        .create_dataset("timestamps")
        .with_f64_data(&timestamps);

    write_string_dataset(&mut group, "summaries", &sessions.summaries);

    let finished = group.finish();
    builder.add_group(finished);
    Ok(())
}

fn build_knowledge_group(
    builder: &mut clawhdf5::FileBuilder,
    knowledge: &KnowledgeCache,
) -> Result<(), MemoryError> {
    let mut group = builder.create_group("knowledge_graph");

    // Entities
    let entity_ids: Vec<i64> = knowledge.entities.iter().map(|e| e.id as i64).collect();
    group
        .create_dataset("entity_ids")
        .with_i64_data(&entity_ids);

    let entity_names: Vec<String> = knowledge.entities.iter().map(|e| e.name.clone()).collect();
    write_string_dataset(&mut group, "entity_names", &entity_names);

    let entity_types: Vec<String> = knowledge
        .entities
        .iter()
        .map(|e| e.entity_type.clone())
        .collect();
    write_string_dataset(&mut group, "entity_types", &entity_types);

    let emb_idxs: Vec<i64> = knowledge.entities.iter().map(|e| e.embedding_idx).collect();
    group
        .create_dataset("entity_emb_idxs")
        .with_i64_data(&emb_idxs);

    // Relations
    let rel_srcs: Vec<i64> = knowledge.relations.iter().map(|r| r.src as i64).collect();
    group
        .create_dataset("relation_srcs")
        .with_i64_data(&rel_srcs);

    let rel_tgts: Vec<i64> = knowledge.relations.iter().map(|r| r.tgt as i64).collect();
    group
        .create_dataset("relation_tgts")
        .with_i64_data(&rel_tgts);

    let rel_types: Vec<String> = knowledge
        .relations
        .iter()
        .map(|r| r.relation.clone())
        .collect();
    write_string_dataset(&mut group, "relation_types", &rel_types);

    let rel_weights: Vec<f32> = knowledge.relations.iter().map(|r| r.weight).collect();
    group
        .create_dataset("relation_weights")
        .with_f32_data(&rel_weights);

    let rel_ts: Vec<f64> = knowledge.relations.iter().map(|r| r.ts).collect();
    group.create_dataset("relation_ts").with_f64_data(&rel_ts);

    // Aliases
    if !knowledge.alias_strings.is_empty() {
        write_string_dataset(&mut group, "alias_strings", &knowledge.alias_strings);
        group
            .create_dataset("alias_entity_ids")
            .with_i64_data(&knowledge.alias_entity_ids);
    }

    let finished = group.finish();
    builder.add_group(finished);
    Ok(())
}

/// Write a string array as a fixed-length string dataset.
///
/// Uses `Datatype::String` with NullPad encoding. Each string is padded
/// to the length of the longest string in the array.
///
/// When `compress` is true, uses chunked storage with deflate(6) —
/// NullPad strings have high redundancy and compress very well.
/// Payload size (bytes) at or above which a fixed-length string dataset is
/// stored chunked + deflate-compressed. Below this, the chunk B-tree/heap
/// overhead outweighs the savings, so the data is left contiguous.
const STRING_COMPRESS_THRESHOLD: usize = 4096;

fn write_string_dataset(
    group: &mut clawhdf5_format::type_builders::GroupBuilder,
    name: &str,
    strings: &[String],
) {
    if strings.is_empty() {
        // Empty dataset: use 1-byte string type with no data
        let dtype = Datatype::String {
            size: 1,
            padding: StringPadding::NullPad,
            charset: CharacterSet::Utf8,
        };
        group
            .create_dataset(name)
            .with_compound_data(dtype, vec![], 0);
        return;
    }

    let max_len = strings.iter().map(|s| s.len()).max().unwrap_or(0).max(1);
    let mut raw = Vec::with_capacity(strings.len() * max_len);
    for s in strings {
        let mut bytes = s.as_bytes().to_vec();
        bytes.resize(max_len, 0);
        raw.extend_from_slice(&bytes);
    }
    let raw_len = raw.len();

    let dtype = Datatype::String {
        size: max_len as u32,
        padding: StringPadding::NullPad,
        charset: CharacterSet::Utf8,
    };
    let ds = group
        .create_dataset(name)
        .with_compound_data(dtype, raw, strings.len() as u64);

    // Fixed-length NullPad strings have high redundancy (padding + repeated
    // content), so deflate pays off once the payload is large enough to absorb
    // the chunking overhead. Fixed-length string datasets are chunkable like
    // any other fixed-size datatype.
    if strings.len() > 1 && raw_len >= STRING_COMPRESS_THRESHOLD {
        // Target ~64KB chunks for string data.
        let elem_size = max_len as u64;
        let target_chunk = 64 * 1024;
        let rows_per_chunk = (target_chunk / elem_size).max(1).min(strings.len() as u64);
        ds.with_chunks(&[rows_per_chunk]);
        ds.with_deflate(6);
    }
}

/// Validate an HDF5 file has the correct schema and load all data.
/// Read the checkpoint's [`WalMark`] from `/meta`, if it has one.
pub fn read_wal_mark(file: &clawhdf5::File) -> Option<WalMark> {
    let attrs = file.group("meta").ok()?.attrs().ok()?;
    let len = match attrs.get(WAL_APPLIED_LEN_ATTR)? {
        AttrValue::I64(v) => u64::try_from(*v).ok()?,
        _ => return None,
    };
    let crc = match attrs.get(WAL_APPLIED_CRC_ATTR)? {
        AttrValue::I64(v) => u32::try_from(*v).ok()?,
        _ => return None,
    };
    Some(WalMark { len, crc })
}

/// Read the checkpoint bookkeeping from `/meta`.
pub fn read_checkpoint_meta(file: &clawhdf5::File) -> CheckpointMeta {
    let ann_generation = file
        .group("meta")
        .ok()
        .and_then(|g| g.attrs().ok())
        .and_then(|attrs| match attrs.get(ANN_GENERATION_ATTR) {
            Some(AttrValue::I64(v)) => Some(*v as u64),
            _ => None,
        });
    CheckpointMeta {
        wal_applied: read_wal_mark(file),
        ann_generation,
    }
}

pub fn validate_and_load(
    file: &clawhdf5::File,
) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> {
    // Read /meta group attributes
    let meta = file
        .group("meta")
        .map_err(|e| MemoryError::Schema(format!("missing /meta group: {e}")))?;
    let attrs = meta
        .attrs()
        .map_err(|e| MemoryError::Schema(format!("cannot read /meta attrs: {e}")))?;

    let schema_version = match attrs.get("schema_version") {
        Some(AttrValue::String(s)) => s.clone(),
        _ => return Err(MemoryError::Schema("missing schema_version attr".into())),
    };
    if schema_version != SCHEMA_VERSION {
        return Err(MemoryError::Schema(format!(
            "schema version mismatch: expected {SCHEMA_VERSION}, got {schema_version}"
        )));
    }

    let created_at = extract_string_attr(&attrs, "created_at")?;
    let agent_id = extract_string_attr(&attrs, "agent_id")?;
    let embedder = extract_string_attr(&attrs, "embedder")?;
    let embedding_dim = extract_i64_attr(&attrs, "embedding_dim")? as usize;
    let chunk_size = extract_i64_attr(&attrs, "chunk_size")? as usize;
    let overlap = extract_i64_attr(&attrs, "overlap")? as usize;

    let config = MemoryConfig {
        path: std::path::PathBuf::new(), // will be set by caller
        agent_id,
        embedder,
        embedding_dim,
        chunk_size,
        overlap,
        float16: optional_bool_attr(&attrs, "float16", false),
        compression: optional_bool_attr(&attrs, "compression", false),
        compression_level: optional_i64_attr(&attrs, "compression_level")
            .and_then(|v| u32::try_from(v).ok())
            .unwrap_or(0),
        compact_threshold: optional_f32_attr(&attrs, "compact_threshold", 0.3),
        hebbian_boost: optional_f32_attr(&attrs, "hebbian_boost", 0.15),
        decay_factor: optional_f32_attr(&attrs, "decay_factor", 0.98),
        created_at,
        wal_enabled: optional_bool_attr(&attrs, "wal_enabled", true),
        wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries")
            .and_then(|v| usize::try_from(v).ok())
            .unwrap_or(500),
    };

    // Load /memory group
    let memory_cache = load_memory_group(file, embedding_dim)?;

    // Load /sessions group
    let session_cache = load_sessions_group(file)?;

    // Load /knowledge_graph group
    let knowledge_cache = load_knowledge_group(file)?;

    Ok((config, memory_cache, session_cache, knowledge_cache))
}

fn load_memory_group(
    file: &clawhdf5::File,
    embedding_dim: usize,
) -> Result<MemoryCache, MemoryError> {
    let group = file
        .group("memory")
        .map_err(|e| MemoryError::Schema(format!("missing /memory group: {e}")))?;

    let chunks = read_string_dataset_from_group(&group, "chunks")?;
    let n = chunks.len();

    let mut cache = MemoryCache::new(embedding_dim);

    if n == 0 {
        return Ok(cache);
    }

    let flat_embeddings = read_f32_dataset(&group, "embeddings")?;
    let source_channels = read_string_dataset_from_group(&group, "source_channel")?;
    let timestamps = read_f64_dataset(&group, "timestamps")?;
    let session_ids = read_string_dataset_from_group(&group, "session_ids")?;
    let tags = read_string_dataset_from_group(&group, "tags")?;
    let tombstones = read_u8_dataset(&group, "tombstones")?;

    // Every per-record dataset must describe exactly `n` records. Without
    // this, a truncated or hand-edited file loads "successfully" and then
    // panics on the first out-of-bounds index during search/delete.
    if embedding_dim == 0 {
        return Err(MemoryError::Schema(format!(
            "/memory has {n} records but embedding_dim is 0"
        )));
    }
    let expected_flat = n.checked_mul(embedding_dim).ok_or_else(|| {
        MemoryError::Schema(format!("/memory size overflow: {n} x {embedding_dim}"))
    })?;
    let check_len = |name: &str, actual: usize, expected: usize| {
        if actual == expected {
            Ok(())
        } else {
            Err(MemoryError::Schema(format!(
                "/memory/{name} has {actual} entries, expected {expected} \
                 ({n} records)"
            )))
        }
    };
    check_len("embeddings", flat_embeddings.len(), expected_flat)?;
    check_len("source_channel", source_channels.len(), n)?;
    check_len("timestamps", timestamps.len(), n)?;
    check_len("session_ids", session_ids.len(), n)?;
    check_len("tags", tags.len(), n)?;
    check_len("tombstones", tombstones.len(), n)?;

    // Norms are derived data: use the stored ones only if they are present
    // and the right length, otherwise recompute from the embeddings.
    let norms = match read_f32_dataset(&group, "norms") {
        Ok(stored) if stored.len() == n => stored,
        _ => flat_embeddings
            .chunks(embedding_dim)
            .map(|chunk| {
                let sq_sum: f32 = chunk.iter().map(|x| x * x).sum();
                sq_sum.sqrt()
            })
            .collect(),
    };

    // No unflattening: the cache stores the buffer as it is on disk.
    // Read activation_weights if present, default to vec![1.0; N] for backward compat
    let activation_weights = match read_f32_dataset(&group, "activation_weights") {
        Ok(w) if w.len() == n => w,
        _ => vec![1.0; n],
    };

    cache.chunks = chunks;
    cache.embeddings.set_flat(embedding_dim, flat_embeddings);
    cache.source_channels = source_channels;
    cache.timestamps = timestamps;
    cache.session_ids = session_ids;
    cache.tags = tags;
    cache.tombstones = tombstones;
    cache.norms = norms;
    cache.activation_weights = activation_weights;

    Ok(cache)
}

fn load_sessions_group(file: &clawhdf5::File) -> Result<SessionCache, MemoryError> {
    let group = file
        .group("sessions")
        .map_err(|e| MemoryError::Schema(format!("missing /sessions group: {e}")))?;

    let ids = read_string_dataset_from_group(&group, "ids")?;
    if ids.is_empty() {
        return Ok(SessionCache::new());
    }

    let start_idxs = read_i64_dataset(&group, "start_idxs")?;
    let end_idxs = read_i64_dataset(&group, "end_idxs")?;
    let channels = read_string_dataset_from_group(&group, "channels")?;
    let timestamps = read_f64_dataset(&group, "timestamps")?;
    let summaries = read_string_dataset_from_group(&group, "summaries")?;

    let mut cache = SessionCache::new();
    for i in 0..ids.len() {
        cache.entries.push(crate::session::SessionEntry {
            id: ids[i].clone(),
            start_idx: start_idxs[i] as u64,
            end_idx: end_idxs[i] as u64,
            channel: channels[i].clone(),
            ts: timestamps[i],
        });
        cache.summaries.push(summaries[i].clone());
    }

    Ok(cache)
}

fn load_knowledge_group(file: &clawhdf5::File) -> Result<KnowledgeCache, MemoryError> {
    let group = file
        .group("knowledge_graph")
        .map_err(|e| MemoryError::Schema(format!("missing /knowledge_graph group: {e}")))?;

    let entity_ids = read_i64_dataset(&group, "entity_ids")?;
    let next_id = entity_ids.iter().max().map(|&m| m as u64 + 1).unwrap_or(0);
    let mut cache = KnowledgeCache::new_with_next_id(next_id);

    if !entity_ids.is_empty() {
        let entity_names = read_string_dataset_from_group(&group, "entity_names")?;
        let entity_types = read_string_dataset_from_group(&group, "entity_types")?;
        let emb_idxs = read_i64_dataset(&group, "entity_emb_idxs")?;

        for i in 0..entity_ids.len() {
            cache.entities.push(crate::knowledge::Entity {
                id: entity_ids[i] as u64,
                name: entity_names[i].clone(),
                name_lower: entity_names[i].to_lowercase(),
                entity_type: entity_types[i].clone(),
                embedding_idx: emb_idxs[i],
                ..Default::default()
            });
        }
    }

    let rel_srcs = read_i64_dataset(&group, "relation_srcs")?;
    if !rel_srcs.is_empty() {
        let rel_tgts = read_i64_dataset(&group, "relation_tgts")?;
        let rel_types = read_string_dataset_from_group(&group, "relation_types")?;
        let rel_weights = read_f32_dataset(&group, "relation_weights")?;
        let rel_ts = read_f64_dataset(&group, "relation_ts")?;

        for i in 0..rel_srcs.len() {
            cache.relations.push(crate::knowledge::Relation {
                src: rel_srcs[i] as u64,
                tgt: rel_tgts[i] as u64,
                relation: rel_types[i].clone(),
                weight: rel_weights[i],
                ts: rel_ts[i],
                ..Default::default()
            });
        }
    }

    // Load aliases (default to empty for backward compat)
    if let Ok(alias_strings) = read_string_dataset_from_group(&group, "alias_strings")
        && let Ok(alias_entity_ids) = read_i64_dataset(&group, "alias_entity_ids")
    {
        cache.alias_strings = alias_strings;
        cache.alias_entity_ids = alias_entity_ids;
    }

    Ok(cache)
}

// --- Helper functions ---

fn extract_string_attr(
    attrs: &std::collections::HashMap<String, AttrValue>,
    name: &str,
) -> Result<String, MemoryError> {
    match attrs.get(name) {
        Some(AttrValue::String(s)) => Ok(s.clone()),
        _ => Err(MemoryError::Schema(format!("missing attr: {name}"))),
    }
}

type MetaAttrs = std::collections::HashMap<String, AttrValue>;

fn optional_i64_attr(attrs: &MetaAttrs, name: &str) -> Option<i64> {
    match attrs.get(name) {
        Some(AttrValue::I64(v)) => Some(*v),
        _ => None,
    }
}

fn optional_bool_attr(attrs: &MetaAttrs, name: &str, default: bool) -> bool {
    optional_i64_attr(attrs, name).map_or(default, |v| v != 0)
}

/// Finite values only: a NaN threshold/decay would poison every comparison.
fn optional_f32_attr(attrs: &MetaAttrs, name: &str, default: f32) -> f32 {
    match attrs.get(name) {
        Some(AttrValue::F64(v)) if v.is_finite() => *v as f32,
        _ => default,
    }
}

fn extract_i64_attr(
    attrs: &std::collections::HashMap<String, AttrValue>,
    name: &str,
) -> Result<i64, MemoryError> {
    match attrs.get(name) {
        Some(AttrValue::I64(v)) => Ok(*v),
        _ => Err(MemoryError::Schema(format!("missing attr: {name}"))),
    }
}

fn read_string_dataset_from_group(
    group: &clawhdf5::Group<'_>,
    name: &str,
) -> Result<Vec<String>, MemoryError> {
    let ds = group
        .dataset(name)
        .map_err(|e| MemoryError::Hdf5(format!("cannot read {name}: {e}")))?;
    let shape = ds
        .shape()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read shape of {name}: {e}")))?;
    if shape.first() == Some(&0) || shape.is_empty() {
        return Ok(Vec::new());
    }
    ds.read_string()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read strings from {name}: {e}")))
}

fn read_f32_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<f32>, MemoryError> {
    let ds = group
        .dataset(name)
        .map_err(|e| MemoryError::Hdf5(format!("cannot read {name}: {e}")))?;
    let shape = ds
        .shape()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read shape of {name}: {e}")))?;
    if shape.first() == Some(&0) {
        return Ok(Vec::new());
    }
    ds.read_f32()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read f32 from {name}: {e}")))
}

fn read_f64_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<f64>, MemoryError> {
    let ds = group
        .dataset(name)
        .map_err(|e| MemoryError::Hdf5(format!("cannot read {name}: {e}")))?;
    let shape = ds
        .shape()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read shape of {name}: {e}")))?;
    if shape.first() == Some(&0) {
        return Ok(Vec::new());
    }
    ds.read_f64()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read f64 from {name}: {e}")))
}

fn read_i64_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<i64>, MemoryError> {
    let ds = group
        .dataset(name)
        .map_err(|e| MemoryError::Hdf5(format!("cannot read {name}: {e}")))?;
    let shape = ds
        .shape()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read shape of {name}: {e}")))?;
    if shape.first() == Some(&0) {
        return Ok(Vec::new());
    }
    ds.read_i64()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read i64 from {name}: {e}")))
}

fn read_u8_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<u8>, MemoryError> {
    let ds = group
        .dataset(name)
        .map_err(|e| MemoryError::Hdf5(format!("cannot read {name}: {e}")))?;
    let shape = ds
        .shape()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read shape of {name}: {e}")))?;
    if shape.first() == Some(&0) {
        return Ok(Vec::new());
    }
    // Read raw bytes - for u8 data we need the raw representation
    let data = ds
        .read_i32()
        .map_err(|e| MemoryError::Hdf5(format!("cannot read u8 from {name}: {e}")))?;
    Ok(data.into_iter().map(|v| v as u8).collect())
}

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

    fn config() -> MemoryConfig {
        MemoryConfig::new(std::path::PathBuf::from("unused.h5"), "agent", 4)
    }

    fn cache_with(n: usize) -> MemoryCache {
        let mut cache = MemoryCache::new(4);
        for i in 0..n {
            cache.push(
                format!("chunk {i}"),
                vec![i as f32 + 1.0, 0.0, 0.0, 0.0],
                "user".into(),
                i as f64,
                "s".into(),
                "t".into(),
            );
        }
        cache
    }

    fn roundtrip(cache: &MemoryCache) -> Result<MemoryCache, MemoryError> {
        let bytes = build_hdf5_file(
            &config(),
            cache,
            &SessionCache::new(),
            &KnowledgeCache::new(),
        )?;
        let file =
            clawhdf5::File::from_bytes(bytes).map_err(|e| MemoryError::Hdf5(e.to_string()))?;
        validate_and_load(&file).map(|(_, cache, _, _)| cache)
    }

    #[test]
    fn behavioural_config_survives_a_reopen() {
        let mut cfg = config();
        cfg.compression = true;
        cfg.compression_level = 7;
        cfg.compact_threshold = 0.5;
        cfg.hebbian_boost = 0.25;
        cfg.decay_factor = 0.9;
        cfg.wal_enabled = false;
        cfg.wal_max_entries = 42;
        let bytes = build_hdf5_file(
            &cfg,
            &cache_with(2),
            &SessionCache::new(),
            &KnowledgeCache::new(),
        )
        .unwrap();
        let file = clawhdf5::File::from_bytes(bytes).unwrap();
        let (loaded, loaded_cache, ..) = validate_and_load(&file).unwrap();
        // The compressed embeddings must also read back intact.
        assert_eq!(loaded_cache.embeddings, cache_with(2).embeddings);
        assert!(loaded.compression);
        assert_eq!(loaded.compression_level, 7);
        assert_eq!(loaded.compact_threshold, 0.5);
        assert_eq!(loaded.hebbian_boost, 0.25);
        assert_eq!(loaded.decay_factor, 0.9);
        assert!(!loaded.wal_enabled);
        assert_eq!(loaded.wal_max_entries, 42);
    }

    #[test]
    fn consistent_store_loads() {
        let loaded = roundtrip(&cache_with(3)).unwrap();
        assert_eq!(loaded.chunks.len(), 3);
        assert_eq!(loaded.norms, vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn wrong_length_norms_are_recomputed_not_trusted() {
        // Regression: the guard used to be `n.len() == n.len()`, so a norms
        // dataset of any length was accepted and corrupted every cosine score.
        let mut cache = cache_with(3);
        cache.norms = vec![99.0];
        let loaded = roundtrip(&cache).unwrap();
        assert_eq!(loaded.norms, vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn mismatched_per_record_datasets_are_schema_errors() {
        type Corrupt = fn(&mut MemoryCache);
        let cases: [(&str, Corrupt); 5] = [
            ("tombstones", |c| c.tombstones.truncate(1)),
            ("timestamps", |c| c.timestamps.truncate(1)),
            ("tags", |c| c.tags.truncate(1)),
            ("session_ids", |c| c.session_ids.truncate(1)),
            ("source_channel", |c| c.source_channels.truncate(1)),
        ];
        for (name, corrupt) in cases {
            let mut cache = cache_with(3);
            corrupt(&mut cache);
            match roundtrip(&cache) {
                Err(MemoryError::Schema(msg)) => {
                    assert!(msg.contains(name), "{name}: unexpected message {msg}")
                }
                other => panic!("{name}: expected Schema error, got {:?}", other.map(|_| ())),
            }
        }
    }
}