yantrikdb 0.16.0

//! Graph traversal utilities for entity-augmented recall.

use std::collections::{HashMap, HashSet, VecDeque};

use rusqlite::{params, Connection};

use crate::error::Result;

// ── Word-boundary entity matching ──

/// Tokenize text into lowercase words, splitting on non-alphanumeric
/// chars — INCLUDING the apostrophe (wheel C5a, 2026-08-06).
///
/// The apostrophe was exempted to keep contractions whole. Measured in
/// production, the exemption did the opposite of its job, twice over:
/// `Taylor's` survived as one token that never matched entity `taylor`
/// (a possessive silently disabled entity resolution on the default
/// path — one apostrophe changed ~77% of top-5 on true minimal pairs),
/// and contractions became first-class phantom entities (`Don't` at 96
/// mentions; `Pranab's` held 748 mentions — 35% of that person's
/// references — mistyped and unreachable). Splitting symmetrically is
/// safe because entity names pass through this SAME tokenizer:
/// `O'Brien` becomes `[o, brien]` on both the entity and query side,
/// so contiguous multi-token matching still holds. The persisted
/// pollution needs the C5b alias migration; this stops new pollution.
pub fn tokenize(text: &str) -> Vec<String> {
    text.split(|c: char| !c.is_alphanumeric())
        .filter(|s| !s.is_empty())
        .map(|s| s.to_lowercase())
        .collect()
}

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

    #[test]
    fn possessive_resolves_to_the_bare_entity() {
        // The hermes minimal pair: "Taylor's" must yield token "taylor".
        let toks = tokenize("What is Taylor's role?");
        assert!(toks.contains(&"taylor".to_string()), "{toks:?}");
        assert_eq!(
            tokenize("What is Taylor's role?")
                .iter()
                .filter(|t| *t == "taylor")
                .count(),
            tokenize("What is Taylor s role?")
                .iter()
                .filter(|t| *t == "taylor")
                .count(),
            "possessive and plain forms must tokenize alike"
        );
    }

    #[test]
    fn apostrophe_names_match_symmetrically() {
        // Entity and text pass the same tokenizer, so O'Brien matches
        // whether or not either side carries the apostrophe intact.
        let text_tokens = tokenize("A meeting with O'Brien about the launch");
        assert!(entity_matches_text("O'Brien", &text_tokens));
        assert!(entity_matches_text("o brien", &text_tokens));
    }

    #[test]
    fn contractions_stop_being_coherent_tokens() {
        // The exemption promoted Don't to a 96-mention phantom entity.
        assert_eq!(tokenize("Don't"), vec!["don", "t"]);
    }
}

/// Check if an entity name appears as whole-word(s) in pre-tokenized text.
/// Single-word entities require exact token match.
/// Multi-word entities require contiguous token sequence match.
pub fn entity_matches_text(entity: &str, text_tokens: &[String]) -> bool {
    let entity_tokens = tokenize(entity);
    if entity_tokens.is_empty() {
        return false;
    }
    if entity_tokens.len() == 1 {
        text_tokens.iter().any(|t| t == &entity_tokens[0])
    } else {
        text_tokens
            .windows(entity_tokens.len())
            .any(|window| window.iter().zip(entity_tokens.iter()).all(|(w, e)| w == e))
    }
}

// ── Heuristic proper-noun extraction ──

/// English function/pronoun/auxiliary words that should be stripped from the
/// start or end of a capitalized chunk. A sentence-initial "The" or "Our" is
/// capitalized by position, not because it names an entity.
const ENTITY_STOPWORDS: &[&str] = &[
    "The",
    "A",
    "An",
    "I",
    "We",
    "You",
    "He",
    "She",
    "It",
    "They",
    "This",
    "That",
    "These",
    "Those",
    "My",
    "Your",
    "His",
    "Her",
    "Its",
    "Our",
    "Their",
    "But",
    "And",
    "Or",
    "So",
    "If",
    "When",
    "Where",
    "What",
    "Who",
    "Why",
    "How",
    "Is",
    "Are",
    "Was",
    "Were",
    "Be",
    "Been",
    "Being",
    "Have",
    "Has",
    "Had",
    "Do",
    "Does",
    "Did",
    "Of",
    "In",
    "On",
    "At",
    "To",
    "For",
    "With",
    "From",
    "By",
    "As",
    "Than",
    "Then",
    "Also",
    "Just",
    "Only",
    "Very",
    "Much",
    // Added 2026-08-13 with the case fix below. These were absent in EVERY
    // case, so they became entities regardless of the comparison bug.
    "Not",
    "No",
    "Nor",
    "Most",
    "More",
    "Less",
    "Some",
    "Any",
    "All",
    "Each",
    "Every",
    "Both",
    "Such",
    "Same",
    "Other",
    "Another",
    "Yet",
    "Still",
    "Because",
    "While",
    "After",
    "Before",
    "During",
    "Since",
    "Until",
    "Between",
    "Through",
    "About",
    "Into",
    "Over",
    "Under",
    "Again",
    "Once",
    "Here",
    "There",
    "Now",
    "Thus",
    "However",
    "Therefore",
    "Note",
    "See",
    "Can",
    "Could",
    "Will",
    "Would",
    "Should",
    "May",
    "Might",
    "Must",
    "Let",
    "Get",
    "Got",
];

/// Bare month names. Not function words — a different class, and stoplisted
/// for a different reason: a month alone is not the thing a sentence is about,
/// but it appears in nearly every dated record, so as a graph node it links
/// everything to everything. Observed doing exactly that: a query for
/// "encryption at rest and key rotation" retrieved an unrelated record whose
/// stated join was `graph-connected via June`.
///
/// This is a blunt instrument. The principled fix for entities that are real
/// but uselessly common is inverse-document-frequency node weighting plus a
/// hub-degree penalty, so a node's retrieval weight falls as it connects more
/// of the corpus. Until that exists, a month is more noise than signal.
const AMBIGUOUS_COMMON_ENTITIES: &[&str] = &[
    "January",
    "February",
    "March",
    "April",
    "May",
    "June",
    "July",
    "August",
    "September",
    "October",
    "November",
    "December",
    "Monday",
    "Tuesday",
    "Wednesday",
    "Thursday",
    "Friday",
    "Saturday",
    "Sunday",
];

/// Is this token unusable as an entity?
///
/// **Case-INSENSITIVE, and that is the whole point.** This compared with
/// `ENTITY_STOPWORDS.contains(&tok)` — an exact string match against the
/// capitalized forms — so `"At"` was stripped while `"AT"` sailed through and
/// became an entity. `tokenize()` then lowercases it to `"at"`, and
/// `entity_matches_text` compares lowercased tokens, so the phantom entity
/// `AT` matched EVERY query containing the word "at".
///
/// Measured consequence on a ~900-record production store: the query
/// "encryption at rest and key rotation" returned a real-estate tax analysis
/// in the top 3, joined by `claims_match: AT -acquired-> 25 (anchor AT)`.
/// Anchors `NOT`, `THE`, `DID`, `Most` and `June` were seen the same way.
fn is_entity_stopword(tok: &str) -> bool {
    ENTITY_STOPWORDS.iter().any(|s| s.eq_ignore_ascii_case(tok))
        || AMBIGUOUS_COMMON_ENTITIES
            .iter()
            .any(|s| s.eq_ignore_ascii_case(tok))
}

/// A name is at most this many words. Beyond it, a "capitalized chunk" is a
/// run of prose, not an entity — the heuristic groups CONSECUTIVE capitalized
/// words with no upper bound, so a heading becomes one long entity.
const MAX_ENTITY_TOKENS: usize = 6;

/// Three or more ALL-CAPS words in a row is emphasis or a heading, not a name.
/// Genuine all-caps names are short: `NASA`, `IBM`, `HNSW`, `IBM WATSON`.
const MAX_ALLCAPS_TOKENS: usize = 2;

fn is_all_caps_token(tok: &str) -> bool {
    tok.chars().any(|c| c.is_alphabetic())
        && tok.chars().all(|c| !c.is_alphabetic() || c.is_uppercase())
}

