hypersteeldb 0.5.3

//! Gazetteer tier — deterministic HIGH-RESOLUTION token extraction. The SPLADE heads fire on single
//! wordpieces, so multi-word entities ("Amazon EKS", "Direct Connect Gateway", "zero-ETL") shatter into
//! coarse fragments. This tier matches a closed vocabulary of surface→token entries (mined from the
//! corpus + discovered facet examples) and emits the WHOLE-ENTITY `facet/value` token, so the bitmap
//! carries discriminative tokens the analytics programs (rank/crosstab/structure) can reason over.
//! Used symmetrically by ingest AND query so index and question share one high-res token space.
//! Ported from `src/ingest/gazetteer.ts` (`71e1714`).
//!
//! Match strategy is dependency-free (no regex crate): text is tokenised into maximal alphanumeric
//! words with byte spans, then a longest-first sliding window (up to the gazetteer's max phrase length)
//! is looked up against a normalised surface→token map. Both surface keys and text windows are
//! normalised the same way (lowercase words joined by a single space), so "zero-ETL" ↔ "zero etl".

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;

#[derive(Deserialize, Serialize)]
struct GazEntry {
    surface: String,
    token: String,
}

pub struct GazHit {
    pub token: String,
    pub span: (usize, usize),
}

pub struct Gazetteer {
    /// normalised surface (lowercase words joined by single space) → token
    map: HashMap<String, String>,
    /// longest surface, in words — the window size to try first
    max_words: usize,
    /// Subset of `map` learned at ingest (the growing overlay). Persisted separately from the shipped
    /// base so growth accumulates across runs without rewriting the base vocabulary.
    learned: HashMap<String, String>,
}

impl Gazetteer {
    /// Load a `gazetteer.json` (array of `{surface, token}`). Longest surfaces win on overlap.
    pub fn load(path: &Path) -> Result<Gazetteer, Box<dyn std::error::Error + Send + Sync>> {
        let entries: Vec<GazEntry> = serde_json::from_slice(&std::fs::read(path)?)?;
        let mut g = Gazetteer::empty();
        for e in entries {
            g.insert(&e.surface, e.token, false);
        }
        Ok(g)
    }

    /// An empty gazetteer — the starting point for a pure learned overlay (no shipped base present).
    pub fn empty() -> Gazetteer {
        Gazetteer { map: HashMap::new(), max_words: 1, learned: HashMap::new() }
    }

    fn insert(&mut self, surface: &str, token: String, learned: bool) -> bool {
        let key = normalize(surface);
        if key.is_empty() || token.is_empty() {
            return false;
        }
        self.max_words = self.max_words.max(key.split(' ').count());
        let is_new = !self.map.contains_key(&key);
        self.map.entry(key.clone()).or_insert_with(|| token.clone());
        if learned {
            self.learned.entry(key).or_insert(token);
        }
        is_new
    }

    /// Learn a whole-entity surface→token from ingest. Returns true if it was new. The overlay is what
    /// `save_overlay` persists; base entries are untouched.
    pub fn learn(&mut self, surface: &str, token: impl Into<String>) -> bool {
        self.insert(surface, token.into(), true)
    }

    /// Register an entity span under a **canonical token** derived from its content-word core, and learn
    /// `surface → canonical`. This is conservative entity normalization: leading determiners/possessives
    /// ("the/our/a") and a small set of trailing event/status words ("outage/failure/…") are stripped, so
    /// case/determiner/role variants of one entity map to the SAME canonical `facet/core` token (their
    /// postings merge), while distinct cores never merge. Returns the canonical token to emit for this
    /// span, or `None` if the surface has no content core. Ported intent from roadmap #7.
    pub fn register(&mut self, surface: &str, facet: &str) -> Option<String> {
        let core = canonical_core(surface)?;
        let token = format!("{facet}/{core}");
        self.insert(surface, token.clone(), true);
        Some(token)
    }

    /// Merge a previously-saved overlay file into this gazetteer (its entries count as learned so a
    /// later `save_overlay` round-trips them).
    pub fn merge_overlay(&mut self, path: &Path) {
        if let Ok(bytes) = std::fs::read(path) {
            if let Ok(entries) = serde_json::from_slice::<Vec<GazEntry>>(&bytes) {
                for e in entries {
                    self.insert(&e.surface, e.token, true);
                }
            }
        }
    }

    /// Persist the learned overlay (union with whatever is already on disk) as a `gazetteer.json` array.
    pub fn save_overlay(&self, path: &Path) -> std::io::Result<()> {
        // Union with any existing overlay so concurrent/earlier growth isn't lost.
        let mut merged: HashMap<String, String> = HashMap::new();
        if let Ok(bytes) = std::fs::read(path) {
            if let Ok(prev) = serde_json::from_slice::<Vec<GazEntry>>(&bytes) {
                for e in prev {
                    merged.insert(normalize(&e.surface), e.token);
                }
            }
        }
        for (k, v) in &self.learned {
            merged.insert(k.clone(), v.clone());
        }
        let mut entries: Vec<GazEntry> = merged.into_iter().map(|(surface, token)| GazEntry { surface, token }).collect();
        entries.sort_by(|a, b| a.surface.cmp(&b.surface));
        let json = serde_json::to_vec_pretty(&entries).map_err(std::io::Error::other)?;
        std::fs::write(path, json)
    }

    /// Number of learned (overlay) entries.
    pub fn learned_len(&self) -> usize {
        self.learned.len()
    }

