mushroomdb-rules 0.6.2

use crate::def::{Predicate, MAX_KEYMATCH_LIST};
use crate::hnsw::HnswIndex;
use core_storage::{list_tokens, Value, ValueKey};
use std::collections::{BTreeMap, BTreeSet};

// ---------------------------------------------------------------------------
// IVF-Flat constants (Plan 11 T4)
// ---------------------------------------------------------------------------

/// Minimum number of k-means clusters. Keeps probing meaningful even for
/// small sides (< 16 vectors).
pub const IVF_K_MIN: usize = 4;

/// Maximum number of k-means clusters. Bounds centroid memory and fit time.
pub const IVF_K_MAX: usize = 1024;

/// Fixed number of k-means iterations per fit (deterministic convergence).
pub const IVF_ITERATIONS: usize = 12;

/// Probe denominator: P = max(1, ceil(k / IVF_PROBE_DENOM)) centroids queried
/// per lookup.  k=4 → P=1; k=64 → P=4; k=1024 → P=64.
pub const IVF_PROBE_DENOM: usize = 16;

/// Rebuild an approximate rule when dst-side IVF drift exceeds this count.
/// Drift is only known after apply, so the WAL path issues `RebuildRule` as a
/// second commit (not a pre-WAL Batch).
pub const IVF_DRIFT_REBUILD: u64 = 256;

thread_local! {
    static IVF_DRIFT_REBUILD_OVERRIDE: std::cell::Cell<Option<u64>> =
        const { std::cell::Cell::new(None) };
}

pub(crate) fn ivf_drift_rebuild_threshold() -> u64 {
    IVF_DRIFT_REBUILD_OVERRIDE.with(|c| c.get().unwrap_or(IVF_DRIFT_REBUILD))
}

/// Run `f` with a temporary IVF dst-drift rebuild threshold.
/// Restores the previous override (including across panics).
pub fn with_ivf_drift_rebuild<R>(threshold: u64, f: impl FnOnce() -> R) -> R {
    IVF_DRIFT_REBUILD_OVERRIDE.with(|c| {
        let prev = c.replace(Some(threshold));
        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
        c.set(prev);
        match out {
            Ok(v) => v,
            Err(p) => std::panic::resume_unwind(p),
        }
    })
}

/// k = ceil(sqrt(n)) clamped to [IVF_K_MIN, IVF_K_MAX].
pub fn cluster_k(n: usize) -> usize {
    if n == 0 {
        return IVF_K_MIN;
    }
    let k = (n as f64).sqrt().ceil() as usize;
    k.clamp(IVF_K_MIN, IVF_K_MAX)
}

/// P = max(1, ceil(k / IVF_PROBE_DENOM)).
pub fn probe_count(k: usize) -> usize {
    k.div_ceil(IVF_PROBE_DENOM).max(1)
}

/// L2-normalize `xs`. Returns `None` for the zero vector (skipped, not clustered).
fn l2_normalize(xs: &[f64]) -> Option<Vec<f64>> {
    let n = xs.iter().map(|x| x * x).sum::<f64>().sqrt();
    if n == 0.0 {
        return None;
    }
    Some(xs.iter().map(|x| x / n).collect())
}

/// Squared Euclidean distance between two equal-length slices.
/// Returns `f64::MAX` on dimension mismatch so callers always have a valid order.
fn l2_sq(a: &[f64], b: &[f64]) -> f64 {
    if a.len() != b.len() {
        return f64::MAX;
    }
    a.iter().zip(b.iter()).map(|(x, y)| (x - y) * (x - y)).sum()
}

/// Index of the nearest centroid to `xs` by L2 distance (minimum squared).
/// Returns 0 when `centroids` is empty.
pub fn nearest_centroid(centroids: &[Vec<f64>], xs: &[f64]) -> usize {
    centroids
        .iter()
        .enumerate()
        .min_by(|(_, a), (_, b)| {
            l2_sq(xs, a)
                .partial_cmp(&l2_sq(xs, b))
                .unwrap_or(std::cmp::Ordering::Equal)
        })
        .map(|(i, _)| i)
        .unwrap_or(0)
}

/// FNV-1a 64-bit hash — stable, documented, NOT DefaultHasher.
/// Used to seed k-means so the same rule name always produces the same
/// clusters on the same data (WAL replay identity).
pub fn fnv1a_u64(data: &[u8]) -> u64 {
    const FNV_OFFSET: u64 = 14_695_981_039_346_656_037;
    const FNV_PRIME: u64 = 1_099_511_628_211;
    let mut h = FNV_OFFSET;
    for &b in data {
        h ^= b as u64;
        h = h.wrapping_mul(FNV_PRIME);
    }
    h
}

/// Seeded LCG step — Knuth multiplicative; used for centroid init and empty
/// cluster reseeding.
#[inline]
fn lcg_next(state: u64) -> u64 {
    state
        .wrapping_mul(6_364_136_223_846_793_005)
        .wrapping_add(1_442_695_040_888_963_407)
}

/// Fit k-means over `vecs` (node_id, vector) pairs.
///
/// - Each vector is L2-normalized before clustering (zero vectors skipped).
/// - `k` is clamped to `min(k, vecs.len())` so we never request more centroids
///   than vectors.
/// - Centroids are initialised by seeded LCG selection without replacement.
/// - 12 iterations; empty clusters are deterministically reseeded from the
///   full dataset.
/// - Returns a `Vec<Vec<f64>>` of k centroids (same length as `xs` entries).
pub fn kmeans_fit(vecs: &[(u32, Vec<f64>)], k: usize, seed: u64) -> Vec<Vec<f64>> {
    let vecs: Vec<(u32, Vec<f64>)> = vecs
        .iter()
        .filter_map(|(id, xs)| l2_normalize(xs).map(|n| (*id, n)))
        .collect();
    if vecs.is_empty() || k == 0 {
        return vec![];
    }
    let n = vecs.len();
    let k = k.min(n);
    let dim = vecs[0].1.len();
    if dim == 0 {
        return vec![];
    }

    // --- Centroid initialisation: pick k distinct indices via seeded LCG ---
    let mut state = seed;
    let mut used = vec![false; n];
    let mut init_idxs: Vec<usize> = Vec::with_capacity(k);
    let mut attempts = 0usize;
    while init_idxs.len() < k && attempts < n * 4 {
        state = lcg_next(state);
        let idx = (state >> 33) as usize % n;
        if !used[idx] {
            used[idx] = true;
            init_idxs.push(idx);
        }
        attempts += 1;
    }
    // If LCG didn't yield k distinct indices (pathological: n very small or
    // many collisions), fill sequentially.
    if init_idxs.len() < k {
        for (i, in_use) in used.iter().enumerate().take(n) {
            if !in_use {
                init_idxs.push(i);
                if init_idxs.len() == k {
                    break;
                }
            }
        }
    }
    let mut centroids: Vec<Vec<f64>> = init_idxs.iter().map(|&i| vecs[i].1.clone()).collect();
    let mut assignments = vec![0usize; n];

    // --- k-means iterations ---
    for iter in 0..IVF_ITERATIONS {
        // Assignment step
        for (j, (_, xs)) in vecs.iter().enumerate() {
            assignments[j] = nearest_centroid(&centroids, xs);
        }

        // Update step: accumulate sums and counts
        let mut sums = vec![vec![0.0f64; dim]; k];
        let mut counts = vec![0usize; k];
        for (j, (_, xs)) in vecs.iter().enumerate() {
            let c = assignments[j];
            counts[c] += 1;
            for d in 0..dim {
                sums[c][d] += xs[d];
            }
        }

        // Compute new centroids; collect empty ones for reseed
        let mut new_centroids = vec![vec![0.0f64; dim]; k];
        let mut empty: Vec<usize> = Vec::new();
        for c in 0..k {
            if counts[c] == 0 {
                empty.push(c);
            } else {
                for d in 0..dim {
                    new_centroids[c][d] = sums[c][d] / counts[c] as f64;
                }
            }
        }

        // Deterministic empty-cluster reseed: pick a vector from the dataset
        // seeded by (original seed XOR iteration XOR empty-cluster-index).
        for (ei, ec) in empty.into_iter().enumerate() {
            let reseed =
                seed ^ (iter as u64).wrapping_mul(0x9E37) ^ (ei as u64).wrapping_mul(0x1234_5679);
            let mut rs = lcg_next(reseed);
            rs = lcg_next(rs);
            let pick = (rs >> 33) as usize % n;
            new_centroids[ec] = vecs[pick].1.clone();
        }

        centroids = new_centroids;
    }

    centroids
}

#[cfg(test)]
thread_local! {
    static VECTOR_DIM_REJECT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
    static VECTOR_EARLY_EXIT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
}

fn vector_dim_reject_enabled() -> bool {
    #[cfg(test)]
    {
        VECTOR_DIM_REJECT.with(|c| c.get())
    }
    #[cfg(not(test))]
    {
        true
    }
}