/// Does this capitalized run read as prose rather than a name?
///
/// Found by censusing a live store after fixing the stoplist. The extractor
/// had no length bound, so entire ALL-CAPS sentences became single entities:
///
///   "THINGS I MISSED THAT CODEX FOUND BY READING THE CODE"
///   "USER MUST UPDATE MCP CONFIG"
///   "HERMES REMOTE DESKTOP LIVE VERIFICATION PASSED 2026 08 13"
///   "REAL ESTATE TAX ANALYSIS"
///
/// Every one of those is a node in the knowledge graph, and three came from
/// memories written that same day — an agent that writes ALL-CAPS headings
/// pollutes its own graph, which is a self-reinforcing failure a human author
/// would never trigger.
///
/// Two bounds, deliberately kept separate because they catch different shapes:
/// a token cap for runaway mixed-case runs, and an all-caps cap for headings.
/// Both err toward keeping short candidates, since a missed entity costs one
/// retrieval path while a phantom entity costs precision on EVERY query that
/// happens to share one of its words.
fn is_prose_run(chunk: &[String]) -> bool {
    if chunk.len() > MAX_ENTITY_TOKENS {
        return true;
    }
    chunk.iter().filter(|t| is_all_caps_token(t)).count() > MAX_ALLCAPS_TOKENS
}

/// Would today's extractor refuse to mint this entity name?
///
/// The rules above stop NEW pollution, but a store written by an older engine
/// still holds the phantoms — `AT` with 10 mentions, `REAL ESTATE TAX
/// ANALYSIS`, `USER MUST UPDATE MCP CONFIG` — and they keep degrading recall
/// until something removes them. [`crate::graph_index::GraphIndex`] applies
/// this at load, so a store heals by being opened rather than by running a
/// destructive migration: nothing is deleted, and reverting the rules restores
/// the old behaviour exactly.
///
/// Names a caller deliberately created via `relate()` are NEVER judged by this
/// — that check lives at the call site, which is the only place that knows
/// provenance.
pub fn is_rejected_entity_name(name: &str) -> bool {
    let toks: Vec<String> = name.split_whitespace().map(|s| s.to_string()).collect();
    if toks.is_empty() {
        return true;
    }
    // No alphabetic character anywhere: "546", "15", "3.6", "2026-08-16".
    // The capitalized-chunk extractor can never mint these, but the claims
    // lane can and did — a first-hand probe on the production store found
    // `claims_match: 15 -leads-> LOG (anchor 15)` SURVIVING the stopword
    // heal, because this predicate only knew function words. A subject
    // with no letters names nothing; it anchors neither claims nor
    // conflicts. (The conflict detector carried its own copy of this
    // guard; centralizing it here makes every surface that consults this
    // predicate — graph load, claims lane, conflict admission — agree.)
    if !name.chars().any(|c| c.is_alphabetic()) {
        return true;
    }
    // Wholly made of function words / bare months: "AT", "June", "THE Most".
    if toks.iter().all(|t| is_entity_stopword(t)) {
        return true;
    }
    is_prose_run(&toks)
}

/// Strip fenced code blocks and inline code spans before entity extraction.
///
/// The capitalized-chunk heuristic below cannot tell `String`, `User` or
/// `GET` in a code sample from `Alice`, `Anthropic` or `NASA` in prose — both
/// are capitalized or all-caps tokens. Measured on a code-bearing
/// conversation corpus (BEAM, 2026-08-11): ~360 records produced **5,550
/// entities**, roughly 15 per record, and every conflict the detector then
/// raised was a false positive keyed on the entity `GET` — pairing two
/// adjacent chunks of the same turn because both quoted a Flask route.
/// Garbage entities do not merely add noise: they invent `entity`-scoped
/// conflicts, inflate `mention_count`, and give graph expansion spurious
/// bridges between unrelated records.
///
/// Prose is the right domain for a proper-noun heuristic; code is not. This
/// removes ``` fences and `inline spans` (keeping a space so word chunks do
/// not weld across the removal) and leaves everything else untouched, so
/// entities named in the surrounding narrative are still captured.
///
/// KNOWN LIMITS, deliberate rather than overlooked (adversarial review,
/// 2026-08-11). This is a heuristic guard on a heuristic extractor; the
/// failure it prevents (fabricated entities) is worse than the failure it
/// allows (a missed entity), so every ambiguous case resolves toward
/// dropping:
/// - An UNTERMINATED fence drops the remaining text. This case is COMMON,
///   not exotic: callers chunk long documents, and a chunk routinely begins
///   inside a fenced block or ends with one open — so the tail genuinely is
///   code more often than it is prose. Keeping it would readmit exactly the
///   identifiers this function exists to remove.
/// - Escaped backticks (``\` ``) are treated as delimiters, so a span
///   between two of them is dropped. Costs a missed entity, never a false
///   one.
/// - Tilde fences and 4-space indented blocks are NOT recognised; code in
///   those forms still reaches the extractor. Fixing that needs a markdown
///   parser, which this deliberately is not.
fn strip_code(text: &str) -> std::borrow::Cow<'_, str> {
    if !text.contains('`') {
        return std::borrow::Cow::Borrowed(text);
    }
    let mut out = String::with_capacity(text.len());
    let mut rest = text;
    // Always act on the EARLIEST marker. Checking for a fence first is wrong:
    // in "`GET` text ```block```" the fence is found at 11 and the inline tick
    // at 0, so a fence-first branch emits "`GET` text " verbatim as prose and
    // the identifier this function exists to remove survives.
    while let Some(t) = rest.find('`') {
        out.push_str(&rest[..t]);
        out.push(' ');
        let after = &rest[t..];
        if let Some(body) = after.strip_prefix("```") {
            // A fenced region may contain single backticks; consume it whole
            // so they cannot be mis-paired as inline spans.
            match body.find("```") {
                Some(end) => rest = &body[end + 3..],
                None => return std::borrow::Cow::Owned(out), // unterminated: drop the tail
            }
        } else {
            let body = &after[1..];
            match body.find('`') {
                Some(end) => rest = &body[end + 1..],
                None => {
                    // Unterminated single backtick: keep the remainder as
                    // prose rather than discarding real text.
                    out.push_str(body);
                    return std::borrow::Cow::Owned(out);
                }
            }
        }
    }
    out.push_str(rest);
    std::borrow::Cow::Owned(out)
}

/// Extract candidate proper-noun entities from free-form text using a
/// capitalized-chunk heuristic. Groups consecutive capitalized words into
/// multi-word entities ("Alice Chen", "San Francisco", "Acme Corp") and
/// strips leading/trailing English stopwords.
///
/// Code spans and fenced blocks are removed first (see [`strip_code`]) —
/// identifiers are not proper nouns, and treating them as entities poisons
/// conflict detection and graph expansion.
///
/// This is intentionally not a full NER — it captures the common case of
/// people, companies, places, and products well enough that conflict
/// detection can fire without requiring users to call `/v1/relate` for every
/// entity. Acronyms, lowercase entities, and ambiguous mentions still need
/// explicit `relate()` calls to enter the graph.
pub fn extract_heuristic_entities(text: &str) -> Vec<String> {
    let stripped = strip_code(text);
    extract_heuristic_entities_inner(stripped.as_ref())
}

fn extract_heuristic_entities_inner(text: &str) -> Vec<String> {
    let mut entities: Vec<String> = Vec::new();
    let mut chunk: Vec<String> = Vec::new();

    let flush = |chunk: &mut Vec<String>, out: &mut Vec<String>| {
        while !chunk.is_empty() && is_entity_stopword(&chunk[0]) {
            chunk.remove(0);
        }
        // Trailing-stopword strip skips single-character tokens so multi-word
        // entities like "Series A" or "Version B" keep their letter suffix
        // (A is a stopword but is also a valid version designator when trailing).
        while let Some(last) = chunk.last() {
            if is_entity_stopword(last) && last.chars().count() > 1 {
                chunk.pop();
            } else {
                break;
            }
        }
        if !chunk.is_empty() && !is_prose_run(chunk) {
            let candidate = chunk.join(" ");
            let alpha_chars = candidate.chars().filter(|c| c.is_alphanumeric()).count();
            if alpha_chars >= 2 {
                out.push(candidate);
            }
        }
        chunk.clear();
    };

    for word in text
        .split(|c: char| !c.is_alphanumeric() && c != '\'')
        .filter(|s| !s.is_empty())
    {
        // A possessive clitic belongs to the grammar around a name, not to
        // the entity's identity. End the current chunk at the owner so
        // "Sol's Q2 plan" yields "Sol" and "Q2", rather than minting the
        // phantom compound "Sol's Q2". Internal apostrophes remain intact:
        // O'Brien stays O'Brien, and O'Brien's canonicalizes to O'Brien.
        let possessive = word
            .strip_suffix("'s")
            .or_else(|| word.strip_suffix("'S"))
            .or_else(|| word.strip_suffix('\''))
            .filter(|bare| !bare.is_empty());
        let entity_word = possessive.unwrap_or(word);
        let first = entity_word.chars().next().unwrap();
        let starts_upper = first.is_uppercase();
        let is_all_caps = entity_word.len() > 1
            && entity_word
                .chars()
                .all(|c| !c.is_alphabetic() || c.is_uppercase());

        let joins_chunk = if chunk.is_empty() {
            // Open a new chunk only on capitalized or all-caps tokens.
            starts_upper || is_all_caps
        } else {
            // Continue an existing chunk on capitalized words or short letter-suffixes
            // (e.g., "Series A", "Version B").
            starts_upper || is_all_caps || (entity_word.len() == 1 && first.is_ascii_uppercase())
        };

        if joins_chunk {
            chunk.push(entity_word.to_string());
            if possessive.is_some() {
                flush(&mut chunk, &mut entities);
            }
        } else {
            flush(&mut chunk, &mut entities);
        }
    }
    flush(&mut chunk, &mut entities);

    // Deduplicate while preserving first-appearance order.
    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
    entities.retain(|e| seen.insert(e.clone()));
    entities
}