    pub fn len(&self) -> usize {
        self.map.len()
    }
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Non-overlapping high-res token hits (leftmost, longest-first) over `text`, with char/byte spans.
    pub fn extract(&self, text: &str) -> Vec<GazHit> {
        let words = words_with_spans(text);
        let mut out = Vec::new();
        let mut i = 0;
        while i < words.len() {
            let mut matched = false;
            let hi = (self.max_words).min(words.len() - i);
            for n in (1..=hi).rev() {
                let phrase: Vec<&str> = words[i..i + n].iter().map(|(w, _, _)| w.as_str()).collect();
                let key = phrase.join(" ");
                if let Some(tok) = self.map.get(&key) {
                    let span = (words[i].1, words[i + n - 1].2);
                    out.push(GazHit { token: tok.clone(), span });
                    i += n;
                    matched = true;
                    break;
                }
            }
            if !matched {
                i += 1;
            }
        }
        out
    }
}

/// Normalise a surface string to lowercase alphanumeric words joined by single spaces.
fn normalize(s: &str) -> String {
    words_with_spans(s).into_iter().map(|(w, _, _)| w).collect::<Vec<_>>().join(" ")
}

/// Leading words that are never part of an entity's identity (determiners/possessives). Stripped so
/// "the redshift cluster" ↔ "redshift cluster".
const LEADING_DROP: &[&str] =
    &["the", "a", "an", "our", "your", "its", "their", "this", "that", "these", "those", "my", "his", "her"];

/// Trailing event/status words that describe something happening TO an entity, not the entity itself.
/// Conservatively small so we never strip identity-bearing words. "connect gateway outage" ↔ "connect
/// gateway".
const TRAILING_DROP: &[&str] =
    &["outage", "failure", "error", "issue", "incident", "problem", "alert", "warning", "downtime"];

/// Content-word core of a surface for canonicalization: normalise, drop leading determiners/possessives
/// and trailing event/status words, join with '-'. Returns `None` if nothing identity-bearing remains.
/// Deterministic, so any variant with the same core produces the same canonical slug.
fn canonical_core(surface: &str) -> Option<String> {
    let mut words: Vec<String> = words_with_spans(surface).into_iter().map(|(w, _, _)| w).collect();
    while words.first().map(|w| LEADING_DROP.contains(&w.as_str())).unwrap_or(false) {
        words.remove(0);
    }
    while words.len() > 1 && words.last().map(|w| TRAILING_DROP.contains(&w.as_str())).unwrap_or(false) {
        words.pop();
    }
    if words.is_empty() {
        return None;
    }
    Some(words.join("-"))
}

/// Maximal alphanumeric runs, lowercased, with their byte spans in the source.
fn words_with_spans(text: &str) -> Vec<(String, usize, usize)> {
    let mut out = Vec::new();
    let mut start: Option<usize> = None;
    let mut buf = String::new();
    for (i, ch) in text.char_indices() {
        if ch.is_alphanumeric() {
            if start.is_none() {
                start = Some(i);
            }
            buf.extend(ch.to_lowercase());
        } else if let Some(s) = start.take() {
            out.push((std::mem::take(&mut buf), s, i));
        }
    }
    if let Some(s) = start.take() {
        out.push((buf, s, text.len()));
    }
    out
}

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

    #[test]
    fn learn_extract_and_overlay_roundtrip() {
        let dir = std::env::temp_dir().join(format!("steeldb-gaz-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("overlay.json");
        let _ = std::fs::remove_file(&path);

        // Learn two multi-word entities, extract, and persist.
        let mut g = Gazetteer::empty();
        assert!(g.learn("Model X", "ent/model-x"));
        assert!(g.learn("Direct Connect Gateway", "aws-service/direct-connect-gateway"));
        assert!(!g.learn("model  x", "ent/model-x")); // same normalised surface → not new
        assert_eq!(g.learned_len(), 2);
        let hits = g.extract("the Model X uses a Direct Connect Gateway");
        let toks: Vec<&str> = hits.iter().map(|h| h.token.as_str()).collect();
        assert!(toks.contains(&"ent/model-x"));
        assert!(toks.contains(&"aws-service/direct-connect-gateway"));
        g.save_overlay(&path).unwrap();

        // A fresh gazetteer that merges the overlay recognises the same entities (query-side growth).
        let mut g2 = Gazetteer::empty();
        g2.merge_overlay(&path);
        let toks2: Vec<String> = g2.extract("Model X").into_iter().map(|h| h.token).collect();
        assert_eq!(toks2, vec!["ent/model-x".to_string()]);

        // Saving again unions rather than duplicates.
        g2.save_overlay(&path).unwrap();
        let entries: Vec<GazEntry> = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
        assert_eq!(entries.len(), 2);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn register_collapses_alias_variants_to_canonical() {
        let mut g = Gazetteer::empty();
        // Determiner, case, and trailing event-word variants of ONE entity.
        let a = g.register("The Payment Gateway", "ent").unwrap();
        let b = g.register("payment gateway", "ent").unwrap();
        let c = g.register("Our Payment Gateway", "ent").unwrap();
        let d = g.register("Payment Gateway outage", "ent").unwrap();
        assert_eq!(a, "ent/payment-gateway");
        assert_eq!(b, a);
        assert_eq!(c, a);
        assert_eq!(d, a, "trailing event word should be stripped");

        // A genuinely different core must NOT collapse into the same token (no over-merge).
        let other = g.register("Payment Processor", "ent").unwrap();
        assert_ne!(other, a);

        // Query-side: any variant surface extracts the canonical token.
        let toks: Vec<String> = g.extract("we upgraded our payment gateway last week").into_iter().map(|h| h.token).collect();
        assert!(toks.contains(&"ent/payment-gateway".to_string()));
    }
}