pub(crate) fn vector_early_exit_enabled() -> bool {
    #[cfg(test)]
    {
        VECTOR_EARLY_EXIT.with(|c| c.get())
    }
    #[cfg(not(test))]
    {
        true
    }
}

/// Force the ScanAll dim fast-reject on or off. Identity-proof hook.
#[cfg(test)]
pub fn with_vector_dim_reject<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
    VECTOR_DIM_REJECT.with(|c| {
        let prev = c.replace(enabled);
        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
        c.set(prev);
        match out {
            Ok(v) => v,
            Err(p) => std::panic::resume_unwind(p),
        }
    })
}

/// Force the checkpointed Cauchy-Schwarz early-exit on or off. Identity-proof hook.
#[cfg(test)]
pub fn with_vector_early_exit<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
    VECTOR_EARLY_EXIT.with(|c| {
        let prev = c.replace(enabled);
        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
        c.set(prev);
        match out {
            Ok(v) => v,
            Err(p) => std::panic::resume_unwind(p),
        }
    })
}

#[derive(Debug, Default)]
pub struct SideIndex {
    by_key: BTreeMap<ValueKey, BTreeSet<u32>>,
    /// Per-node `(dim, L2 norm)` for `ScanAll` members. Maintained by the
    /// same insert/remove choke-points as `by_key`. Cosine still reads live
    /// props; `dim` is a fast-reject; `norm` is the primary freshness gate for
    /// the checkpointed Cauchy-Schwarz early-exit (Plan 11 T3).
    vec_meta: BTreeMap<u32, (u32, f64)>,
    /// Per-node checkpointed suffix norms for the Cauchy-Schwarz early-exit.
    /// `ckpts[i]` = L2 norm of `xs[i * dim / 8 ..]`.
    /// `ckpts[0]` = full L2 norm; `ckpts[7]` = norm of the last eighth.
    /// Built at index-insert, torn out at index-remove — maintained in lockstep
    /// with `vec_meta` by the same choke-points.
    /// Memory: 8 × 8 = 64 bytes per indexed vector (6.4 MB at 100k vectors).
    vec_checkpoints: BTreeMap<u32, [f64; 8]>,
    /// Per-node first element (`xs[0]`) for heuristic permutation detection.
    /// A permuted vector can share `(dim, norm)` with the indexed one but
    /// differs at `xs[0]` in virtually all realistic cases, so comparing this
    /// one extra f64 (8 bytes per vector) breaks same-norm permutation aliasing
    /// cheaply.  This is heuristic hardening — not a proof — but eliminates
    /// the energy-distribution construction identified in the Plan 11 T3 review.
    vec_anchor: BTreeMap<u32, f64>,

    // --- IVF-Flat fields (Plan 11 T4; only populated for VectorClusters specs) ---
    /// Raw vectors stored for IVF fitting and assignment-on-insert.
    /// Populated at insert-time, torn out at remove-time.
    /// Memory: O(n × dim) per indexed side — present only for approximate rules.
    ivf_raw: BTreeMap<u32, Vec<f64>>,
    /// Fitted k-means centroids (empty until after first `fit_ivf_clusters` call).
    ivf_centroids: Vec<Vec<f64>>,
    /// Per-node cluster assignment post-fit.
    /// `by_key[ivf_cluster_key(cluster)] → {node_ids}`.
    ivf_clusters: BTreeMap<u32, usize>,
    /// Count of vector inserts/removes since last fit.  When dst-side drift
    /// exceeds [`IVF_DRIFT_REBUILD`] on an approximate rule, apply queues a
    /// `RebuildRule` second commit (fit resets this to zero).
    pub ivf_drift: u64,

    // --- HNSW fields (default for approximate: true + VectorSimilar) ---
    /// HNSW graph; `None` until `init_hnsw` is called.
    hnsw: Option<HnswIndex>,
    /// All node ids inserted via `CandidateSpec::Hnsw`.
    /// Used as a full-scan fallback when `hnsw` is `None` or has no entry point.
    hnsw_tracked: BTreeSet<u32>,
}

#[derive(Debug, Default)]
pub struct RuleIndex {
    pub src_side: SideIndex,
    pub dst_side: SideIndex,
}

#[derive(Debug)]
pub enum CandidateSpec<'a> {
    ByKey,
    Scalar {
        field: &'a str,
    },
    Tokens {
        field: &'a str,
    },
    /// Src side of a `KeyMatch`-rooted rule: the FK field's scalar value, or —
    /// when that value is a list — one bucket per **string** element (the first
    /// [`MAX_KEYMATCH_LIST`] in stored order, non-strings skipped).
    ///
    /// The reverse lookup always probes with a single key (the destination
    /// node's key), so a multi-valued FK has to fan out at index time: a src
    /// node listing n keys sits in n buckets and is found through any of them.
    /// A scalar value indexes exactly as [`CandidateSpec::Scalar`] does.
    ScalarOrElements {
        field: &'a str,
    },
    NumericBucket {
        field: &'a str,
        tolerance: f64,
    },
    GeoGrid {
        field: &'a str,
        km: f64,
    },
    ScanAll {
        field: &'a str,
    },
    /// IVF-Flat approximate candidate selection (legacy; still supported as
    /// direct fallback — no longer the default for `approximate: true`).
    ///
    /// k-means fitted over the indexed side's vectors; candidates are members
    /// of the P = `max(1, ceil(k/16))` nearest centroids to the query vector.
    /// NOT a superset of true positives — recall floor governs correctness.
    VectorClusters {
        field: &'a str,
        min: f64,
    },
    /// HNSW approximate candidate selection (default for `approximate: true`).
    ///
    /// Returns the `k` nearest vectors by cosine similarity from the in-tree
    /// HNSW graph.  Falls back to returning all tracked nodes when the graph
    /// has no entry point (e.g. before any node is inserted, or when used
    /// without calling `init_hnsw`).
    Hnsw {
        field: &'a str,
        /// Number of approximate candidates to return; typically
        /// `max(max_edges, 64)` from the owning `RuleDef`.
        k: usize,
    },
    /// Union of multiple candidate specs, used for `Any` predicates.
    ///
    /// Each branch of the `Any` predicate contributes its own candidate set
    /// (key index, token index, numeric bucket, etc.); the resulting candidate
    /// set is their union.  Insert and remove recurse into every child spec so
    /// the index stays coherent for all branches simultaneously.
    Union(Vec<CandidateSpec<'a>>),
    /// Intersection of multiple candidate specs, used for `All` predicates.
    ///
    /// Each conjunct contributes its own candidate set; the result is their
    /// intersection (empty child → empty). `ScanAll` children are skipped at
    /// probe time (they are the universe); if every child is `ScanAll`, the
    /// spec stays a full scan. Insert/remove recurse into every child.
    Intersect(Vec<CandidateSpec<'a>>),
}

/// Returns the exact candidate strategy derived from `p`.
///
/// `All(parts)` returns `Intersect` of each part's spec. A leading
/// `VectorSimilar` is `ScanAll` and is skipped at probe time when another
/// conjunct has an index; candidates stay a superset of true matches.
///
/// `Any(parts)` returns `Union` of each branch's candidate spec — the correct
/// superset for OR semantics.
///
/// # Panics
///
/// Panics on `All([])` or `Any([])`. Predicates must pass `RuleDef::validate()` first.
pub fn candidate_spec(p: &Predicate) -> CandidateSpec<'_> {
    match p {
        Predicate::KeyMatch { .. } => CandidateSpec::ByKey,
        Predicate::FieldEqual { field } => CandidateSpec::Scalar { field },
        Predicate::Overlap { field, .. } => CandidateSpec::Tokens { field },
        Predicate::NumericWithin { field, tolerance } => CandidateSpec::NumericBucket {
            field,
            tolerance: *tolerance,
        },
        Predicate::GeoRadius { field, km } => CandidateSpec::GeoGrid { field, km: *km },
        Predicate::VectorSimilar { field, .. } => CandidateSpec::ScanAll { field },
        Predicate::All(parts) => {
            debug_assert!(
                !parts.is_empty(),
                "candidate_spec requires a validated predicate"
            );
            CandidateSpec::Intersect(parts.iter().map(candidate_spec).collect())
        }
        Predicate::Any(parts) => {
            debug_assert!(
                !parts.is_empty(),
                "candidate_spec requires a validated predicate"
            );
            CandidateSpec::Union(parts.iter().map(candidate_spec).collect())
        }
    }
}