// ── Heuristic relation extraction (RFC 006 Phase 1) ──

/// A candidate relation extracted from text by pattern matching.
#[derive(Debug, Clone)]
pub struct RelationCandidate {
    pub src: String,
    pub rel_type: String,
    pub dst: String,
    pub polarity: i32,           // 1=positive, -1=negative
    pub modality: String,        // asserted, reported, hypothetical, denied
    pub confidence_band: String, // low, medium, high
}

/// Relation patterns: keyword phrases that appear BETWEEN two entities
/// and indicate a specific relationship. Each pattern maps to a rel_type.
const RELATION_PATTERNS: &[(&[&str], &str)] = &[
    // Role-based (entity A <pattern> entity B → rel_type)
    (
        &["is the ceo of", "is ceo of", "serves as ceo of"],
        "ceo_of",
    ),
    (
        &["is the cto of", "is cto of", "serves as cto of"],
        "cto_of",
    ),
    (
        &["is the cfo of", "is cfo of", "serves as cfo of"],
        "cfo_of",
    ),
    (
        &["is the founder of", "is founder of", "co-founded"],
        "founded",
    ),
    (&["founded"], "founded"),
    (&["leads", "heads", "runs", "manages", "directs"], "leads"),
    (
        &[
            "works at",
            "works for",
            "employed at",
            "employed by",
            "joined",
        ],
        "works_at",
    ),
    // Location/origin
    (&["was born in", "born in"], "born_in"),
    (
        &[
            "is headquartered in",
            "headquartered in",
            "is based in",
            "based in",
            "located in",
        ],
        "headquartered_in",
    ),
    // Personal
    (&["is married to", "married to", "wed to"], "married_to"),
    // Corporate
    (
        &["acquired", "bought", "purchased", "took over"],
        "acquired",
    ),
    (
        &[
            "is a subsidiary of",
            "subsidiary of",
            "is owned by",
            "owned by",
        ],
        "subsidiary_of",
    ),
    // Language/skill
    (&["speaks", "is fluent in"], "speaks"),
    // Generic membership/part-of
    (
        &["is a member of", "member of", "belongs to", "part of"],
        "member_of",
    ),
    (&["reports to"], "reports_to"),
];

/// Possessive/appositive reverse patterns: "ORG's CEO, PERSON" or "ORG's CEO PERSON"
const REVERSE_ROLE_PATTERNS: &[(&str, &str)] = &[
    ("ceo", "ceo_of"),
    ("cto", "cto_of"),
    ("cfo", "cfo_of"),
    ("founder", "founded"),
    ("president", "leads"),
    ("director", "leads"),
    ("head", "leads"),
];

/// Extract candidate relations from text using entities as anchors.
///
/// For each ordered pair of entities (A before B in text), examines the
/// text between them for relation-indicating keywords. Also checks for
/// negation cues in the window to set polarity, and tense cues to infer
/// past-tense (which callers can use for valid_to).
///
/// Returns high-precision, low-recall candidates — only emits when a
/// clear keyword pattern matches. Designed for the RFC 006 Phase 1
/// relation whitelist.
pub fn extract_heuristic_relations(text: &str, entities: &[String]) -> Vec<RelationCandidate> {
    if entities.len() < 2 {
        return vec![];
    }

    let text_lower = text.to_lowercase();
    let mut candidates: Vec<RelationCandidate> = Vec::new();

    // Find position of each entity in the text (case-insensitive)
    let mut entity_positions: Vec<(usize, &str)> = Vec::new();
    for entity in entities {
        let entity_lower = entity.to_lowercase();
        if let Some(pos) = text_lower.find(&entity_lower) {
            entity_positions.push((pos, entity.as_str()));
        }
    }
    entity_positions.sort_by_key(|(pos, _)| *pos);

    // For each adjacent pair, check the text between them
    for i in 0..entity_positions.len() {
        for j in (i + 1)..entity_positions.len() {
            let (pos_a, entity_a) = entity_positions[i];
            let (pos_b, entity_b) = entity_positions[j];

            // Skip pairs too far apart (likely different sentences)
            if pos_b - pos_a > 150 {
                continue;
            }

            let between_start = pos_a + entity_a.to_lowercase().len();
            let between_end = pos_b;
            if between_start >= between_end || between_end > text_lower.len() {
                continue;
            }

            let between = text_lower[between_start..between_end].trim();
            if between.is_empty() {
                continue;
            }

            // Check negation in the between-window, then strip negation
            // words so pattern matching still works on "is NOT the CEO of"
            let has_negation = NEGATION_CUES
                .iter()
                .any(|cue| between.split_whitespace().any(|w| w == *cue));
            let polarity = if has_negation { -1 } else { 1 };
            let between_stripped: String = between
                .split_whitespace()
                .filter(|w| !NEGATION_CUES.contains(w))
                .collect::<Vec<_>>()
                .join(" ");

            // Check modality cues
            let modality = if MODALITY_CUES.iter().any(|cue| between.contains(cue)) {
                "reported"
            } else {
                "asserted"
            };

            // Match forward patterns: entity_a <pattern> entity_b
            // Uses between_stripped (negation removed) for matching.
            for (patterns, rel_type) in RELATION_PATTERNS {
                for pattern in *patterns {
                    if between_stripped.contains(pattern) {
                        candidates.push(RelationCandidate {
                            src: entity_a.to_string(),
                            rel_type: rel_type.to_string(),
                            dst: entity_b.to_string(),
                            polarity,
                            modality: modality.to_string(),
                            confidence_band: "medium".to_string(),
                        });
                        break; // one match per pattern group per pair
                    }
                }
            }

            // Check possessive/appositive reverse: "Acme's CEO, Alice" → ceo_of(Alice, Acme)
            for (role_keyword, rel_type) in REVERSE_ROLE_PATTERNS {
                let possessive = format!("'s {}", role_keyword);
                let possessive2 = format!("s {}", role_keyword);
                if between_stripped.contains(&possessive) || between_stripped.contains(&possessive2)
                {
                    // Reversed: entity_a is the org, entity_b is the person
                    candidates.push(RelationCandidate {
                        src: entity_b.to_string(), // person
                        rel_type: rel_type.to_string(),
                        dst: entity_a.to_string(), // org
                        polarity,
                        modality: modality.to_string(),
                        confidence_band: "medium".to_string(),
                    });
                    break;
                }
            }
        }
    }

    // Deduplicate: same (src, rel_type, dst) keeps highest confidence
    let mut seen = std::collections::HashSet::new();
    candidates.retain(|c| seen.insert((c.src.clone(), c.rel_type.clone(), c.dst.clone())));

    candidates
}

// ── Text feature analysis (Phase 0 audit data for RFC 006) ──

/// Cues that indicate a statement is negated. Window-scanned around pattern
/// matches to flag `polarity=negative` in v0.6.0. In v0.5.13 we only count
/// occurrences for audit telemetry.
const NEGATION_CUES: &[&str] = &[
    "not", "no", "never", "denied", "refuted", "isn't", "wasn't", "aren't", "weren't", "doesn't",
    "didn't", "disputes", "denies",
];

/// Cues that indicate a statement has temporal scope. Used to flag that a
/// memory would benefit from `valid_from` / `valid_to` qualifiers.
const TEMPORAL_CUES: &[&str] = &[
    "was",
    "were",
    "until",
    "before",
    "after",
    "since",
    "during",
    "former",
    "current",
    "currently",
    "previously",
    "recently",
    "now",
    "then",
    "later",
    "earlier",
    "ago",
    "yesterday",
    "tomorrow",
];

/// Cues that indicate modality (hypothetical, reported, quoted).
const MODALITY_CUES: &[&str] = &[
    "may",
    "might",
    "allegedly",
    "reportedly",
    "rumor",
    "rumored",
    "said",
    "claims",
    "according",
    "stated",
    "announced",
];

/// Compound-sentence separators that a v0.6.0 extractor should split on
/// before running patterns. Counting these at audit time tells us how many
/// real-world memories contain multiple claims per write.
const COMPOUND_MARKERS: &[&str] = &[
    "; ",
    ", then ",
    ", subsequently ",
    " but ",
    " however ",
    " although ",
];

/// Text features collected for extraction-audit telemetry (RFC 006 Phase 0).
/// Captures everything the v0.6.0 extractor would need to know without
/// changing any storage behavior — purely observational.
#[derive(Debug, Clone, Default)]
pub struct TextFeatures {
    pub char_length: usize,
    pub sentence_count: usize,
    pub entity_count: usize,
    pub negation_cue_count: usize,
    pub temporal_cue_count: usize,
    pub modality_cue_count: usize,
    pub has_compound_markers: bool,
    pub likely_assertion: bool,
}

/// Compute text features for extraction audit. Pure function, no I/O.
pub fn analyze_text_features(text: &str, extracted_entities: &[String]) -> TextFeatures {
    let lower = text.to_lowercase();
    let tokens: Vec<&str> = text
        .split(|c: char| !c.is_alphanumeric() && c != '\'')
        .filter(|s| !s.is_empty())
        .collect();
    let tokens_lower: Vec<String> = tokens.iter().map(|t| t.to_lowercase()).collect();

    let sentence_count = text
        .chars()
        .filter(|c| matches!(c, '.' | '!' | '?'))
        .count()
        .max(1);

    let negation_cue_count = tokens_lower
        .iter()
        .filter(|t| NEGATION_CUES.contains(&t.as_str()))
        .count();

    let temporal_cue_count = tokens_lower
        .iter()
        .filter(|t| TEMPORAL_CUES.contains(&t.as_str()))
        .count();

    let modality_cue_count = tokens_lower
        .iter()
        .filter(|t| MODALITY_CUES.contains(&t.as_str()))
        .count();

    let has_compound_markers = COMPOUND_MARKERS.iter().any(|m| lower.contains(m));

    // Rough "assertion?" signal: not a question, has at least 2 tokens, not
    // pure modality/rumor. Used to estimate what fraction of agent writes
    // the v0.6.0 extractor should try to process at all.
    let likely_assertion =
        !text.trim_end().ends_with('?') && tokens.len() >= 2 && modality_cue_count == 0;

    TextFeatures {
        char_length: text.chars().count(),
        sentence_count,
        entity_count: extracted_entities.len(),
        negation_cue_count,
        temporal_cue_count,
        modality_cue_count,
        has_compound_markers,
        likely_assertion,
    }
}

// ── Entity type classification ──

/// Tech terms that should NOT be classified as person names even if title-cased/all-caps.
const TECH_BLOCKLIST: &[&str] = &[
    "faiss",
    "onnx",
    "scann",
    "redis",
    "kafka",
    "docker",
    "kubernetes",
    "react",
    "python",
    "rust",
    "java",
    "swift",
    "flutter",
    "pytorch",
    "tensorflow",
    "numpy",
    "pandas",
    "spark",
    "hadoop",
    "nginx",
    "postgres",
    "mysql",
    "sqlite",
    "graphql",
    "grpc",
    "oauth",
    "jwt",
    "html",
    "css",
    "api",
    "sdk",
    "ml",
    "ai",
    "gpu",
    "cpu",
    "ram",
    "ssd",
    "aws",
    "gcp",
    "claude",
    "openai",
    "anthropic",
    "gemini",
    "llama",
    "ollama",
];

/// Words that indicate the entity is NOT a person when used as first word.
const NON_PERSON_PREFIXES: &[&str] = &[
    "project",
    "team",
    "company",
    "group",
    "department",
    "org",
    "the",
    "operation",
    "task",
    "plan",
    "system",
    "service",
    "app",
    "tool",
    "code",
    "server",
    "client",
    "api",
    "db",
    "database",
    "agent",
    "model",
    "version",
    "release",
    "build",
    "deploy",
    "config",
];

/// Classify an entity name into a type: "person", "tech", or "unknown".
/// This is a name-only heuristic — prefer `classify_with_relationship()` when
/// relationship context is available.
pub fn classify_entity_type(name: &str) -> &'static str {
    let trimmed = name.trim();
    if trimmed.is_empty() {
        return "unknown";
    }
    let lower = trimmed.to_lowercase();

    // Check tech blocklist
    if TECH_BLOCKLIST.contains(&lower.as_str()) {
        return "tech";
    }

    // All-caps multi-char → tech (e.g., "FAISS", "ONNX")
    if trimmed.len() > 1
        && trimmed
            .chars()
            .all(|c| c.is_uppercase() || !c.is_alphabetic())
    {
        return "tech";
    }

    // Multi-word title-case (e.g., "Priya Sharma", "Sarah Chen") → likely person
    // But NOT if the first word is a non-person prefix (e.g., "Project Athena", "Claude Code")
    if trimmed.contains(' ') {
        let words: Vec<&str> = trimmed.split_whitespace().collect();
        if words.len() == 2
            && words
                .iter()
                .all(|w| w.chars().next().map(|c| c.is_uppercase()).unwrap_or(false))
        {
            let first_lower = words[0].to_lowercase();
            if NON_PERSON_PREFIXES.contains(&first_lower.as_str()) {
                return "unknown";
            }
            // Also reject if any word is in tech blocklist
            if words
                .iter()
                .any(|w| TECH_BLOCKLIST.contains(&w.to_lowercase().as_str()))
            {
                return "tech";
            }
            return "person";
        }
    }

    // Single-word classification is unreliable (Bangalore, Flipkart, Arjun all
    // look the same). Return "unknown" and let relationship context decide.
    "unknown"
}

/// Relationship types that imply both src and dst are persons.
const PERSON_PERSON_RELS: &[&str] = &[
    "married_to",
    "mother_of",
    "father_of",
    "daughter_of",
    "son_of",
    "sister_of",
    "brother_of",
    "sibling_of",
    "parent_of",
    "child_of",
    "knows",
    "friends_with",
    "met",
    "dating",
    "engaged_to",
    "mentors",
    "mentored_by",
    "reports_to",
    "manages",
    "colleagues",
    "roommate",
    "neighbor",
    "called",
    "texted",
    "messaged",
    "date_night",
];

/// Relationship types where dst is a place.
const PLACE_DST_RELS: &[&str] = &[
    "lives_in",
    "born_in",
    "grew_up_in",
    "located_in",
    "based_in",
    "visited",
    "moved_to",
    "traveled_to",
    "from",
];

/// Relationship types where dst is an organization / institution.
const ORG_DST_RELS: &[&str] = &[
    "works_at",
    "works_for",
    "employed_at",
    "employed_by",
    "studied_at",
    "attended",
    "enrolled_in",
    "graduated_from",
    "member_of",
    "belongs_to",
    "founded",
];

/// Relationship types where dst is tech/tool (src is project or person).
const TECH_DST_RELS: &[&str] = &[
    "built_with",
    "uses",
    "depends_on",
    "integrates",
    "requires",
    "written_in",
    "coded_in",
    "implemented_with",
    "powered_by",
    "runs_on",
    "compiled_with",
];

/// Relationship types where dst is infrastructure.
const INFRA_DST_RELS: &[&str] = &[
    "deployed_on",
    "hosted_on",
    "deployed_to",
    "hosted_at",
    "runs_on_infra",
    "served_by",
];

/// Relationship types where src is a person and dst is a project/thing.
const PERSON_PROJECT_RELS: &[&str] = &[
    "works_on",
    "contributes_to",
    "maintains",
    "leads",
    "created",
    "built",
    "designed",
    "architected",
    "owns",
];

/// Relationship types where src is a project and dst is a project (dependency).
const PROJECT_PROJECT_RELS: &[&str] = &[
    "depends_on_project",
    "extends",
    "forks",
    "replaces",
    "supersedes",
    "derived_from",
];

/// Relationship types where dst is an event or activity.
const EVENT_DST_RELS: &[&str] = &[
    "attended_event",
    "participated_in",
    "scheduled_for",
    "presented_at",
    "spoke_at",
];

/// Relationship types where dst is a concept/topic.
const CONCEPT_DST_RELS: &[&str] = &[
    "interested_in",
    "studies",
    "researches",
    "specializes_in",
    "expert_in",
    "learning",
    "teaches",
];