/// Approximate candidate strategy: like `candidate_spec` but replaces
/// `ScanAll` with `CandidateSpec::Hnsw` for `VectorSimilar`-rooted predicates.
///
/// Used when `RuleDef::approximate == true`.  `All` is `Intersect` of each
/// child's approx spec (not `parts[0]`), so `FieldEqual` / `NumericWithin`
/// conjuncts still probe their indexes.
///
/// `k` is the number of HNSW candidates to return; callers should use
/// `max(max_edges, 64)` from the owning `RuleDef`.  Use the public
/// zero-argument wrapper (`candidate_spec_approx`) for tests that don't
/// need a specific k (defaults to 64).
///
/// `Any` predicates cannot be `approximate=true` (validate() rejects them),
/// so `Any` falls through to `candidate_spec` (exact Union path).
///
/// # Panics
///
/// Panics on `All([])` or `Any([])`. Predicates must pass `RuleDef::validate()` first.
pub fn candidate_spec_approx(p: &Predicate) -> CandidateSpec<'_> {
    candidate_spec_approx_with_k(p, 64)
}

/// Like `candidate_spec_approx` but with an explicit HNSW candidate count `k`.
pub fn candidate_spec_approx_with_k(p: &Predicate, k: usize) -> CandidateSpec<'_> {
    match p {
        Predicate::VectorSimilar { field, .. } => CandidateSpec::Hnsw { field, k },
        Predicate::All(parts) => {
            debug_assert!(
                !parts.is_empty(),
                "candidate_spec_approx requires a validated predicate"
            );
            CandidateSpec::Intersect(
                parts
                    .iter()
                    .map(|p| candidate_spec_approx_with_k(p, k))
                    .collect(),
            )
        }
        other => candidate_spec(other),
    }
}

pub(crate) fn as_finite_f64(v: &Value) -> Option<f64> {
    match v {
        Value::Int(i) => Some(*i as f64),
        Value::Float(f) if f.is_finite() => Some(*f),
        _ => None,
    }
}

fn as_latlon(v: &Value) -> Option<(f64, f64)> {
    let Value::List(items) = v else {
        return None;
    };
    if items.len() != 2 {
        return None;
    }
    let lat = as_finite_f64(&items[0])?;
    let lon = as_finite_f64(&items[1])?;
    if (-90.0..=90.0).contains(&lat) && (-180.0..=180.0).contains(&lon) {
        Some((lat, lon))
    } else {
        None
    }
}

pub(crate) fn as_numeric_list(v: &Value) -> Option<Vec<f64>> {
    let Value::List(items) = v else {
        return None;
    };
    if items.is_empty() {
        return None;
    }
    items.iter().map(as_finite_f64).collect()
}

fn vec_dim_norm(v: &Value) -> Option<(u32, f64)> {
    let xs = as_numeric_list(v)?;
    let mut n2 = 0.0;
    for x in &xs {
        n2 += *x * *x;
    }
    Some((xs.len() as u32, n2.sqrt()))
}

/// Checkpointed suffix norms for Cauchy-Schwarz early exit.
///
/// `ckpts[i]` = L2 norm of `xs[boundary(i)..]` where `boundary(i) = i * dim / 8`.
/// `ckpts[0]` equals the full L2 norm; `ckpts[7]` is the last eighth's norm.
/// Multiple checkpoints may share the same boundary for dim < 8 (correct but no-op).
fn compute_ckpts(xs: &[f64]) -> [f64; 8] {
    let dim = xs.len();
    let mut ckpts = [0.0f64; 8];
    if dim == 0 {
        return ckpts;
    }
    // boundaries[i] = i * dim / 8 (integer division).
    let boundaries: [usize; 8] = std::array::from_fn(|i| i * dim / 8);
    let mut suffix_sq = 0.0f64;
    // Walk right-to-left; ci is the highest checkpoint not yet recorded.
    let mut ci = 7i32;
    for j in (0..dim).rev() {
        suffix_sq += xs[j] * xs[j];
        // Assign all checkpoints whose boundary equals j.
        while ci >= 0 && boundaries[ci as usize] == j {
            ckpts[ci as usize] = suffix_sq.sqrt();
            ci -= 1;
        }
    }
    ckpts
}

fn floor_to_i64(x: f64) -> i64 {
    let floored = x.floor();
    if !floored.is_finite() {
        return 0;
    }
    if floored >= i64::MAX as f64 {
        i64::MAX
    } else if floored <= i64::MIN as f64 {
        i64::MIN
    } else {
        floored as i64
    }
}

/// Two values within `tolerance` always land in adjacent buckets
/// (`|floor(a/tol) − floor(b/tol)| ≤ 1`), so probing `{b−1, b, b+1}` is a
/// superset of every evaluate-match.
fn numeric_index_key(v: f64, tolerance: f64) -> Option<ValueKey> {
    if !tolerance.is_finite() || tolerance < 0.0 {
        return None;
    }
    if tolerance == 0.0 {
        let v = if v == 0.0 { 0.0_f64 } else { v };
        return Some(ValueKey::FloatBits(v.to_bits()));
    }
    Some(ValueKey::Int(floor_to_i64(v / tolerance)))
}

fn numeric_probe_keys(v: f64, tolerance: f64) -> BTreeSet<ValueKey> {
    match numeric_index_key(v, tolerance) {
        None => BTreeSet::new(),
        Some(k @ ValueKey::FloatBits(_)) => BTreeSet::from([k]),
        Some(ValueKey::Int(b)) => BTreeSet::from([
            ValueKey::Int(b.saturating_sub(1)),
            ValueKey::Int(b),
            ValueKey::Int(b.saturating_add(1)),
        ]),
        Some(other) => BTreeSet::from([other]),
    }
}

fn geo_cell(lat: f64, lon: f64, km: f64) -> Option<(i64, i64, f64, i64)> {
    if !km.is_finite() || km <= 0.0 {
        return None;
    }
    let cell_deg = (km / 111.0).max(1e-6);
    let gx = floor_to_i64(lat / cell_deg);
    // Longitude wraps; lat does not (validated range, no pole crossing
    // within the supported |lat|≲87 envelope — see cos clamp below).
    let lon_cells = (360.0 / cell_deg).ceil() as i64;
    let lon_cells = lon_cells.max(1);
    let gy = floor_to_i64(lon / cell_deg).rem_euclid(lon_cells);
    Some((gx, gy, cell_deg, lon_cells))
}

fn geo_index_key(lat: f64, lon: f64, km: f64) -> Option<ValueKey> {
    let (gx, gy, _, _) = geo_cell(lat, lon, km)?;
    Some(ValueKey::Str(format!("{gx}|{gy}")))
}

fn geo_probe_keys(lat: f64, lon: f64, km: f64) -> BTreeSet<ValueKey> {
    let Some((gx, gy, cell_deg, lon_cells)) = geo_cell(lat, lon, km) else {
        return BTreeSet::new();
    };
    // Cos clamp keeps the probe a superset up to |lat| ≈ 87.
    let cos_lat = lat.to_radians().cos().max(0.05);
    let n = ((km / (111.0 * cos_lat)) / cell_deg).ceil();
    let n = if n.is_finite() {
        floor_to_i64(n).max(0)
    } else {
        0
    };
    let mut out = BTreeSet::new();
    for dx in -1..=1 {
        for dy in -n..=n {
            let cx = gx.saturating_add(dx);
            let cy = gy.saturating_add(dy).rem_euclid(lon_cells);
            out.insert(ValueKey::Str(format!("{cx}|{cy}")));
        }
    }
    out
}

/// Vector candidates are a deliberate full scan of opposite-side
/// vector-bearing nodes; ANN is Plan 8+.
const SCAN_ALL_SENTINEL: ValueKey = ValueKey::Bool(true);

/// IVF cluster buckets in `by_key`. SOH prefix keeps them off the Int space
/// used by `NumericBucket` / integer `FieldEqual` and off token/geo Str keys.
fn ivf_cluster_key(cluster: usize) -> ValueKey {
    ValueKey::Str(format!("\u{1}ivf:{cluster}"))
}

/// `ScanAll` is the universe in an `Intersect`: skip it when another child
/// has an index. Nested `Intersect` of only `ScanAll` is itself a universe.
fn spec_is_scan_all_universe(spec: &CandidateSpec<'_>) -> bool {
    match spec {
        CandidateSpec::ScanAll { .. } => true,
        CandidateSpec::Intersect(parts) => {
            !parts.is_empty() && parts.iter().all(spec_is_scan_all_universe)
        }
        _ => false,
    }
}

/// `ByKey` is resolved by `compute_desired` (FK id lookup), not `by_key`.
/// Nested `Intersect` of only `ByKey` is likewise external.
fn spec_is_bykey_external(spec: &CandidateSpec<'_>) -> bool {
    match spec {
        CandidateSpec::ByKey => true,
        CandidateSpec::Intersect(parts) => {
            !parts.is_empty() && parts.iter().all(spec_is_bykey_external)
        }
        _ => false,
    }
}