/// Classify entity types using relationship semantics.
/// Returns (src_type, dst_type) — either may be "unknown" if not inferable.
pub fn classify_with_relationship(
    src: &str,
    dst: &str,
    rel_type: &str,
) -> (&'static str, &'static str) {
    let rel_lower = rel_type.to_lowercase();
    let rel = rel_lower.as_str();

    // Person-person relationships
    if PERSON_PERSON_RELS.contains(&rel) {
        return ("person", "person");
    }

    // Person → Place relationships
    if PLACE_DST_RELS.contains(&rel) {
        return ("person", "place");
    }

    // Person → Organization relationships
    if ORG_DST_RELS.contains(&rel) {
        return ("person", "organization");
    }

    // * → Tech/Tool relationships (src type from name heuristic)
    if TECH_DST_RELS.contains(&rel) {
        let src_type = classify_entity_type(src);
        return (
            if src_type == "unknown" {
                "project"
            } else {
                src_type
            },
            "tech",
        );
    }

    // * → Infrastructure relationships
    if INFRA_DST_RELS.contains(&rel) {
        let src_type = classify_entity_type(src);
        return (
            if src_type == "unknown" {
                "project"
            } else {
                src_type
            },
            "infrastructure",
        );
    }

    // Person → Project relationships
    if PERSON_PROJECT_RELS.contains(&rel) {
        return ("person", "project");
    }

    // Project → Project relationships
    if PROJECT_PROJECT_RELS.contains(&rel) {
        return ("project", "project");
    }

    // * → Event relationships
    if EVENT_DST_RELS.contains(&rel) {
        return (classify_entity_type(src), "event");
    }

    // Person → Concept/Topic relationships
    if CONCEPT_DST_RELS.contains(&rel) {
        return ("person", "concept");
    }

    // Fall back to name-based heuristics
    (classify_entity_type(src), classify_entity_type(dst))
}

/// Given a set of memory RIDs, find all entities those memories are linked to.
pub fn entities_for_memories(conn: &Connection, rids: &[&str]) -> Result<Vec<String>> {
    if rids.is_empty() {
        return Ok(vec![]);
    }
    let placeholders: String = (0..rids.len())
        .map(|i| format!("?{}", i + 1))
        .collect::<Vec<_>>()
        .join(",");
    let sql = format!(
        "SELECT DISTINCT entity_name FROM memory_entities WHERE memory_rid IN ({placeholders})"
    );
    let mut stmt = conn.prepare(&sql)?;
    let param_values: Vec<Box<dyn rusqlite::types::ToSql>> = rids
        .iter()
        .map(|r| Box::new(r.to_string()) as Box<dyn rusqlite::types::ToSql>)
        .collect();
    let params_ref: Vec<&dyn rusqlite::types::ToSql> =
        param_values.iter().map(|p| p.as_ref()).collect();
    let entities = stmt
        .query_map(params_ref.as_slice(), |row| row.get(0))?
        .collect::<std::result::Result<Vec<String>, _>>()?;
    Ok(entities)
}

/// Given a set of entity names, find all memory RIDs connected to those entities.
pub fn memories_for_entities(conn: &Connection, entity_names: &[&str]) -> Result<HashSet<String>> {
    if entity_names.is_empty() {
        return Ok(HashSet::new());
    }
    let placeholders: String = (0..entity_names.len())
        .map(|i| format!("?{}", i + 1))
        .collect::<Vec<_>>()
        .join(",");
    let sql = format!(
        "SELECT DISTINCT memory_rid FROM memory_entities WHERE entity_name IN ({placeholders})"
    );
    let mut stmt = conn.prepare(&sql)?;
    let param_values: Vec<Box<dyn rusqlite::types::ToSql>> = entity_names
        .iter()
        .map(|e| Box::new(e.to_string()) as Box<dyn rusqlite::types::ToSql>)
        .collect();
    let params_ref: Vec<&dyn rusqlite::types::ToSql> =
        param_values.iter().map(|p| p.as_ref()).collect();
    let rids = stmt
        .query_map(params_ref.as_slice(), |row| row.get(0))?
        .collect::<std::result::Result<HashSet<String>, _>>()?;
    Ok(rids)
}

/// Expand entity set N hops via the edges table (BFS).
/// Returns (entity_name, hops_from_seed, cumulative_edge_weight).
/// Seeds are returned with hops=0 and weight=1.0.
pub fn expand_entities_nhop(
    conn: &Connection,
    seeds: &[&str],
    max_hops: u8,
    max_entities: usize,
) -> Result<Vec<(String, u8, f64)>> {
    let mut result: Vec<(String, u8, f64)> = Vec::new();
    let mut visited: HashMap<String, (u8, f64)> = HashMap::new();

    // Initialize with seeds
    for s in seeds {
        visited.insert(s.to_string(), (0, 1.0));
        result.push((s.to_string(), 0, 1.0));
    }

    let mut frontier: VecDeque<(String, u8, f64)> =
        seeds.iter().map(|s| (s.to_string(), 0u8, 1.0f64)).collect();

    while let Some((entity, hops, weight)) = frontier.pop_front() {
        if hops >= max_hops || result.len() >= max_entities {
            break;
        }

        // Find neighbors via edges (both directions)
        let mut stmt = conn.prepare(
            "SELECT src, dst, weight FROM edges WHERE (src = ?1 OR dst = ?1) AND tombstoned = 0",
        )?;
        let neighbors: Vec<(String, f64)> = stmt
            .query_map(params![entity], |row| {
                let src: String = row.get(0)?;
                let dst: String = row.get(1)?;
                let w: f64 = row.get(2)?;
                let neighbor = if src == entity { dst } else { src };
                Ok((neighbor, w))
            })?
            .collect::<std::result::Result<Vec<_>, _>>()?;

        for (neighbor, edge_weight) in neighbors {
            if visited.contains_key(&neighbor) {
                continue;
            }
            if result.len() >= max_entities {
                break;
            }
            let cumulative = weight * edge_weight;
            let next_hops = hops + 1;
            visited.insert(neighbor.clone(), (next_hops, cumulative));
            result.push((neighbor.clone(), next_hops, cumulative));
            if next_hops < max_hops {
                frontier.push_back((neighbor, next_hops, cumulative));
            }
        }
    }

    Ok(result)
}

/// Compute graph proximity score for a memory based on its entity connections.
/// Returns the maximum proximity across all entities the memory is linked to.
/// proximity = cumulative_weight / 2^hops  (steeper decay to stay discriminative)
/// Seeds (hops=0) → 1.0, 1-hop → 0.5, 2-hop → 0.25
pub fn graph_proximity(
    conn: &Connection,
    memory_rid: &str,
    expanded_entities: &HashMap<String, (u8, f64)>,
) -> Result<f64> {
    let mem_entities: Vec<String> = conn
        .prepare("SELECT entity_name FROM memory_entities WHERE memory_rid = ?1")?
        .query_map(params![memory_rid], |row| row.get(0))?
        .collect::<std::result::Result<Vec<_>, _>>()?;

    let mut max_proximity = 0.0f64;
    for entity in &mem_entities {
        if let Some(&(hops, weight)) = expanded_entities.get(entity) {
            let prox = weight / f64::powf(2.0, hops as f64);
            if prox > max_proximity {
                max_proximity = prox;
            }
        }
    }
    Ok(max_proximity)
}

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

    #[test]
    fn test_extract_heuristic_entities_basic_names() {
        let got = extract_heuristic_entities("Alice Chen is the CEO of Acme Corp");
        assert!(got.contains(&"Alice Chen".to_string()), "got: {:?}", got);
        assert!(got.contains(&"Acme Corp".to_string()), "got: {:?}", got);
        // CEO is all-caps standalone — should appear as an entity candidate.
        assert!(got.contains(&"CEO".to_string()), "got: {:?}", got);
    }

    #[test]
    fn test_extract_heuristic_entities_strips_sentence_start() {
        let got = extract_heuristic_entities("The database backend is PostgreSQL");
        assert_eq!(got, vec!["PostgreSQL".to_string()]);
    }

    #[test]
    fn test_extract_heuristic_entities_multi_word_place() {
        let got = extract_heuristic_entities("Acme is headquartered in San Francisco");
        assert!(got.contains(&"Acme".to_string()), "got: {:?}", got);
        assert!(got.contains(&"San Francisco".to_string()), "got: {:?}", got);
    }

    #[test]
    fn test_extract_heuristic_entities_single_letter_suffix() {
        let got = extract_heuristic_entities("Series A funding was 20 million dollars");
        assert!(got.contains(&"Series A".to_string()), "got: {:?}", got);
    }

    #[test]
    fn test_extract_heuristic_entities_dedupe() {
        let got = extract_heuristic_entities("Alice met Alice at the cafe");
        let alice_count = got.iter().filter(|e| *e == "Alice").count();
        assert_eq!(alice_count, 1);
    }

    #[test]
    fn test_extract_heuristic_entities_empty_on_lowercase() {
        let got = extract_heuristic_entities("the quick brown fox jumps over the lazy dog");
        assert!(got.is_empty(), "got: {:?}", got);
    }

    // ── Relation extraction tests ──

    #[test]
    fn test_extract_relations_ceo_of() {
        let entities = vec!["Alice Chen".to_string(), "Acme Corp".to_string()];
        let rels = extract_heuristic_relations("Alice Chen is the CEO of Acme Corp", &entities);
        assert_eq!(rels.len(), 1, "got: {:?}", rels);
        assert_eq!(rels[0].src, "Alice Chen");
        assert_eq!(rels[0].rel_type, "ceo_of");
        assert_eq!(rels[0].dst, "Acme Corp");
        assert_eq!(rels[0].polarity, 1);
    }

    #[test]
    fn test_extract_relations_works_at() {
        let entities = vec!["Bob".to_string(), "Google".to_string()];
        let rels = extract_heuristic_relations("Bob works at Google as an engineer", &entities);
        assert!(
            rels.iter().any(|r| r.rel_type == "works_at"),
            "got: {:?}",
            rels
        );
    }

    #[test]
    fn test_extract_relations_headquartered() {
        let entities = vec!["Acme".to_string(), "San Francisco".to_string()];
        let rels = extract_heuristic_relations("Acme is headquartered in San Francisco", &entities);
        assert!(
            rels.iter().any(|r| r.rel_type == "headquartered_in"),
            "got: {:?}",
            rels
        );
    }

    #[test]
    fn test_extract_relations_negation_detected() {
        let entities = vec!["Alice".to_string(), "Acme".to_string()];
        let rels = extract_heuristic_relations("Alice is not the CEO of Acme", &entities);
        assert_eq!(rels.len(), 1);
        assert_eq!(rels[0].polarity, -1, "negation should set polarity to -1");
    }

    #[test]
    fn test_extract_relations_no_match_unrelated() {
        let entities = vec!["Alice".to_string(), "Bob".to_string()];
        let rels = extract_heuristic_relations("Alice and Bob went for coffee", &entities);
        assert!(
            rels.is_empty(),
            "should not extract relation from unrelated text, got: {:?}",
            rels
        );
    }

    #[test]
    fn test_extract_relations_multiple_pairs() {
        let entities = vec![
            "Alice".to_string(),
            "Acme".to_string(),
            "San Francisco".to_string(),
        ];
        let rels = extract_heuristic_relations(
            "Alice is the CEO of Acme which is headquartered in San Francisco",
            &entities,
        );
        assert!(
            rels.len() >= 2,
            "should find CEO + headquartered, got: {:?}",
            rels
        );
    }

    #[test]
    fn test_extract_relations_needs_two_entities() {
        let entities = vec!["Alice".to_string()];
        let rels = extract_heuristic_relations("Alice is the CEO", &entities);
        assert!(
            rels.is_empty(),
            "cannot extract relation with only one entity"
        );
    }

    // ── Text feature analysis tests ──

    #[test]
    fn test_analyze_text_features_basic_assertion() {
        let entities = vec!["Alice Chen".to_string(), "Acme Corp".to_string()];
        let f = analyze_text_features("Alice Chen is the CEO of Acme Corp", &entities);
        assert_eq!(f.entity_count, 2);
        assert_eq!(f.negation_cue_count, 0);
        assert_eq!(f.modality_cue_count, 0);
        assert!(f.likely_assertion);
        assert!(!f.has_compound_markers);
    }

    #[test]
    fn test_analyze_text_features_negation() {
        let f = analyze_text_features("Alice is not the CEO of Acme", &[]);
        assert_eq!(f.negation_cue_count, 1);
    }

    #[test]
    fn test_analyze_text_features_temporal() {
        let f = analyze_text_features("Alice was previously the CEO before 2024", &[]);
        assert!(f.temporal_cue_count >= 2, "got: {}", f.temporal_cue_count);
    }

    #[test]
    fn test_analyze_text_features_modality_suppresses_assertion() {
        let f = analyze_text_features("Alice may become CEO allegedly", &[]);
        assert!(f.modality_cue_count >= 2);
        assert!(!f.likely_assertion);
    }

    #[test]
    fn test_analyze_text_features_compound() {
        let f = analyze_text_features("Alice was CEO until 2024; then Bob took over", &[]);
        assert!(f.has_compound_markers);
    }

    #[test]
    fn test_analyze_text_features_question_not_assertion() {
        let f = analyze_text_features("Who is the CEO of Acme?", &[]);
        assert!(!f.likely_assertion);
    }

    #[test]
    fn test_extract_heuristic_entities_distinct_people() {
        // Regression guard for the false-merge case that motivated this:
        // two sentences structurally similar but referring to different people.
        let a = extract_heuristic_entities("Alice Chen is the CEO of Acme Corp");
        let b = extract_heuristic_entities("Sarah Kim is the CTO of Acme Corp");
        let a_set: std::collections::HashSet<_> = a.iter().collect();
        let b_set: std::collections::HashSet<_> = b.iter().collect();
        // They share Acme Corp but differ on person name — disjointness on people.
        assert!(a_set.contains(&"Alice Chen".to_string()));
        assert!(b_set.contains(&"Sarah Kim".to_string()));
        assert!(!a_set.contains(&"Sarah Kim".to_string()));
        assert!(!b_set.contains(&"Alice Chen".to_string()));
    }

    fn setup_db() -> YantrikDB {
        let db = YantrikDB::new(":memory:", 4).unwrap();
        // Create entities and edges
        db.relate("Alice", "Bob", "knows", 1.0).unwrap();
        db.relate("Bob", "Charlie", "knows", 0.8).unwrap();
        db.relate("Alice", "ProjectX", "works_on", 1.0).unwrap();
        db.relate("Dave", "ProjectX", "works_on", 0.9).unwrap();

        // Record memories and link to entities
        let emb = vec![1.0f32, 0.0, 0.0, 0.0];
        let r1 = db
            .record(
                "Alice discussed the plan",
                "episodic",
                0.5,
                0.0,
                604800.0,
                &serde_json::json!({}),
                &emb,
                "default",
                0.8,
                "general",
                "user",
                None,
            )
            .unwrap();
        let r2 = db
            .record(
                "Bob reviewed the code",
                "episodic",
                0.5,
                0.0,
                604800.0,
                &serde_json::json!({}),
                &emb,
                "default",
                0.8,
                "general",
                "user",
                None,
            )
            .unwrap();
        let r3 = db
            .record(
                "Charlie deployed to production",
                "episodic",
                0.5,
                0.0,
                604800.0,
                &serde_json::json!({}),
                &emb,
                "default",
                0.8,
                "general",
                "user",
                None,
            )
            .unwrap();

        db.link_memory_entity(&r1, "Alice").unwrap();
        db.link_memory_entity(&r1, "ProjectX").unwrap();
        db.link_memory_entity(&r2, "Bob").unwrap();
        db.link_memory_entity(&r3, "Charlie").unwrap();

        db
    }

    #[test]
    fn test_entities_for_memories() {
        let db = setup_db();
        // Get the first memory's rid
        let rid: String = db
            .conn()
            .query_row(
                "SELECT rid FROM memories ORDER BY created_at LIMIT 1",
                [],
                |row| row.get(0),
            )
            .unwrap();

        let entities = entities_for_memories(&*db.conn(), &[&rid]).unwrap();
        assert!(entities.contains(&"Alice".to_string()));
        assert!(entities.contains(&"ProjectX".to_string()));
    }

    #[test]
    fn test_memories_for_entities() {
        let db = setup_db();
        let rids = memories_for_entities(&*db.conn(), &["Alice"]).unwrap();
        assert_eq!(rids.len(), 1); // Only the Alice memory is linked
    }

    #[test]
    fn test_expand_1hop() {
        let db = setup_db();
        let expanded = expand_entities_nhop(&*db.conn(), &["Alice"], 1, 30).unwrap();
        let names: HashSet<String> = expanded.iter().map(|(n, _, _)| n.clone()).collect();
        // Alice (seed) + Bob (knows) + ProjectX (works_on)
        assert!(names.contains("Alice"));
        assert!(names.contains("Bob"));
        assert!(names.contains("ProjectX"));
    }

    #[test]
    fn test_expand_2hop() {
        let db = setup_db();
        let expanded = expand_entities_nhop(&*db.conn(), &["Alice"], 2, 30).unwrap();
        let names: HashSet<String> = expanded.iter().map(|(n, _, _)| n.clone()).collect();
        // 2-hop from Alice: Alice->Bob->Charlie, Alice->ProjectX->Dave
        assert!(names.contains("Charlie"));
        assert!(names.contains("Dave"));
    }

    #[test]
    fn test_expand_budget_limit() {
        let db = setup_db();
        let expanded = expand_entities_nhop(&*db.conn(), &["Alice"], 2, 3).unwrap();
        assert!(expanded.len() <= 3);
    }

    #[test]
    fn test_no_tombstoned_edges() {
        let db = setup_db();
        // Tombstone the Alice->Bob edge
        db.conn()
            .execute(
                "UPDATE claims SET tombstoned = 1 WHERE src = 'Alice' AND dst = 'Bob'",
                [],
            )
            .unwrap();
        let expanded = expand_entities_nhop(&*db.conn(), &["Alice"], 1, 30).unwrap();
        let names: HashSet<String> = expanded.iter().map(|(n, _, _)| n.clone()).collect();
        // Bob should NOT be reachable via tombstoned edge
        assert!(!names.contains("Bob"));
        // ProjectX should still be reachable
        assert!(names.contains("ProjectX"));
    }

    #[test]
    fn test_graph_proximity_score() {
        let db = setup_db();
        let rid: String = db
            .conn()
            .query_row(
                "SELECT rid FROM memories ORDER BY created_at LIMIT 1",
                [],
                |row| row.get(0),
            )
            .unwrap();

        let mut expanded = HashMap::new();
        expanded.insert("Alice".to_string(), (0u8, 1.0f64));
        expanded.insert("ProjectX".to_string(), (1u8, 1.0f64));

        let prox = graph_proximity(&*db.conn(), &rid, &expanded).unwrap();
        // Alice is hops=0 → proximity = 1.0 / (0+1) = 1.0
        assert!((prox - 1.0).abs() < 1e-10);
    }

    // ── Word-boundary matching tests ──

    #[test]
    fn test_tokenize_basic() {
        let tokens = tokenize("What is Sarah working on?");
        assert_eq!(tokens, vec!["what", "is", "sarah", "working", "on"]);
    }

    #[test]
    fn test_tokenize_splits_apostrophes() {
        // INVERTED 2026-08-06 (wheel C5a). This test used to pin the
        // apostrophe exemption — the convicted defect itself: it kept
        // "daughter's" whole, which meant "Taylor's" never matched
        // entity "taylor" and possessive queries silently lost entity
        // resolution on the default path. The possessive must now
        // yield the bare token so the entity is reachable.
        let tokens = tokenize("daughter's school play");
        assert_eq!(tokens, vec!["daughter", "s", "school", "play"]);
    }

    #[test]
    fn test_entity_matches_single_word() {
        let tokens = tokenize("Sarah discussed the plan with Mike");
        assert!(entity_matches_text("Sarah", &tokens));
        assert!(entity_matches_text("Mike", &tokens));
        assert!(!entity_matches_text("Sara", &tokens)); // partial ≠ match
    }

    #[test]
    fn test_entity_matches_multi_word() {
        let tokens = tokenize("The data pipeline crashed during migration");
        assert!(entity_matches_text("data pipeline", &tokens));
        assert!(!entity_matches_text("data migration", &tokens)); // non-contiguous
    }

    #[test]
    fn test_entity_no_substring_false_positive() {
        let tokens = tokenize("The database was updated successfully");
        // "data" should NOT match inside "database"
        assert!(!entity_matches_text("data", &tokens));
    }

    #[test]
    fn test_entity_matches_case_insensitive() {
        let tokens = tokenize("We evaluated FAISS for vector search");
        assert!(entity_matches_text("FAISS", &tokens));
        assert!(entity_matches_text("faiss", &tokens));
    }

    // ── Entity type classification tests ──

    #[test]
    fn test_classify_name_only_ambiguous() {
        // Single-word title-case is now "unknown" without relationship context
        assert_eq!(classify_entity_type("Sarah"), "unknown");
        assert_eq!(classify_entity_type("Bangalore"), "unknown");
        assert_eq!(classify_entity_type("Flipkart"), "unknown");
    }

    #[test]
    fn test_classify_name_multi_word_person() {
        // Multi-word title-case full names are still "person"
        assert_eq!(classify_entity_type("Sarah Chen"), "person");
        assert_eq!(classify_entity_type("Priya Sharma"), "person");
    }

    #[test]
    fn test_classify_tech_blocklist() {
        assert_eq!(classify_entity_type("FAISS"), "tech");
        assert_eq!(classify_entity_type("ONNX"), "tech");
        assert_eq!(classify_entity_type("Redis"), "tech");
        assert_eq!(classify_entity_type("Python"), "tech");
    }

    #[test]
    fn test_classify_tech_allcaps() {
        assert_eq!(classify_entity_type("GPU"), "tech");
        assert_eq!(classify_entity_type("API"), "tech");
    }

    #[test]
    fn test_classify_unknown() {
        assert_eq!(classify_entity_type("recommendation engine"), "unknown");
        assert_eq!(classify_entity_type("data pipeline"), "unknown");
        assert_eq!(classify_entity_type("sleep patterns"), "unknown");
    }

    // ── Relationship-based classification tests ──

    #[test]
    fn test_classify_with_rel_person_person() {
        let (s, d) = classify_with_relationship("Arjun", "Priya", "married_to");
        assert_eq!(s, "person");
        assert_eq!(d, "person");
    }

    #[test]
    fn test_classify_with_rel_person_place() {
        let (s, d) = classify_with_relationship("Priya", "Bangalore", "lives_in");
        assert_eq!(s, "person");
        assert_eq!(d, "place");
    }

    #[test]
    fn test_classify_with_rel_person_org() {
        let (s, d) = classify_with_relationship("Priya", "Flipkart", "works_at");
        assert_eq!(s, "person");
        assert_eq!(d, "organization");
    }

    #[test]
    fn test_classify_with_rel_tech_dst() {
        // "uses" implies dst is tech; FAISS is tech by name heuristic
        let (s, d) = classify_with_relationship("FAISS", "data pipeline", "uses");
        assert_eq!(s, "tech");
        assert_eq!(d, "tech");
    }

    #[test]
    fn test_classify_with_rel_built_with() {
        // "built_with" → src defaults to "project" if unknown, dst is tech
        let (s, d) = classify_with_relationship("MyApp", "React", "built_with");
        assert_eq!(s, "project");
        assert_eq!(d, "tech");
    }

    #[test]
    fn test_classify_with_rel_deployed_on() {
        let (s, d) = classify_with_relationship("MyApp", "AWS", "deployed_on");
        assert_eq!(s, "project");
        assert_eq!(d, "infrastructure");
    }

    #[test]
    fn test_classify_with_rel_works_on() {
        let (s, d) = classify_with_relationship("Pranab", "YantrikDB", "works_on");
        assert_eq!(s, "person");
        assert_eq!(d, "project");
    }

    #[test]
    fn test_classify_with_rel_fallback() {
        // Truly unknown relationship → falls back to name heuristics
        let (s, d) = classify_with_relationship("FAISS", "data pipeline", "related_to");
        assert_eq!(s, "tech");
        assert_eq!(d, "unknown");
    }
}

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

    #[test]
    fn code_identifiers_do_not_become_entities() {
        // The BEAM failure, reduced: a Flask route in a fenced block made
        // GET/POST/String entities, and the conflict detector then paired
        // unrelated chunks that merely both quoted a route.
        let text = "Alice deployed the service.\n\n```python\n\
                    @app.route('/login', methods=['GET', 'POST'])\n\
                    def login():\n    data = LoginSchema(String)\n```\n\
                    She reported it to Acme Corp.";
        let got = extract_heuristic_entities(text);
        for bad in ["GET", "POST", "String", "LoginSchema"] {
            assert!(
                !got.iter().any(|e| e.contains(bad)),
                "code identifier {bad:?} leaked into entities: {got:?}"
            );
        }
        // Prose entities on both sides of the block survive.
        assert!(
            got.iter().any(|e| e == "Alice"),
            "lost prose entity: {got:?}"
        );
        assert!(
            got.iter().any(|e| e.contains("Acme")),
            "lost prose entity after the block: {got:?}"
        );
    }

    #[test]
    fn inline_spans_are_stripped_without_welding_neighbours() {
        let got = extract_heuristic_entities("Bob set `MAX_RETRIES` Carol reviewed it");
        assert!(!got.iter().any(|e| e.contains("MAX_RETRIES")), "{got:?}");
        // The space substituted for the span must keep Bob and Carol apart
        // rather than producing a single "Bob Carol" chunk.
        assert!(got.iter().any(|e| e == "Bob"), "{got:?}");
        assert!(got.iter().any(|e| e == "Carol"), "{got:?}");
        assert!(!got.iter().any(|e| e == "Bob Carol"), "welded: {got:?}");
    }

    #[test]
    fn inline_span_before_a_fence_is_still_stripped() {
        // Regression: a fence-first branch emitted everything preceding the
        // fence verbatim, so an inline `GET` earlier in the same text
        // survived — the exact identifier this function exists to remove.
        let text = "`GET` Alice then
```python
class User: pass
```
done";
        let got = extract_heuristic_entities(text);
        assert!(
            !got.iter().any(|e| e.contains("GET")),
            "inline span leaked: {got:?}"
        );
        assert!(
            !got.iter().any(|e| e.contains("User")),
            "fence leaked: {got:?}"
        );
        assert!(got.iter().any(|e| e == "Alice"), "prose lost: {got:?}");
    }

    #[test]
    fn text_without_backticks_is_unchanged() {
        let plain = "Alice Chen is the CEO of Acme Corp";
        assert_eq!(
            extract_heuristic_entities(plain),
            extract_heuristic_entities_inner(plain),
            "no-backtick path must be byte-identical to the pre-change behavior"
        );
        assert!(matches!(strip_code(plain), std::borrow::Cow::Borrowed(_)));
    }

    #[test]
    fn unterminated_markers_do_not_drop_prose() {
        // A stray single backtick must not swallow the rest of the memory.
        let got = extract_heuristic_entities("Dave noted ` then Erin shipped it");
        assert!(
            got.iter().any(|e| e == "Erin"),
            "prose lost after stray tick: {got:?}"
        );
    }
}

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

    /// The bug, stated as a test.
    ///
    /// `ENTITY_STOPWORDS.contains(&tok)` is an exact match against the
    /// capitalized forms, so ALL-CAPS function words were never stripped.
    /// They then tokenize to lowercase and match every query containing that
    /// ordinary word, which is how `AT` became a graph anchor joining
    /// unrelated records.
    #[test]
    fn all_caps_function_words_are_not_entities() {
        for text in [
            "AT the meeting we shipped it",
            "THE release went out",
            "DID the migration finish",
            "NOT a real entity here",
        ] {
            let ents = extract_heuristic_entities(text);
            for bad in ["AT", "THE", "DID", "NOT"] {
                assert!(
                    !ents.iter().any(|e| e == bad),
                    "{bad:?} became an entity from {text:?} -> {ents:?}"
                );
            }
        }
    }

    /// Mixed case must not smuggle them either.
    #[test]
    fn mixed_case_function_words_are_not_entities() {
        let ents = extract_heuristic_entities("aT tHe meeting, dId anything ship");
        assert!(
            !ents.iter().any(|e| e.eq_ignore_ascii_case("at")
                || e.eq_ignore_ascii_case("the")
                || e.eq_ignore_ascii_case("did")),
            "mixed-case function word survived: {ents:?}"
        );
    }

    /// Words absent from the list in EVERY case, found the same way.
    #[test]
    fn newly_listed_function_words_are_not_entities() {
        let ents = extract_heuristic_entities("Most of it shipped. Not all. More later.");
        for bad in ["Most", "Not", "More"] {
            assert!(
                !ents.iter().any(|e| e == bad),
                "{bad:?} became an entity -> {ents:?}"
            );
        }
    }

    /// A bare month is in nearly every dated record, so as a node it links
    /// everything to everything. Observed as `graph-connected via June`.
    #[test]
    fn bare_month_names_are_not_entities() {
        let ents = extract_heuristic_entities("June was busy. We shipped in March.");
        for bad in ["June", "March"] {
            assert!(
                !ents.iter().any(|e| e == bad),
                "{bad:?} became an entity -> {ents:?}"
            );
        }
    }

    /// THE OTHER DIRECTION, which is what makes this a real gate rather than a
    /// blanket suppressor: real entities must still be extracted, including
    /// ones that merely CONTAIN a stopword, and ones that are legitimately
    /// capitalized after a stripped leading stopword.
    #[test]
    fn real_entities_still_extracted() {
        let ents = extract_heuristic_entities(
            "At Yantrik Systems we met Alice Chen about the Boston office.",
        );
        for good in ["Yantrik Systems", "Alice Chen", "Boston"] {
            assert!(
                ents.iter()
                    .any(|e| e.contains(good) || good.contains(e.as_str())),
                "real entity {good:?} was lost -> {ents:?}"
            );
        }
    }

    /// An all-caps token that is NOT a function word is still an entity —
    /// otherwise the fix would eat acronyms, which are exactly the kind of
    /// short high-signal name a memory system must keep.
    #[test]
    fn all_caps_acronyms_survive() {
        let ents = extract_heuristic_entities("The NASA contract and the HNSW index shipped.");
        assert!(
            ents.iter().any(|e| e.contains("NASA")),
            "NASA was stripped as if it were a function word -> {ents:?}"
        );
    }
}

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

    /// The census that motivated this: real entity names taken verbatim from a
    /// live store, every one of them a graph node joining unrelated records.
    #[test]
    fn all_caps_headings_are_not_entities() {
        for text in [
            "THINGS I MISSED THAT CODEX FOUND BY READING THE CODE follow.",
            "USER MUST UPDATE MCP CONFIG before restarting.",
            "REAL ESTATE TAX ANALYSIS was attached.",
            "HERMES REMOTE DESKTOP LIVE VERIFICATION PASSED today.",
        ] {
            for e in extract_heuristic_entities(text) {
                let caps = e
                    .split_whitespace()
                    .filter(|t| is_all_caps_token(t))
                    .count();
                assert!(
                    caps <= MAX_ALLCAPS_TOKENS,
                    "heading became entity {e:?} from {text:?}"
                );
            }
        }
    }

    /// A runaway mixed-case run is prose too.
    #[test]
    fn overlong_capitalized_runs_are_not_entities() {
        let ents =
            extract_heuristic_entities("Recall Return Unrelated Records Root Cause Found Today");
        assert!(
            ents.iter()
                .all(|e| e.split_whitespace().count() <= MAX_ENTITY_TOKENS),
            "overlong run survived -> {ents:?}"
        );
    }

    /// THE OTHER DIRECTION. Short acronyms and ordinary names are the whole
    /// point of the extractor and must be untouched.
    #[test]
    fn short_acronyms_and_names_survive() {
        let ents = extract_heuristic_entities(
            "NASA and IBM Watson met Alice Chen at Yantrik Systems in San Francisco.",
        );
        for good in [
            "NASA",
            "IBM Watson",
            "Alice Chen",
            "Yantrik Systems",
            "San Francisco",
        ] {
            assert!(
                ents.iter().any(|e| e.contains(good)),
                "real entity {good:?} lost -> {ents:?}"
            );
        }
    }

    /// Two all-caps tokens is a name, not a heading — the boundary must not
    /// slide down and start eating them.
    #[test]
    fn two_token_all_caps_names_survive() {
        let ents = extract_heuristic_entities("The NASA JPL team shipped it.");
        assert!(
            ents.iter().any(|e| e.contains("NASA JPL")),
            "two-token acronym name lost -> {ents:?}"
        );
    }
}

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

    #[test]
    fn possessives_are_canonicalized_before_becoming_entities() {
        let ents = extract_heuristic_entities(
            "Pranab's benchmark compared Reddit's API with Sol's Q2 plan.",
        );
        for canonical in ["Pranab", "Reddit", "Sol", "Q2"] {
            assert!(
                ents.iter().any(|e| e == canonical),
                "canonical {canonical:?} missing from {ents:?}"
            );
        }
        assert!(
            ents.iter()
                .all(|e| !e.ends_with("'s") && !e.ends_with('\'')),
            "possessive phantom survived: {ents:?}"
        );
    }

    #[test]
    fn apostrophes_inside_names_are_preserved() {
        let ents = extract_heuristic_entities("O'Brien met D'Arcy about O'Brien's release.");
        assert!(ents.iter().any(|e| e == "O'Brien"), "{ents:?}");
        assert!(ents.iter().any(|e| e == "D'Arcy"), "{ents:?}");
        assert!(!ents.iter().any(|e| e == "O'Brien's"), "{ents:?}");
    }

    #[test]
    fn capitalized_contractions_do_not_create_bare_phantoms() {
        let ents = extract_heuristic_entities("Let's begin. It's ready. What's next?");
        for bad in ["Let", "It", "What"] {
            assert!(!ents.iter().any(|e| e == bad), "{bad:?} survived: {ents:?}");
        }
    }
}