impl SideIndex {
    fn index_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
        match spec {
            CandidateSpec::ByKey => BTreeSet::new(),
            CandidateSpec::Scalar { field } => get(field)
                .as_ref()
                .and_then(ValueKey::from_value)
                .into_iter()
                .collect(),
            CandidateSpec::Tokens { field } => get(field)
                .as_ref()
                .and_then(list_tokens)
                .unwrap_or_default(),
            CandidateSpec::ScalarOrElements { field } => match get(field) {
                Some(Value::List(items)) => items
                    .iter()
                    .take(MAX_KEYMATCH_LIST)
                    .filter(|v| matches!(v, Value::Str(_)))
                    .filter_map(ValueKey::from_value)
                    .collect(),
                Some(v) => ValueKey::from_value(&v).into_iter().collect(),
                None => BTreeSet::new(),
            },
            CandidateSpec::NumericBucket { field, tolerance } => get(field)
                .as_ref()
                .and_then(as_finite_f64)
                .and_then(|v| numeric_index_key(v, *tolerance))
                .into_iter()
                .collect(),
            CandidateSpec::GeoGrid { field, km } => get(field)
                .as_ref()
                .and_then(as_latlon)
                .and_then(|(lat, lon)| geo_index_key(lat, lon, *km))
                .into_iter()
                .collect(),
            CandidateSpec::ScanAll { field } => get(field)
                .as_ref()
                .and_then(as_numeric_list)
                .map(|_| SCAN_ALL_SENTINEL)
                .into_iter()
                .collect(),
            // VectorClusters uses ivf_raw / ivf_clusters, not by_key. The
            // insert() path returns early before reaching index_keys for this
            // variant, so this arm is unreachable at runtime; it must be
            // present to satisfy exhaustiveness.
            CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
            // Hnsw uses the separate hnsw / hnsw_tracked fields, not by_key.
            CandidateSpec::Hnsw { .. } => BTreeSet::new(),
            // Union: each branch contributes its own index keys; the result is
            // their union.  VectorClusters/Hnsw children are handled by the
            // early-return in insert()/remove().
            CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
                let mut out = BTreeSet::new();
                for s in specs {
                    out.extend(Self::index_keys(s, get));
                }
                out
            }
        }
    }

    fn probe_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
        match spec {
            CandidateSpec::ByKey
            | CandidateSpec::Scalar { .. }
            | CandidateSpec::Tokens { .. }
            | CandidateSpec::ScalarOrElements { .. } => Self::index_keys(spec, get),
            CandidateSpec::NumericBucket { field, tolerance } => get(field)
                .as_ref()
                .and_then(as_finite_f64)
                .map(|v| numeric_probe_keys(v, *tolerance))
                .unwrap_or_default(),
            CandidateSpec::GeoGrid { field, km } => get(field)
                .as_ref()
                .and_then(as_latlon)
                .map(|(lat, lon)| geo_probe_keys(lat, lon, *km))
                .unwrap_or_default(),
            CandidateSpec::ScanAll { field } => get(field)
                .as_ref()
                .and_then(as_numeric_list)
                .map(|_| SCAN_ALL_SENTINEL)
                .into_iter()
                .collect(),
            // VectorClusters probing is handled by ivf_candidates(), not probe_keys().
            CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
            // Hnsw probing is handled by hnsw_candidates(), not probe_keys().
            CandidateSpec::Hnsw { .. } => BTreeSet::new(),
            // Union / Intersect: probe each child. `candidates()` intersects
            // Intersect node-sets; mixing keys here is only for insert/remove.
            CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
                let mut out = BTreeSet::new();
                for s in specs {
                    out.extend(Self::probe_keys(s, get));
                }
                out
            }
        }
    }

    pub fn insert(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
        // Union / Intersect: recurse into each child spec. insert() is
        // idempotent for ScanAll metadata (same-value overwrite).
        if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
            for s in specs {
                self.insert(s, node, get);
            }
            return;
        }
        // Hnsw: maintain hnsw_tracked for fallback, and hnsw graph if initialized.
        if let CandidateSpec::Hnsw { field, .. } = spec {
            if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
                self.hnsw_tracked.insert(node);
                if let Some(h) = &mut self.hnsw {
                    h.insert(node, &xs);
                }
            }
            return;
        }
        // VectorClusters: IVF path — separate from the by_key / ScanAll path.
        if let CandidateSpec::VectorClusters { field, .. } = spec {
            if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
                self.ivf_raw.insert(node, xs.clone());
                if !self.ivf_centroids.is_empty() {
                    // Assign in cosine space (centroids are unit-norm). Skip zeros.
                    if let Some(unit) = l2_normalize(&xs) {
                        let c = nearest_centroid(&self.ivf_centroids, &unit);
                        self.ivf_clusters.insert(node, c);
                        self.by_key
                            .entry(ivf_cluster_key(c))
                            .or_default()
                            .insert(node);
                    }
                    self.ivf_drift = self.ivf_drift.saturating_add(1);
                }
            }
            return;
        }

        for k in Self::index_keys(spec, get) {
            self.by_key.entry(k).or_default().insert(node);
        }
        if let CandidateSpec::ScanAll { field } = spec {
            if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
                let mut n2 = 0.0f64;
                for x in &xs {
                    n2 += x * x;
                }
                let norm = n2.sqrt();
                self.vec_meta.insert(node, (xs.len() as u32, norm));
                self.vec_checkpoints.insert(node, compute_ckpts(&xs));
                // xs is non-empty (as_numeric_list rejects empty lists).
                self.vec_anchor.insert(node, xs[0]);
            }
        }
    }

    pub fn remove(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
        // Union / Intersect: recurse into each child spec.
        if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
            for s in specs {
                self.remove(s, node, get);
            }
            return;
        }
        // Hnsw: remove from hnsw_tracked and hnsw graph.  Increment ivf_drift
        // as a deletion counter so maybe_queue_ivf_rebuild fires at the same
        // cadence it did for IVF rules; the resulting rebuild re-scans all nodes
        // and optionally compacts the HNSW graph.
        if let CandidateSpec::Hnsw { field, .. } = spec {
            if get(field).as_ref().and_then(as_numeric_list).is_some() {
                self.hnsw_tracked.remove(&node);
                if let Some(h) = &mut self.hnsw {
                    h.remove(node);
                }
                self.ivf_drift = self.ivf_drift.saturating_add(1);
            }
            return;
        }
        // VectorClusters: remove from ivf_raw and by_key cluster bucket.
        // Removal shifts cluster membership (the centroid stays but its member set
        // shrinks), which is a form of drift; increment the counter so callers can
        // decide when to trigger a rebuild.
        if let CandidateSpec::VectorClusters { .. } = spec {
            if self.ivf_raw.remove(&node).is_some() {
                self.ivf_drift = self.ivf_drift.saturating_add(1);
                if let Some(c) = self.ivf_clusters.remove(&node) {
                    let key = ivf_cluster_key(c);
                    if let Some(s) = self.by_key.get_mut(&key) {
                        s.remove(&node);
                        if s.is_empty() {
                            self.by_key.remove(&key);
                        }
                    }
                }
            }
            return;
        }

        for k in Self::index_keys(spec, get) {
            if let Some(set) = self.by_key.get_mut(&k) {
                set.remove(&node);
                if set.is_empty() {
                    self.by_key.remove(&k);
                }
            }
        }
        if let CandidateSpec::ScanAll { field } = spec {
            if get(field).as_ref().and_then(as_numeric_list).is_some() {
                self.vec_meta.remove(&node);
                self.vec_checkpoints.remove(&node);
                self.vec_anchor.remove(&node);
            }
        }
    }

    /// Cached vector dimension for a `ScanAll` member, if present.
    pub fn vec_dim(&self, node: u32) -> Option<u32> {
        self.vec_meta.get(&node).map(|(d, _)| *d)
    }

    /// Cached `(dim, L2 norm)` for tests / debug.
    pub fn vec_meta(&self, node: u32) -> Option<(u32, f64)> {
        self.vec_meta.get(&node).copied()
    }

    /// Cached checkpoints for tests / debug.
    pub fn vec_ckpts(&self, node: u32) -> Option<&[f64; 8]> {
        self.vec_checkpoints.get(&node)
    }

    /// Returns `(cached_norm, &checkpoints)` if the cached state matches the
    /// live vector under all three freshness checks.
    ///
    /// # Stale-cache gate
    ///
    /// Stale checkpoints (from a vector that differs from `live`) can produce
    /// **false rejects** — the Cauchy-Schwarz suffix bound may be under-tight
    /// for the live vector's actual energy distribution.  Three guards defend
    /// against this in ascending selectivity order:
    ///
    /// 1. **Dim check** — `cached_dim == live.len()`.  Different lengths →
    ///    immediate fallback.
    /// 2. **Norm check** — recomputes L2 norm with the same sequential
    ///    accumulation used at insert time so bits are identical for an unchanged
    ///    vector.  Changed norm → fallback.
    /// 3. **Anchor check** — compares `xs[0]` against the cached first element.
    ///    A permuted vector can share `(dim, norm)` with the indexed one but
    ///    differ at `xs[0]`, breaking the most realistic same-norm aliasing
    ///    attack.  This is **heuristic hardening**, not a proof: a permutation
    ///    that preserves `xs[0]` would still pass, but is vanishingly unlikely
    ///    in practice.
    ///
    /// The real coherence guarantee is structural: checkpoint rebuilds flow
    /// through the same insert/remove choke-points as `vec_meta`, so in
    /// normal single-writer operation the cache is always coherent.  These
    /// gates are belt-and-suspenders against bugs in those choke-points.
    pub(crate) fn fresh_ckpts_for<'a>(
        &'a self,
        node: u32,
        live: &[f64],
    ) -> Option<(f64, &'a [f64; 8])> {
        let &(dim, norm) = self.vec_meta.get(&node)?;
        if dim != live.len() as u32 {
            return None;
        }
        // Compute the live norm with the same sequential accumulation used at
        // insert time so the bits are identical when the vector is unchanged.
        let live_norm = {
            let mut n2 = 0.0f64;
            for x in live {
                n2 += x * x;
            }
            n2.sqrt()
        };
        if norm != live_norm {
            return None; // stale — fall back to brute-force evaluate()
        }
        // Heuristic anchor check: first element breaks same-norm permutation
        // aliasing in virtually all realistic cases.  dim > 0 guaranteed (dim
        // was stored from non-empty xs; live.len() == dim > 0).
        let live_anchor = live[0];
        let &cached_anchor = self.vec_anchor.get(&node)?;
        if live_anchor != cached_anchor {
            return None;
        }
        let ckpts = self.vec_checkpoints.get(&node)?;
        Some((norm, ckpts))
    }

    pub fn candidates(
        &self,
        spec: &CandidateSpec,
        get: &dyn Fn(&str) -> Option<Value>,
    ) -> BTreeSet<u32> {
        // Hnsw: approximate nearest-neighbor search.
        if let CandidateSpec::Hnsw { field, k } = spec {
            return self.hnsw_candidates(field, *k, get);
        }
        // VectorClusters: probe the P nearest centroids.
        if let CandidateSpec::VectorClusters { field, .. } = spec {
            return self.ivf_candidates(field, get);
        }
        // Union: take the union of candidates from each child spec.
        if let CandidateSpec::Union(specs) = spec {
            return specs.iter().flat_map(|s| self.candidates(s, get)).collect();
        }
        if let CandidateSpec::Intersect(specs) = spec {
            return self.intersect_candidates(specs, get);
        }

        let mut out = BTreeSet::new();
        for k in Self::probe_keys(spec, get) {
            if let Some(set) = self.by_key.get(&k) {
                out.extend(set.iter().copied());
            }
        }
        // Exact: VectorSimilar evaluate is None when dims differ.
        if vector_dim_reject_enabled() {
            if let CandidateSpec::ScanAll { field } = spec {
                if let Some((dim, _)) = get(field).as_ref().and_then(vec_dim_norm) {
                    out.retain(|id| self.vec_meta.get(id).is_none_or(|(d, _)| *d == dim));
                }
            }
        }
        out
    }

    /// Intersect child candidate sets. `ScanAll` is the universe (skipped);
    /// if every child is `ScanAll`, fall back to `ScanAll`. `ByKey` is resolved
    /// outside the index. Empty child → empty.
    fn intersect_candidates(
        &self,
        specs: &[CandidateSpec<'_>],
        get: &dyn Fn(&str) -> Option<Value>,
    ) -> BTreeSet<u32> {
        let mut restrictive = Vec::new();
        let mut scan_alls = Vec::new();
        for s in specs {
            if spec_is_scan_all_universe(s) {
                scan_alls.push(s);
            } else if spec_is_bykey_external(s) {
                continue;
            } else {
                restrictive.push(s);
            }
        }
        let to_intersect: &[&CandidateSpec<'_>] = if !restrictive.is_empty() {
            &restrictive
        } else if !scan_alls.is_empty() {
            &scan_alls
        } else {
            return BTreeSet::new();
        };
        let mut iter = to_intersect.iter();
        let Some(first) = iter.next() else {
            return BTreeSet::new();
        };
        let mut acc = self.candidates(first, get);
        if acc.is_empty() {
            return acc;
        }
        for s in iter {
            let other = self.candidates(s, get);
            if other.is_empty() {
                return BTreeSet::new();
            }
            acc = acc.intersection(&other).copied().collect();
            if acc.is_empty() {
                return acc;
            }
        }
        acc
    }

    /// IVF candidate lookup: find the P nearest centroids to the query vector,
    /// return the union of their cluster members.
    fn ivf_candidates(&self, field: &str, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<u32> {
        let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
            return BTreeSet::new();
        };
        if self.ivf_centroids.is_empty() {
            // Not yet fitted (e.g. empty side at create time, or no data).
            // Fall back to full scan so early crash-recovery states don't drop recall
            // to zero when too few vectors were inserted for IVF to be meaningful.
            return self.ivf_raw.keys().copied().collect();
        }
        // When n ≤ k (actual centroid count), every node is its own centroid;
        // P probes only return the src's own cluster (which excludes itself),
        // yielding zero candidates. Full scan is correct and O(n) for these
        // tiny sets — this covers n < IVF_K_MIN and the exact n == k edge case.
        if self.ivf_raw.len() <= self.ivf_centroids.len() {
            return self.ivf_raw.keys().copied().collect();
        }
        let k = self.ivf_centroids.len();
        let p = probe_count(k);

        // Probe in cosine space (same as centroid fit). Zero query → no candidates
        // (cosine with a zero vector is undefined; exact evaluate also returns None).
        let Some(xs) = l2_normalize(&xs) else {
            return BTreeSet::new();
        };

        // Rank centroids by L2 distance to the unit query; take top-P.
        let mut dists: Vec<(usize, f64)> = self
            .ivf_centroids
            .iter()
            .enumerate()
            .map(|(i, c)| (i, l2_sq(&xs, c)))
            .collect();
        dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));

        let mut out = BTreeSet::new();
        for (ci, _) in dists.iter().take(p) {
            let key = ivf_cluster_key(*ci);
            if let Some(nodes) = self.by_key.get(&key) {
                out.extend(nodes.iter().copied());
            }
        }
        out
    }

    /// Fit (or re-fit) the IVF k-means index for this side using all currently
    /// stored raw vectors.  Called by the engine after reindexing all nodes in
    /// `create_rule` and `rebuild`.
    ///
    /// `rule_name` is hashed via FNV-1a to produce a stable seed, ensuring the
    /// same rule+data always yields the same clusters (WAL replay identity).
    ///
    /// Clears all existing cluster assignments and by_key cluster entries, then
    /// assigns every non-zero vector (L2-normalized) to its nearest new centroid.
    /// Resets `ivf_drift` to zero.
    pub fn fit_ivf_clusters(&mut self, rule_name: &str) {
        if self.ivf_raw.is_empty() {
            self.ivf_centroids.clear();
            self.ivf_clusters.clear();
            self.ivf_drift = 0;
            return;
        }

        // Clear old cluster → node mappings from by_key (namespaced IVF keys).
        for c in self.ivf_clusters.values() {
            self.by_key.remove(&ivf_cluster_key(*c));
        }
        self.ivf_clusters.clear();

        // Gather vectors in deterministic order (BTreeMap → sorted by node id).
        let vecs: Vec<(u32, Vec<f64>)> = self
            .ivf_raw
            .iter()
            .map(|(&id, xs)| (id, xs.clone()))
            .collect();

        let n = vecs.len();
        let k = cluster_k(n);
        let seed = fnv1a_u64(rule_name.as_bytes());

        self.ivf_centroids = kmeans_fit(&vecs, k, seed);

        // Assign in cosine space (skip zeros; they stay in ivf_raw but unclustered).
        for (node, xs) in &vecs {
            let Some(unit) = l2_normalize(xs) else {
                continue;
            };
            let c = nearest_centroid(&self.ivf_centroids, &unit);
            self.ivf_clusters.insert(*node, c);
            self.by_key
                .entry(ivf_cluster_key(c))
                .or_default()
                .insert(*node);
        }
        self.ivf_drift = 0;
    }

    /// Number of fitted centroids (0 = not yet fitted).
    pub fn ivf_k(&self) -> usize {
        self.ivf_centroids.len()
    }

    /// Cluster assignment for a node (None if not fitted or node not in index).
    pub fn ivf_cluster_of(&self, node: u32) -> Option<usize> {
        self.ivf_clusters.get(&node).copied()
    }

    /// Export IVF state for snapshot persistence: (centroids, clusters, drift).
    ///
    /// The caller stores this in the V4 snapshot and passes it back to
    /// `load_ivf_state` on the next open, avoiding a full k-means re-fit.
    pub fn export_ivf_state(&self) -> (Vec<Vec<f64>>, BTreeMap<u32, usize>, u64) {
        (
            self.ivf_centroids.clone(),
            self.ivf_clusters.clone(),
            self.ivf_drift,
        )
    }

    /// Restore IVF state from a V4 snapshot.
    ///
    /// This must be called AFTER the normal `insert()` pass (which populates
    /// `ivf_raw`) but INSTEAD OF `fit_ivf_clusters`.  It:
    ///   1. Removes any stale cluster-key entries from `by_key`.
    ///   2. Installs the persisted centroids and drift counter.
    ///   3. Rebuilds `by_key` cluster buckets from the persisted assignments.
    ///
    /// Nodes present in `ivf_raw` but absent from `clusters` (e.g. inserted
    /// post-snapshot via WAL replay before this is called) are left unassigned;
    /// `on_node_changed` will assign them to the nearest centroid incrementally.
    pub fn load_ivf_state(
        &mut self,
        centroids: Vec<Vec<f64>>,
        clusters: BTreeMap<u32, usize>,
        drift: u64,
    ) {
        // Precondition: ivf_clusters is empty when called from reindex_all_load_ivf (indexes reset to default); loop is defensive for any future direct-call path.
        // Remove old cluster bucket entries from by_key.
        for c in self.ivf_clusters.values() {
            self.by_key.remove(&ivf_cluster_key(*c));
        }
        self.ivf_clusters.clear();

        self.ivf_centroids = centroids;
        self.ivf_drift = drift;

        // Rebuild by_key from persisted assignments (only for nodes still in ivf_raw).
        for (&node, &c) in &clusters {
            if !self.ivf_raw.contains_key(&node) {
                // Node was removed post-snapshot (WAL replay deleted it).  Skip.
                continue;
            }
            self.ivf_clusters.insert(node, c);
            self.by_key
                .entry(ivf_cluster_key(c))
                .or_default()
                .insert(node);
        }
    }

    // -----------------------------------------------------------------------
    // HNSW methods
    // -----------------------------------------------------------------------

    /// Initialise the HNSW graph for this side, seeding it with `FNV-1a(rule_name)`.
    ///
    /// Must be called before inserting nodes via `CandidateSpec::Hnsw`.
    /// Idempotent: calling again with the same name replaces the existing graph.
    pub fn init_hnsw(&mut self, rule_name: &str) {
        let seed = fnv1a_u64(rule_name.as_bytes());
        self.hnsw = Some(HnswIndex::new(seed));
    }

    /// HNSW candidate lookup: `k`-nearest-neighbor search using the built graph.
    ///
    /// Falls back to returning all tracked nodes when the HNSW is absent or
    /// empty (e.g. before any insert or when used without `init_hnsw`).
    fn hnsw_candidates(
        &self,
        field: &str,
        k: usize,
        get: &dyn Fn(&str) -> Option<Value>,
    ) -> BTreeSet<u32> {
        let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
            return BTreeSet::new();
        };
        if let Some(h) = &self.hnsw {
            if !h.is_empty() {
                return h.search(&xs, k).into_iter().map(|(id, _)| id).collect();
            }
        }
        // Fallback: full scan of all tracked nodes (superset of true positives).
        self.hnsw_tracked.clone()
    }

    /// Export the HNSW graph as an opaque bincoded blob.
    ///
    /// Returns an empty `Vec` when the HNSW is not initialized.
    pub fn export_hnsw_blob(&self) -> Vec<u8> {
        self.hnsw
            .as_ref()
            .and_then(|h| bincode::serialize(h).ok())
            .unwrap_or_default()
    }

    /// Restore the HNSW graph from a previously exported blob.
    ///
    /// The `hnsw_tracked` set is populated from the restored graph's node ids
    /// so candidates/remove work correctly after restore.
    /// Silently ignores empty or corrupt blobs (HNSW stays uninitialized).
    pub fn load_hnsw_blob(&mut self, blob: &[u8]) {
        if blob.is_empty() {
            return;
        }
        if let Ok(h) = bincode::deserialize::<HnswIndex>(blob) {
            // Repopulate hnsw_tracked from the loaded graph.
            self.hnsw_tracked = h.node_ids();
            self.hnsw = Some(h);
        }
    }

    /// True when the HNSW graph has been initialized and contains at least one node.
    pub fn has_hnsw(&self) -> bool {
        self.hnsw.as_ref().is_some_and(|h| !h.is_empty())
    }

    /// Borrow the HNSW index, if initialized.
    pub fn hnsw_ref(&self) -> Option<&HnswIndex> {
        self.hnsw.as_ref()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::def::Predicate;
    use core_storage::Value;
    use std::collections::{BTreeMap, HashMap};

    fn getter(map: &HashMap<String, Value>) -> impl Fn(&str) -> Option<Value> + '_ {
        move |f: &str| map.get(f).cloned()
    }

    #[test]
    fn kmeans_centroids_are_unit_norm() {
        let vecs = vec![(0, vec![3.0, 0.0, 0.0]), (1, vec![0.0, 4.0, 0.0])];
        let cents = kmeans_fit(&vecs, 2, 1);
        for c in cents {
            let n = c.iter().map(|x| x * x).sum::<f64>().sqrt();
            assert!((n - 1.0).abs() < 1e-9, "{n}");
        }
    }

    /// Raw L2 would put `[3,0,0]` on a nearby large centroid while cosine (and
    /// the unit vector `[1,0,0]`) prefer the x-axis centroid. Assignment must
    /// L2-normalize first so scale-equivalent vectors share a cluster.
    ///
    /// Tests IVF directly (via `CandidateSpec::VectorClusters`) since
    /// `candidate_spec_approx` now returns `CandidateSpec::Hnsw`.
    #[test]
    fn scaled_vector_joins_same_ivf_cluster_as_unit() {
        // Use VectorClusters directly to test IVF cluster assignment.
        let spec = CandidateSpec::VectorClusters {
            field: "emb",
            min: 0.5,
        };
        let mut idx = SideIndex::default();
        idx.load_ivf_state(
            vec![vec![1.0, 0.0, 0.0], vec![2.5, 0.1, 0.0]],
            BTreeMap::new(),
            0,
        );
        idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0, 0.0])));
        idx.insert(&spec, 2, &getter(&emb(&[3.0, 0.0, 0.0])));
        assert_eq!(
            idx.ivf_cluster_of(1),
            idx.ivf_cluster_of(2),
            "scale-equivalent vectors must share an IVF cluster; got {:?} vs {:?}",
            idx.ivf_cluster_of(1),
            idx.ivf_cluster_of(2)
        );
        assert_eq!(idx.ivf_cluster_of(1), Some(0));
    }

    #[test]
    fn scalar_index_buckets_by_value() {
        let pred = Predicate::FieldEqual {
            field: "ind".into(),
        };
        let spec = candidate_spec(&pred);
        let mut idx = SideIndex::default();
        let a: HashMap<_, _> = [("ind".to_string(), Value::Str("arch".into()))].into();
        let b: HashMap<_, _> = [("ind".to_string(), Value::Str("law".into()))].into();
        idx.insert(&spec, 1, &getter(&a));
        idx.insert(&spec, 2, &getter(&b));
        idx.insert(&spec, 3, &getter(&a));
        let c = idx.candidates(&spec, &getter(&a));
        assert_eq!(c.into_iter().collect::<Vec<_>>(), vec![1, 3]);
        idx.remove(&spec, 3, &getter(&a));
        assert_eq!(idx.candidates(&spec, &getter(&a)).len(), 1);
        // node without the field indexes nothing and matches nothing
        let empty: HashMap<String, Value> = HashMap::new();
        idx.insert(&spec, 9, &getter(&empty));
        assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
    }

    #[test]
    fn token_index_unions_buckets() {
        let mk =
            |items: &[&str]| Value::List(items.iter().map(|s| Value::Str((*s).into())).collect());
        let pred = Predicate::Overlap {
            field: "tags".into(),
            min: 0.5,
        };
        let spec = candidate_spec(&pred);
        let mut idx = SideIndex::default();
        let a: HashMap<_, _> = [("tags".to_string(), mk(&["x", "y"]))].into();
        let b: HashMap<_, _> = [("tags".to_string(), mk(&["y", "z"]))].into();
        let c: HashMap<_, _> = [("tags".to_string(), mk(&["q"]))].into();
        idx.insert(&spec, 1, &getter(&a));
        idx.insert(&spec, 2, &getter(&b));
        idx.insert(&spec, 3, &getter(&c));
        let probe: HashMap<_, _> = [("tags".to_string(), mk(&["y"]))].into();
        assert_eq!(
            idx.candidates(&spec, &getter(&probe))
                .into_iter()
                .collect::<Vec<_>>(),
            vec![1, 2]
        );
        idx.remove(&spec, 2, &getter(&b));
        assert_eq!(
            idx.candidates(&spec, &getter(&probe))
                .into_iter()
                .collect::<Vec<_>>(),
            vec![1]
        );
    }

    #[test]
    fn all_intersects_parts_and_bykey_indexes_nothing() {
        let all = Predicate::All(vec![
            Predicate::FieldEqual {
                field: "ind".into(),
            },
            Predicate::Overlap {
                field: "tags".into(),
                min: 0.5,
            },
        ]);
        match candidate_spec(&all) {
            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
            other => panic!("{other:?}"),
        }
        let km = Predicate::KeyMatch { field: "fk".into() };
        assert!(matches!(candidate_spec(&km), CandidateSpec::ByKey));
        let mut idx = SideIndex::default();
        let a: HashMap<_, _> = [("fk".to_string(), Value::Str("c1".into()))].into();
        idx.insert(&candidate_spec(&km), 1, &getter(&a));
        assert!(idx.candidates(&candidate_spec(&km), &getter(&a)).is_empty());
    }

    fn year(v: Value) -> HashMap<String, Value> {
        [("year".to_string(), v)].into()
    }

    fn loc(lat: f64, lon: f64) -> HashMap<String, Value> {
        [(
            "loc".to_string(),
            Value::List(vec![Value::Float(lat), Value::Float(lon)]),
        )]
        .into()
    }

    fn emb(vals: &[f64]) -> HashMap<String, Value> {
        [(
            "emb".to_string(),
            Value::List(vals.iter().copied().map(Value::Float).collect()),
        )]
        .into()
    }

    fn bucket_int(spec: &CandidateSpec, map: &HashMap<String, Value>) -> Option<i64> {
        match SideIndex::index_keys(spec, &getter(map)).into_iter().next() {
            Some(ValueKey::Int(b)) => Some(b),
            _ => None,
        }
    }

    #[test]
    fn numeric_bucket_adjacency_and_far_value() {
        let pred = Predicate::NumericWithin {
            field: "year".into(),
            tolerance: 2.0,
        };
        let spec = candidate_spec(&pred);
        assert!(matches!(
            spec,
            CandidateSpec::NumericBucket {
                field: "year",
                tolerance
            } if tolerance == 2.0
        ));

        let v10 = year(Value::Float(10.0));
        let v119 = year(Value::Float(11.9));
        let v99 = year(Value::Float(9.9));
        let v141 = year(Value::Float(14.1));

        let b10 = bucket_int(&spec, &v10).unwrap();
        let b119 = bucket_int(&spec, &v119).unwrap();
        let b99 = bucket_int(&spec, &v99).unwrap();
        // 10.0 and 11.9 share a bucket; 9.9 is adjacent (forces ±1 probe).
        assert!((b10 - b119).abs() <= 1);
        assert!((b10 - b99).abs() <= 1);

        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&v10));
        idx.insert(&spec, 2, &getter(&v119));
        idx.insert(&spec, 3, &getter(&v141));
        idx.insert(&spec, 4, &getter(&v99));
        let hits = idx.candidates(&spec, &getter(&v10));
        assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2, 4]);
    }

    #[test]
    fn numeric_tol_zero_int_float_collide() {
        let pred = Predicate::NumericWithin {
            field: "year".into(),
            tolerance: 0.0,
        };
        let spec = candidate_spec(&pred);
        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&year(Value::Int(2))));
        assert_eq!(
            idx.candidates(&spec, &getter(&year(Value::Float(2.0))))
                .into_iter()
                .collect::<Vec<_>>(),
            vec![1]
        );
        assert!(idx
            .candidates(&spec, &getter(&year(Value::Float(2.1))))
            .is_empty());
    }

    #[test]
    fn numeric_tol_zero_signed_zero_collides() {
        let pred = Predicate::NumericWithin {
            field: "year".into(),
            tolerance: 0.0,
        };
        let spec = candidate_spec(&pred);
        let neg = year(Value::Float(-0.0));
        let pos = year(Value::Float(0.0));
        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&neg));
        assert_eq!(
            idx.candidates(&spec, &getter(&pos))
                .into_iter()
                .collect::<Vec<_>>(),
            vec![1]
        );
        let mut idx2 = SideIndex::default();
        idx2.insert(&spec, 2, &getter(&pos));
        assert_eq!(
            idx2.candidates(&spec, &getter(&neg))
                .into_iter()
                .collect::<Vec<_>>(),
            vec![2]
        );
    }

    #[test]
    fn geo_grid_same_cell_cross_cell_and_far_city() {
        let pred = Predicate::GeoRadius {
            field: "loc".into(),
            km: 400.0,
        };
        let spec = candidate_spec(&pred);
        assert!(matches!(
            spec,
            CandidateSpec::GeoGrid {
                field: "loc",
                km
            } if km == 400.0
        ));

        let paris = loc(48.8566, 2.3522);
        let london = loc(51.5074, -0.1278);
        let nearby = loc(48.9, 2.4); // same cell as Paris at km=400
        let ny = loc(40.7128, -74.0060);

        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&paris));
        idx.insert(&spec, 2, &getter(&london));
        idx.insert(&spec, 3, &getter(&nearby));
        idx.insert(&spec, 4, &getter(&ny));

        let from_paris = idx.candidates(&spec, &getter(&paris));
        assert!(from_paris.contains(&1), "same-cell self");
        assert!(from_paris.contains(&3), "same-cell neighbor");
        assert!(from_paris.contains(&2), "cross-cell Paris↔London ~343.5 km");
        assert!(!from_paris.contains(&4), "New York not in 400 km probe");
    }

    #[test]
    fn geo_grid_high_latitude_probe_is_superset() {
        let pred = Predicate::GeoRadius {
            field: "loc".into(),
            km: 340.0,
        };
        let spec = candidate_spec(&pred);
        let reyk = loc(64.1466, -21.9426);
        let lat = 64.0_f64;
        let dlon = 300.0 / (111.0 * lat.to_radians().cos());
        let east = loc(lat, -21.9426 + dlon);

        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&reyk));
        idx.insert(&spec, 2, &getter(&east));
        let hits = idx.candidates(&spec, &getter(&reyk));
        assert!(
            hits.contains(&2),
            "300 km east of Reykjavik must stay in the high-lat probe"
        );
    }

    #[test]
    fn geo_grid_antimeridian_wrap_and_evaluate_agree() {
        let pred = Predicate::GeoRadius {
            field: "loc".into(),
            km: 400.0,
        };
        let spec = candidate_spec(&pred);
        let east = loc(70.0, 179.9);
        let west = loc(70.0, -179.9);

        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&east));
        assert!(
            idx.candidates(&spec, &getter(&west)).contains(&1),
            "±180 pair at lat 70 must land in the wrapped probe"
        );

        let sp = |f: &str| east.get(f).cloned();
        let dp = |f: &str| west.get(f).cloned();
        let score = crate::def::evaluate(
            &pred,
            &crate::def::NodeView {
                key: "e",
                props: &sp,
            },
            &crate::def::NodeView {
                key: "w",
                props: &dp,
            },
        );
        assert!(
            score.is_some(),
            "haversine must match across the antimeridian"
        );

        // Wrap must not alias distant longitudes into the Paris probe.
        let paris = loc(48.8566, 2.3522);
        let ny = loc(40.7128, -74.0060);
        let mut idx2 = SideIndex::default();
        idx2.insert(&spec, 4, &getter(&ny));
        assert!(
            !idx2.candidates(&spec, &getter(&paris)).contains(&4),
            "New York still not in the Paris probe after wrap"
        );
    }

    #[test]
    fn scan_all_returns_vector_nodes_skips_malformed() {
        let pred = Predicate::VectorSimilar {
            field: "emb".into(),
            min: 0.5,
        };
        let spec = candidate_spec(&pred);
        assert!(matches!(spec, CandidateSpec::ScanAll { field: "emb" }));

        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
        idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
        idx.insert(&spec, 3, &getter(&emb(&[1.0, 2.0, 3.0])));
        let empty: HashMap<_, _> = [("emb".to_string(), Value::List(vec![]))].into();
        let text: HashMap<_, _> =
            [("emb".to_string(), Value::List(vec![Value::Str("x".into())]))].into();
        let missing: HashMap<String, Value> = HashMap::new();
        idx.insert(&spec, 4, &getter(&empty));
        idx.insert(&spec, 5, &getter(&text));
        idx.insert(&spec, 6, &getter(&missing));

        let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
        assert_eq!(
            hits.into_iter().collect::<Vec<_>>(),
            vec![1, 2],
            "dim-2 probe must drop the dim-3 member"
        );
        assert_eq!(
            idx.candidates(&spec, &getter(&emb(&[1.0, 2.0, 3.0])))
                .into_iter()
                .collect::<Vec<_>>(),
            vec![3]
        );
        with_vector_dim_reject(false, || {
            assert_eq!(
                idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])))
                    .into_iter()
                    .collect::<Vec<_>>(),
                vec![1, 2, 3],
                "unfiltered ScanAll still returns every vector node"
            );
        });
        assert_eq!(idx.vec_dim(1), Some(2));
        assert_eq!(idx.vec_dim(3), Some(3));
        assert!(idx.vec_meta(1).is_some());
        assert!(idx.vec_dim(4).is_none());
        assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
        assert!(idx.candidates(&spec, &getter(&text)).is_empty());
        assert!(idx.candidates(&spec, &getter(&missing)).is_empty());
        idx.remove(&spec, 1, &getter(&emb(&[1.0, 0.0])));
        assert!(idx.vec_dim(1).is_none());
    }

    #[test]
    fn legacy_specs_probe_keys_equal_index_keys() {
        let a: HashMap<_, _> = [
            ("ind".to_string(), Value::Str("arch".into())),
            (
                "tags".to_string(),
                Value::List(vec![Value::Str("x".into()), Value::Str("y".into())]),
            ),
            ("fk".to_string(), Value::Str("c1".into())),
        ]
        .into();
        let get = getter(&a);
        for pred in [
            Predicate::KeyMatch { field: "fk".into() },
            Predicate::FieldEqual {
                field: "ind".into(),
            },
            Predicate::Overlap {
                field: "tags".into(),
                min: 0.5,
            },
        ] {
            let spec = candidate_spec(&pred);
            assert_eq!(
                SideIndex::index_keys(&spec, &get),
                SideIndex::probe_keys(&spec, &get)
            );
        }
    }

    #[test]
    fn all_vector_then_field_equal_does_not_scan_all() {
        let p = Predicate::All(vec![
            Predicate::VectorSimilar {
                field: "e".into(),
                min: 0.8,
            },
            Predicate::FieldEqual {
                field: "industry".into(),
            },
        ]);
        match candidate_spec(&p) {
            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
            other => panic!("{other:?}"),
        }

        let spec = candidate_spec(&p);
        let mut idx = SideIndex::default();
        let mk = |industry: &str, e: &[f64]| {
            [
                ("industry".to_string(), Value::Str(industry.into())),
                (
                    "e".to_string(),
                    Value::List(e.iter().copied().map(Value::Float).collect()),
                ),
            ]
            .into()
        };
        let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
        let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
        let no_vec: HashMap<_, _> = [("industry".to_string(), Value::Str("tech".into()))].into();
        idx.insert(&spec, 1, &getter(&same));
        idx.insert(&spec, 2, &getter(&other_ind));
        idx.insert(&spec, 3, &getter(&no_vec));

        let hits = idx.candidates(&spec, &getter(&same));
        assert!(hits.contains(&1), "matching industry must stay a candidate");
        assert!(
            !hits.contains(&2),
            "different industry must not be scanned in via VectorSimilar"
        );
        assert!(
            hits.contains(&3),
            "ScanAll is universe: extra Scalar-only candidates are allowed"
        );

        let empty_ind: HashMap<_, _> = mk("finance", &[1.0, 0.0]);
        assert!(
            idx.candidates(&spec, &getter(&empty_ind)).is_empty(),
            "empty Scalar child → empty intersect"
        );
    }

    #[test]
    fn all_approx_vector_then_field_equal_is_intersect() {
        let p = Predicate::All(vec![
            Predicate::VectorSimilar {
                field: "e".into(),
                min: 0.8,
            },
            Predicate::FieldEqual {
                field: "industry".into(),
            },
        ]);
        match candidate_spec_approx(&p) {
            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
            other => panic!("{other:?}"),
        }

        let spec = candidate_spec_approx(&p);
        let mut idx = SideIndex::default();
        // Initialize HNSW so insertions populate the graph.
        idx.init_hnsw("test-rule");
        let mk = |industry: &str, e: &[f64]| {
            [
                ("industry".to_string(), Value::Str(industry.into())),
                (
                    "e".to_string(),
                    Value::List(e.iter().copied().map(Value::Float).collect()),
                ),
            ]
            .into()
        };
        let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
        let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
        idx.insert(&spec, 1, &getter(&same));
        idx.insert(&spec, 2, &getter(&other_ind));
        let hits = idx.candidates(&spec, &getter(&same));
        assert!(hits.contains(&1), "matching industry must stay a candidate");
        assert!(
            !hits.contains(&2),
            "FieldEqual must be probed on the approximate All path"
        );
    }

    #[test]
    fn all_of_scan_all_stays_scan_all() {
        let p = Predicate::All(vec![
            Predicate::VectorSimilar {
                field: "emb".into(),
                min: 0.5,
            },
            Predicate::VectorSimilar {
                field: "emb".into(),
                min: 0.9,
            },
        ]);
        match candidate_spec(&p) {
            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
            other => panic!("{other:?}"),
        }
        let spec = candidate_spec(&p);
        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
        idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
        let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
        assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2]);
    }

    #[test]
    fn any_stays_union() {
        let p = Predicate::Any(vec![
            Predicate::FieldEqual {
                field: "industry".into(),
            },
            Predicate::Overlap {
                field: "tags".into(),
                min: 0.5,
            },
        ]);
        match candidate_spec(&p) {
            CandidateSpec::Union(v) => assert_eq!(v.len(), 2),
            other => panic!("{other:?}"),
        }
    }

    /// Checkpoints are populated at insert, torn out at remove,
    /// and ckpts[0] must equal the full L2 norm.
    #[test]
    fn checkpoint_populated_and_consistent_with_norm() {
        let pred = Predicate::VectorSimilar {
            field: "emb".into(),
            min: 0.8,
        };
        let spec = candidate_spec(&pred);
        let xs = [3.0f64, 4.0]; // norm = 5.0
        let mut idx = SideIndex::default();
        idx.insert(&spec, 1, &getter(&emb(&xs)));

        let ckpts = idx
            .vec_ckpts(1)
            .expect("checkpoints must exist after insert");
        let (_, norm) = idx.vec_meta(1).unwrap();
        assert!(
            (ckpts[0] - norm).abs() < 1e-12,
            "ckpts[0] must equal the full L2 norm; got {} vs {}",
            ckpts[0],
            norm
        );
        assert!(
            (norm - 5.0).abs() < 1e-12,
            "norm of [3,4] must be 5.0, got {norm}"
        );

        // Remove must tear out checkpoints.
        idx.remove(&spec, 1, &getter(&emb(&xs)));
        assert!(
            idx.vec_ckpts(1).is_none(),
            "checkpoints must be removed after remove()"
        );
    }

    /// fresh_ckpts_for returns None when the live vector's norm differs
    /// (freshness gate) and Some when it matches.
    #[test]
    fn fresh_ckpts_for_freshness_gate() {
        let pred = Predicate::VectorSimilar {
            field: "emb".into(),
            min: 0.8,
        };
        let spec = candidate_spec(&pred);
        let xs = [1.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
        let mut idx = SideIndex::default();
        idx.insert(&spec, 7, &getter(&emb(&xs)));

        // Correct live vector → gate passes.
        let result = idx.fresh_ckpts_for(7, &xs);
        assert!(
            result.is_some(),
            "fresh_ckpts_for must succeed with matching live vector"
        );
        let (norm, ckpts) = result.unwrap();
        assert!((norm - 1.0).abs() < 1e-12);
        assert!((ckpts[0] - 1.0).abs() < 1e-12);

        // Wrong norm → gate rejects.
        let wrong = [2.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; // norm = 2.0
        assert!(
            idx.fresh_ckpts_for(7, &wrong).is_none(),
            "freshness gate must reject mismatched norm"
        );

        // Wrong dim → gate rejects.
        let short = [1.0f64, 0.0];
        assert!(
            idx.fresh_ckpts_for(7, &short).is_none(),
            "freshness gate must reject mismatched dim"
        );

        // Missing node → returns None.
        assert!(idx.fresh_ckpts_for(99, &xs).is_none());
    }

    /// Checkpoints for a dim-16 vector: ckpts[i] must be non-increasing
    /// (suffix norms decrease as the suffix shrinks).
    #[test]
    fn checkpoint_suffix_norms_non_increasing() {
        let pred = Predicate::VectorSimilar {
            field: "emb".into(),
            min: 0.5,
        };
        let spec = candidate_spec(&pred);
        let xs: Vec<f64> = (1..=16).map(|i| i as f64).collect();
        let mut idx = SideIndex::default();
        idx.insert(&spec, 42, &getter(&emb(&xs)));

        let ckpts = *idx.vec_ckpts(42).unwrap();
        for c in 0..7 {
            assert!(
                ckpts[c] >= ckpts[c + 1] - 1e-12,
                "suffix norm must be non-increasing: ckpts[{c}]={} < ckpts[{}]={}",
                ckpts[c],
                c + 1,
                ckpts[c + 1]
            );
        }
        // ckpts[7] = suffix norm of the last 2 elements (14..=16).
        let expected_last = (15.0f64 * 15.0 + 16.0 * 16.0).sqrt();
        assert!(
            (ckpts[7] - expected_last).abs() < 1e-9,
            "ckpts[7] should be norm of last segment; got {} vs {}",
            ckpts[7],
            expected_last
        );
    }
}