"""Unified Entity Registry — single source of truth for all entities.

Every subsystem (memory graph, KB wiki, agent tools, vault search) references
entities through this registry. One slug, one canonical name, one set of aliases.

Storage: memory/graph/entities.json (vault-centric, same as before but now
the registry owns it exclusively).
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
import threading
from datetime import UTC, datetime
from difflib import SequenceMatcher
from typing import TYPE_CHECKING, Any

from agntspace.memory.ontology import EntityType, resolve_entity_type
from agntspace.text import sanitize_name as _sanitize_name
from agntspace.text import strip_comma_descriptor as _strip_comma_descriptor

if TYPE_CHECKING:
    from agntspace.vault.reader import VaultReader
    from agntspace.vault.writer import VaultWriter

log = logging.getLogger(__name__)

_ENTITIES_PATH = "memory/graph/entities.json"

# ── Resolution thresholds ──
_EMBEDDING_RESOLVE = 0.82       # Tier 3 cutoff (embedding cosine)
_SEQUENCE_RESOLVE = 0.92        # Tier 4 cutoff (SequenceMatcher, hardened 2026-04-17)


def _on_async_event_loop() -> bool:
    """True when the current thread is running an asyncio event loop.

    Calling a synchronous ``httpx.Client`` request (or any blocking I/O) from
    here would freeze every pending FastAPI request for the duration of the
    HTTP round-trip — 8 s on cold-start Ollama and 50 ms even on warm calls.
    With thousands of entities that's minutes of stacked blockage on the main
    loop. Callers on the loop MUST skip embedding and fall back to the string
    tiers; background work (graph build runs in a pool executor) still runs
    embedding normally because ``get_running_loop()`` raises there.
    """
    try:
        asyncio.get_running_loop()
        return True
    except RuntimeError:
        return False


def _slugify(name: str) -> str:
    """Convert entity name to a stable slug ID."""
    return re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")


def _now_iso() -> str:
    return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")


def _differs_only_in_digits(a: str, b: str) -> bool:
    """Return True if *a* and *b* differ only in their digit characters.

    "oldentity0" vs "oldentity1" → True (same letters, different digits).
    "sarah chen" vs "sarah c."   → False (letters differ).
    """
    alpha_a = re.sub(r"\d", "", a)
    alpha_b = re.sub(r"\d", "", b)
    return alpha_a == alpha_b and a != b


class EntityRegistry:
    """Canonical registry for all entities in the system.

    Every entity has:
    - slug: stable ID derived from name
    - canonical_name: display name
    - entity_type: from ontology EntityType
    - aliases: alternative names
    - authority: explicit > system > inferred (never downgrades)
    - decay_condition: when to forget
    - access tracking: first_seen, last_accessed, access_count
    - status: active or decayed
    """

    def __init__(
        self,
        vault: VaultReader | None = None,
        writer: VaultWriter | None = None,
    ) -> None:
        self._vault = vault
        self._writer = writer
        self._entities: dict[str, dict[str, Any]] = {}
        self._dirty = False

        # Indexes for fast lookup
        self._name_index: dict[str, str] = {}   # lowercase name -> slug
        self._alias_index: dict[str, str] = {}   # lowercase alias -> slug

        # Reentrant lock guards every read and write of the shared state
        # (_entities / _name_index / _alias_index / _embedding_cache / _dirty).
        # The knowledge-graph build runs in a ThreadPoolExecutor while KB
        # ingest + chat entity extraction call add()/resolve() from the main
        # event loop and other worker tasks; without this lock iteration and
        # mutation race, producing "dictionary changed size during iteration"
        # that kills the background graph build permanently. RLock (not Lock)
        # because add() calls resolve() and merge() calls _rebuild_indexes.
        self._lock = threading.RLock()

        # Embedding-based dedup (Tier 4 resolution).
        # Injected post-hoc: ``registry._embedder = vault_embeddings._get_model()``
        # When set, resolve() uses cosine similarity as a final fallback
        # after string-based tiers fail. This catches paraphrases,
        # abbreviations, and name permutations that SequenceMatcher misses.
        self._embedder: object | None = None
        self._embedding_cache: dict[str, Any] = {}  # slug -> embedding vector
        # Circuit breaker: when embedding calls fail (Ollama slow/unreachable),
        # skip the embedding tier for the rest of the process lifetime rather
        # than blocking every entity resolution on 40-second cold-start calls.
        # Reset only by service restart — deliberately NOT time-based because
        # a thrashing cold-start that recovers on retry N is still worse than
        # deterministic FTS-only behavior. The user can re-enable by restarting
        # after warming up the embedding model.
        self._embedding_disabled: bool = False
        # Tier 4 (SequenceMatcher) disable flag — mirrors _embedding_disabled.
        # Set by callers that resolve thousands of names in a tight loop where
        # Tier 4's O(entities × name_length²) cost would saturate CPU for
        # minutes producing no progress. The historical pathology: incremental
        # graph build's _project_wikilink_index calls resolve() ~50 × 2600
        # = 132K times against a 7000-entity registry; each Tier 4 call does
        # 14K+ SequenceMatcher.ratio() comparisons; CPU pinned at 100%, no
        # progress, no log lines, indistinguishable from a deadlock.
        # Tier 1 (exact slug) and Tier 2 (alias index) are O(1) and remain
        # active — they handle every wikilink that was successfully resolved
        # in the original scan, which is all the projection cares about.
        self._fuzzy_disabled: bool = False
        # Cap the per-call cache-refresh batch size. At 3766 entities a full
        # refresh hits Ollama with one giant batch that takes minutes to return
        # cold. Capping at N means first resolve embeds N entities, next N
        # entities get embedded on next resolve, etc — no single call ever
        # blocks the event loop for more than ~1× single-embed latency.
        self._embedding_batch_cap: int = 50

    def load(self) -> int:
        """Load entities from vault JSON. Returns count loaded."""
        if not self._vault:
            return 0

        if not self._vault.exists(_ENTITIES_PATH):
            return 0

        try:
            path = self._vault.paths.resolve(_ENTITIES_PATH) if self._vault.paths else self._vault._vault_dir / _ENTITIES_PATH
            raw = path.read_text(encoding="utf-8")
            data = json.loads(raw)
            with self._lock:
                self._entities = data
                self._rebuild_indexes()
                count = len(self._entities)
            log.debug("EntityRegistry: loaded %d entities", count)
            return count
        except Exception as e:
            log.warning("EntityRegistry: failed to load entities.json: %s", e)
            with self._lock:
                self._entities = {}
            return 0

    def save(self) -> None:
        """Persist entities to vault JSON."""
        if not self._writer:
            return
        # Serialize under lock to prevent iteration during mutation; then
        # release the lock before doing vault I/O so a slow write doesn't
        # block resolve() / add() for seconds.
        with self._lock:
            if not self._dirty:
                return
            entities_json = json.dumps(self._entities, separators=(",", ":"), ensure_ascii=False)
            self._dirty = False
            count = len(self._entities)
        # Internal bookkeeping write — the graph state is a side effect of
        # higher-level decisions (chat turn entity extraction, KB ingest,
        # dream phases, memory consolidation). Those callers gate their
        # own contract checks; here we declare trusted persistence so the
        # VaultWriter runtime guard passes regardless of caller scope.
        self._writer.write(_ENTITIES_PATH, entities_json, trust_reason="graph_persist")
        log.info("EntityRegistry: saved %d entities", count)

    def _rebuild_indexes(self) -> None:
        """Rebuild name and alias indexes from entity data.

        Caller must hold self._lock (internal helper — already invoked from
        load() / merge() / add() which take the lock themselves).
        """
        self._name_index.clear()
        self._alias_index.clear()
        for slug, ent in self._entities.items():
            name_lower = ent["canonical_name"].lower()
            self._name_index[name_lower] = slug
            for alias in ent.get("aliases", []):
                self._alias_index[alias.lower()] = slug

    # ── Resolution ──

    def resolve(self, name: str) -> str | None:
        """Resolve an entity name to its slug. Four-tier matching.

        1. Exact slug or canonical name (case-insensitive) — O(1)
        2. Alias match (via index) — O(1)
        3. Embedding similarity (cosine ≥ 0.82) — semantic match that
           catches abbreviations ("J. Smith" → "John Smith"), name
           permutations ("Chen, Sarah" → "Sarah Chen"), and paraphrases.
           Primary fuzzy tier when embedder is available.
        4. SequenceMatcher (≥ 0.92) — character-level fallback when
           embedder is unavailable (test fixtures, Ollama offline,
           embedding model not pulled). Hardened 2026-04-17 to 0.92
           to avoid false-positive merges between distinct entities.
        """
        name = _sanitize_name(name)
        if not name:
            return None
        name_lower = name.lower().strip()
        slug = _slugify(name)

        with self._lock:
            # Tier 1: exact slug
            if slug in self._entities:
                return slug

            # Tier 1b: exact canonical name (via index)
            if name_lower in self._name_index:
                return self._name_index[name_lower]

            # Tier 2: alias (via index)
            if name_lower in self._alias_index:
                return self._alias_index[name_lower]

        # Tier 3: embedding similarity (primary fuzzy tier). Released the
        # lock before the LLM call so a slow Ollama embed doesn't block
        # other resolves — _resolve_by_embedding reacquires it around
        # the cache iteration itself.
        if len(name) > 3 and self._embedder is not None:
            try:
                match = self._resolve_by_embedding(name)
                if match is not None:
                    return match
            except Exception:
                log.debug("Embedding resolution failed", exc_info=True)

        # Tier 4: SequenceMatcher fallback. Under the lock to prevent
        # mid-iteration mutation from the graph-build thread. Skipped
        # when ``_fuzzy_disabled`` is set — see flag docstring above
        # for the runaway-loop pathology this guards against.
        if len(name) > 5 and not self._fuzzy_disabled:
            with self._lock:
                best_score = 0.0
                best_slug = None
                best_candidate = ""
                for s, ent in self._entities.items():
                    canon = ent["canonical_name"].lower()
                    score = SequenceMatcher(None, name_lower, canon).ratio()
                    if score > best_score:
                        best_score = score
                        best_slug = s
                        best_candidate = canon
                    for alias in ent.get("aliases", []):
                        al = alias.lower()
                        score = SequenceMatcher(None, name_lower, al).ratio()
                        if score > best_score:
                            best_score = score
                            best_slug = s
                            best_candidate = al
                if best_score >= _SEQUENCE_RESOLVE and best_slug:
                    # Guard: names differing only in digits are distinct
                    # entities (Entity0 vs Entity1, Version2 vs Version3).
                    if not _differs_only_in_digits(name_lower, best_candidate):
                        return best_slug

        return None

    def _resolve_by_embedding(
        self, name: str, min_score: float = _EMBEDDING_RESOLVE,
    ) -> str | None:
        """Resolve via cosine similarity on entity name embeddings.

        Embeds the query name and compares against cached embeddings of
        all canonical names + aliases. Returns the slug of the best
        match above *min_score*, or None.

        Concurrency-safe: snapshots the cache before iterating so parallel
        writes from the graph-build thread don't raise
        "dictionary changed size during iteration". Circuit breaker
        (``_embedding_disabled``) trips on first embed failure to avoid
        40-second cold-start stalls on every resolve.
        """
        import math

        model = self._embedder
        with self._lock:
            if model is None or not self._entities or self._embedding_disabled:
                return None

        # Event-loop safety: refuse to run a blocking HTTP embed call from the
        # main asyncio loop. The circuit breaker is NOT tripped — this is a
        # "skip this tier for this caller" decision, not a permanent embed
        # failure. Graph-build callers run in a pool executor and pass this
        # check; chat / ingest / API callers on the main loop fall back to
        # string tiers (which are fast and pure Python).
        if _on_async_event_loop():
            return None

        try:
            query_emb = [float(x) for x in list(model.embed([name]))[0]]
        except Exception:
            # First embed failure trips the breaker — further calls would
            # time out identically and block every resolve for ~40s each.
            with self._lock:
                self._embedding_disabled = True
            log.info(
                "EntityRegistry: embedding tier disabled for this process "
                "(query embed failed — Ollama not reachable or model not pulled). "
                "Resolution falls back to string tiers.",
            )
            return None
        query_sq = sum(x * x for x in query_emb)
        if query_sq < 1e-18:
            return None
        query_norm = math.sqrt(query_sq)

        self._refresh_embedding_cache()

        best_score = 0.0
        best_slug = None
        # Hold the lock across the cache iteration. Pure arithmetic, so
        # contention is minimal; prevents the build thread from mutating
        # _embedding_cache mid-scan (which crashed resolve today at 17:21).
        with self._lock:
            for slug, emb in self._embedding_cache.items():
                emb_list = list(emb) if not isinstance(emb, list) else emb
                if len(emb_list) != len(query_emb):
                    continue
                emb_sq = 0.0
                dot = 0.0
                for a, b in zip(query_emb, emb_list, strict=False):
                    dot += a * float(b)
                    emb_sq += float(b) * float(b)
                if emb_sq < 1e-18:
                    continue
                score = dot / (query_norm * math.sqrt(emb_sq))
                if score > best_score:
                    best_score = score
                    best_slug = slug

        if best_score >= min_score and best_slug:
            log.debug(
                "Embedding resolved '%s' -> '%s' (score=%.3f)",
                name, best_slug, best_score,
            )
            return best_slug
        return None

    def _refresh_embedding_cache(self) -> None:
        """Rebuild embedding cache for entities not yet cached.

        Capped at ``_embedding_batch_cap`` per call. With thousands of
        entities and a cold-start Ollama model, an uncapped batch would
        block resolution for minutes on first call. Capped means we amortize
        the cache fill across many resolves — each call embeds a small batch
        while still making progress. Ambiguous placeholders are excluded
        (they are resolution outputs, never targets).
        """
        model = self._embedder
        with self._lock:
            if model is None or self._embedding_disabled:
                return

        # Event-loop safety: same rule as _resolve_by_embedding. Skipping is
        # safe — the cache fills opportunistically from off-loop callers
        # (graph build, explicit warm-up) and resolves on-loop gracefully
        # miss the embedding tier rather than blocking the loop for seconds.
        if _on_async_event_loop():
            return

        with self._lock:
            # Collect texts to embed (canonical name + aliases, one per slug).
            # Under the lock so we can iterate self._entities + self._embedding_cache
            # without racing the graph-build thread.
            to_embed: list[tuple[str, str]] = []  # (slug, text)
            for slug, ent in self._entities.items():
                if slug in self._embedding_cache:
                    continue
                # Combine canonical name + aliases into one embedding text
                parts = [ent["canonical_name"]]
                parts.extend(ent.get("aliases", [])[:3])  # cap aliases
                to_embed.append((slug, " | ".join(parts)))
                if len(to_embed) >= self._embedding_batch_cap:
                    break

        if not to_embed:
            return

        texts = [t[1] for t in to_embed]
        # Embed outside the lock — a slow Ollama call must not block
        # every resolver behind us for its full duration.
        try:
            embeddings = list(model.embed(texts))
        except Exception:
            # Trip the circuit breaker on the first failure so later calls
            # don't retry. Logged at INFO so the operator sees it without
            # exception spam.
            with self._lock:
                self._embedding_disabled = True
            log.info(
                "EntityRegistry: embedding cache refresh failed "
                "(Ollama not reachable or model not pulled). "
                "Falling back to string resolution tiers.",
            )
            return

        with self._lock:
            for (slug, _), emb in zip(to_embed, embeddings, strict=False):
                self._embedding_cache[slug] = emb

    def invalidate_embedding_cache(self, slug: str | None = None) -> None:
        """Clear cached embeddings. Called after entity add/merge/delete."""
        with self._lock:
            if slug:
                self._embedding_cache.pop(slug, None)
            else:
                self._embedding_cache.clear()

    def backfill_embeddings_blocking(
        self,
        *,
        batch_size: int = 50,
        max_entities: int | None = None,
    ) -> dict:
        """Embed every registered entity that isn't yet in the cache.

        Blocking — caller MUST run via ``asyncio.to_thread`` or in a worker
        thread, NEVER on the asyncio main loop. The event-loop guard in
        ``_refresh_embedding_cache`` is bypassed here on purpose: this is
        an explicit user-triggered fill, not an opportunistic resolver
        side-effect.

        Resets the circuit breaker (``_embedding_disabled``) at the start
        — the user may have just installed Ollama or pulled the model and
        wants a fresh attempt. If the embed call fails again, the breaker
        trips as normal and the loop exits.

        Side-effect-free w.r.t. epistemic semantics: embeddings are a
        similarity lookup tool. No triple is added, removed, repointed,
        or reweighted. No status, authority, decay, or cascade is touched.
        Worst case (all batches fail): the cache is no fuller than before.

        Args:
            batch_size: number of entities to embed per round trip. Bounded
                so that a single failed call doesn't waste a huge batch.
            max_entities: hard cap on total entities scanned in this call.
                ``None`` = no cap, fill the whole registry. Useful for
                throttled / cancellable progressive runs.

        Returns:
            dict with: ``scanned`` (entities pulled into batches),
            ``embedded`` (entities successfully embedded),
            ``failed`` (entities in the failing batch when the breaker
            tripped — they're NOT cached), ``total_entities`` (registry
            size at exit), ``cached_after`` (cache size at exit),
            ``embedder_disabled`` (final state of the circuit breaker),
            ``reason`` (only on early-out: ``"no_embedder"``).
        """
        model = self._embedder
        if model is None:
            return {
                "scanned": 0,
                "embedded": 0,
                "failed": 0,
                "total_entities": len(self._entities),
                "cached_after": len(self._embedding_cache),
                "embedder_disabled": False,
                "reason": "no_embedder",
            }

        # Reset breaker — explicit user-triggered retry. If embeds fail
        # again, the loop trips it back on the first batch failure.
        with self._lock:
            self._embedding_disabled = False

        embedded = 0
        failed = 0
        scanned = 0
        if batch_size < 1:
            batch_size = 1

        while True:
            with self._lock:
                if max_entities is not None and scanned >= max_entities:
                    break
                # Collect next batch — entities not yet cached.
                to_embed: list[tuple[str, str]] = []
                for slug, ent in self._entities.items():
                    if slug in self._embedding_cache:
                        continue
                    parts = [ent["canonical_name"]]
                    parts.extend(ent.get("aliases", [])[:3])
                    to_embed.append((slug, " | ".join(parts)))
                    if len(to_embed) >= batch_size:
                        break

            if not to_embed:
                break

            # Respect max_entities even mid-batch.
            if max_entities is not None:
                remaining = max_entities - scanned
                if remaining < len(to_embed):
                    to_embed = to_embed[:remaining]
                if not to_embed:
                    break

            scanned += len(to_embed)
            texts = [t[1] for t in to_embed]
            try:
                # Embed outside the lock — Ollama latency is real.
                embeddings = list(model.embed(texts))
            except Exception as exc:
                # Trip the breaker so subsequent resolvers don't all stall
                # for the same reason.
                log.info(
                    "EntityRegistry: backfill embed batch failed: %s", exc,
                )
                failed += len(to_embed)
                with self._lock:
                    self._embedding_disabled = True
                break

            with self._lock:
                for (slug, _), emb in zip(to_embed, embeddings, strict=False):
                    self._embedding_cache[slug] = emb
                    embedded += 1

        return {
            "scanned": scanned,
            "embedded": embedded,
            "failed": failed,
            "total_entities": len(self._entities),
            "cached_after": len(self._embedding_cache),
            "embedder_disabled": self._embedding_disabled,
        }

    # ── CRUD ──

    def add(
        self,
        name: str,
        entity_type: str | EntityType = EntityType.THING,
        authority: str = "inferred",
        aliases: list[str] | None = None,
        decay_condition: str = "when:unaccessed:90d",
    ) -> str:
        """Add or update an entity. Returns the slug.

        If entity exists: merges aliases, upgrades authority (never downgrades),
        updates entity_type if more specific.
        """
        # Sanitize name first — reject embedded newlines, markdown markers,
        # and control-only junk. This is the single chokepoint for every
        # entity write; callers that pass garbage get an empty slug back
        # and nothing is stored.
        name = _sanitize_name(name)
        if not name:
            return ""

        # Comma-descriptor strip: catches the LLM extraction error where
        # a role/description gets glued to a proper-noun head with a
        # comma — "Erasmo Zuleta, gobernador de Córdoba" gets stripped
        # to "Erasmo Zuleta" while the original full string is seeded
        # as an alias so prior queries still resolve. Conservative — only
        # fires when the post-comma segment starts with a lowercase
        # letter AND neither part contains another comma. Preserves
        # legitimate names like "University of California, Berkeley" and
        # "Smith, John". This is the deterministic safety net behind the
        # skill-driven Name Purity rule (kb-ingest, memory-consolidate,
        # memory-dream-light/deep/rem). See ``strip_comma_descriptor``.
        clean_name, original_full = _strip_comma_descriptor(name)
        comma_descriptor_alias = ""
        if original_full and clean_name and clean_name != name:
            log.debug(
                "Entity name normalized: '%s' -> '%s' (comma-descriptor stripped; "
                "original preserved as alias)",
                original_full, clean_name,
            )
            name = clean_name
            comma_descriptor_alias = original_full

        # Cross-language canonicalization (spec v1.1 item 6):
        # If the caller passes a non-English form of a known canonical entity
        # (e.g. "Стокгольм", "北京"), swap to the English canonical name so
        # the slug stays ASCII-safe, and remember the foreign form as an
        # alias to seed below.
        _was_non_canonical_alias = False
        try:
            from agntspace.infra.language import canonical_from_alias
            canonical = canonical_from_alias(name)
            if canonical and canonical != name:
                name = canonical
                _was_non_canonical_alias = True
        except Exception:
            pass

        # Normalize entity_type
        if isinstance(entity_type, str):
            etype = resolve_entity_type(entity_type)
        else:
            etype = entity_type

        existing = self.resolve(name)
        now = _now_iso()

        # If the comma-descriptor strip fired AND the canonical short form
        # already exists in the registry, fold the original full string
        # into the existing entity's aliases instead of dropping it.
        merge_aliases = list(aliases or [])
        if comma_descriptor_alias and comma_descriptor_alias not in merge_aliases:
            merge_aliases.append(comma_descriptor_alias)

        # RLock is reentrant — resolve() / _merge_existing() can safely
        # reacquire. Holding the lock across the whole block ensures the
        # slug-collision counter below doesn't race with a concurrent add().
        with self._lock:
            if existing:
                return self._merge_existing(existing, etype, authority, merge_aliases, now)

            # New entity
            slug = _slugify(name)
            base_slug = slug
            counter = 2
            while slug in self._entities:
                slug = f"{base_slug}-{counter}"
                counter += 1

            # Cross-language alias seeding (spec §19 item 6):
            # If `name` is a known canonical entity (or a known alias), auto-add
            # all cross-language forms so "Stockholm" / "Стокгольм" / "斯德哥尔摩"
            # all resolve to the same slug.
            seeded_aliases: list[str] = list(aliases or [])
            # Preserve the pre-strip full string as an alias so prior
            # queries / extractions that used the long form still resolve
            # to this slug. Insert ahead of cross-language forms so the
            # original spelling stays first in the alias list.
            if comma_descriptor_alias and comma_descriptor_alias not in seeded_aliases:
                seeded_aliases.append(comma_descriptor_alias)
            try:
                from agntspace.infra.language import (
                    canonical_from_alias,
                    cross_language_aliases,
                )
                canonical = canonical_from_alias(name) or name
                extra = cross_language_aliases(canonical)
                # If the user passed a non-English form, also add the canonical
                # English form as an alias so lookup works from any direction.
                if canonical != name and canonical not in seeded_aliases:
                    seeded_aliases.append(canonical)
                for form in extra:
                    if form and form != name and form not in seeded_aliases:
                        seeded_aliases.append(form)
            except Exception:
                pass

            self._entities[slug] = {
                "canonical_name": name,
                "entity_type": etype.value,
                "aliases": seeded_aliases,
                "authority": authority,
                "decay_condition": decay_condition,
                "first_seen": now,
                "last_accessed": now,
                "access_count": 0,
                "status": "active",
            }

            # Update indexes
            self._name_index[name.lower()] = slug
            for alias in seeded_aliases:
                self._alias_index[alias.lower()] = slug

            self._dirty = True
            self.invalidate_embedding_cache(slug)
            return slug

    def _merge_existing(
        self,
        slug: str,
        entity_type: EntityType,
        authority: str,
        aliases: list[str] | None,
        now: str,
    ) -> str:
        """Merge new data into an existing entity.

        Caller must hold ``self._lock`` — invoked from ``add()`` which already
        holds it. RLock is reentrant so nested acquisition is safe; we declare
        the invariant so future refactors don't accidentally bypass the lock.
        """
        ent = self._entities[slug]

        # Merge aliases
        if aliases:
            current_aliases = set(ent.get("aliases", []))
            new_aliases = set(aliases) - current_aliases
            if new_aliases:
                current_aliases.update(new_aliases)
                ent["aliases"] = sorted(current_aliases)
                for alias in new_aliases:
                    self._alias_index[alias.lower()] = slug

        # Upgrade authority (never downgrade)
        auth_rank = {"inferred": 0, "system": 1, "explicit": 2}
        if auth_rank.get(authority, 0) > auth_rank.get(ent.get("authority", "inferred"), 0):
            ent["authority"] = authority

        # Upgrade entity_type if more specific (THING -> specific type)
        current_type = resolve_entity_type(ent.get("entity_type", "thing"))
        if current_type == EntityType.THING and entity_type != EntityType.THING:
            ent["entity_type"] = entity_type.value

        # Track re-encounters separately from query access.
        # Entity dedup (re-adding same entity) should NOT bump last_accessed —
        # that defeats unaccessed decay. last_accessed is only set by
        # bump_access() which is called from actual query paths.
        ent["last_encountered"] = now
        self._dirty = True
        self.invalidate_embedding_cache(slug)
        return slug

    def get(self, slug: str) -> dict | None:
        """Get entity data by slug."""
        with self._lock:
            return self._entities.get(slug)

    def get_name(self, slug: str | None) -> str:
        """Get display name for a slug. Returns '?' if not found."""
        if not slug:
            return "?"
        with self._lock:
            ent = self._entities.get(slug)
            return ent["canonical_name"] if ent else slug

    def get_type(self, slug: str) -> EntityType:
        """Get the EntityType for a slug."""
        with self._lock:
            ent = self._entities.get(slug)
            if not ent:
                return EntityType.THING
            return resolve_entity_type(ent.get("entity_type", "thing"))

    def bump_access(self, slug: str) -> None:
        """Record that an entity was accessed (for relevance decay tracking)."""
        with self._lock:
            ent = self._entities.get(slug)
            if ent:
                ent["last_accessed"] = _now_iso()
                ent["access_count"] = ent.get("access_count", 0) + 1
                self._dirty = True

    def mark_verified(self, slug: str) -> None:
        """Record that an entity was epistemically verified (not just accessed).

        Verification events: user approval, authority upgrade, triple re-confirmation.
        This is distinct from bump_access which tracks reads.
        """
        with self._lock:
            ent = self._entities.get(slug)
            if ent:
                ent["last_verified"] = _now_iso()
                self._dirty = True

    def add_alias(self, slug: str, alias: str) -> bool:
        """Add an alias to an existing entity."""
        with self._lock:
            ent = self._entities.get(slug)
            if not ent:
                return False
            aliases = set(ent.get("aliases", []))
            if alias not in aliases:
                aliases.add(alias)
                ent["aliases"] = sorted(aliases)
                self._alias_index[alias.lower()] = slug
                self._dirty = True
            return True

    def restore(self, slug: str, entity_data: dict) -> bool:
        """Resurrect an entity at the given slug from a snapshot.

        Used by ``KnowledgeGraph.unmerge_entities`` to reverse a destructive
        merge. ``entity_data`` is typically the full dict captured from
        ``self.get(slug_remove)`` BEFORE ``merge()`` destroyed the row, so
        restoration is lossless on every field (canonical_name, aliases,
        authority, first_seen, access_count, status, decay_condition,
        last_verified, etc.).

        Returns True on success. Returns False (idempotent no-op) when
        the slug is already occupied by an entity with the same
        canonical_name — the caller can treat this as already-restored.
        Returns False when the slug is occupied by a DIFFERENT entity
        — the caller surfaces that conflict.
        """
        with self._lock:
            existing = self._entities.get(slug)
            canonical = (entity_data or {}).get("canonical_name", "")
            if not canonical:
                return False
            if existing:
                # Idempotent restore: same canonical name = treat as success
                return existing.get("canonical_name") == canonical

            ent = dict(entity_data)
            # Default any missing fields so the row is well-formed even when
            # the snapshot came from an older schema (forward compatibility).
            ent.setdefault("aliases", [])
            ent.setdefault("authority", "system")
            ent.setdefault("status", "active")
            ent.setdefault("first_seen", _now_iso())
            ent.setdefault("last_accessed", _now_iso())
            ent.setdefault("access_count", 0)

            self._entities[slug] = ent

            # Re-establish index entries
            self._name_index[canonical.lower()] = slug
            for alias in ent.get("aliases", []) or []:
                if isinstance(alias, str) and alias:
                    self._alias_index[alias.lower()] = slug

            self._dirty = True
            self.invalidate_embedding_cache(slug)
            log.info("Entity restored at slug='%s' (canonical='%s')", slug, canonical)
            return True

    def merge(self, slug_keep: str, slug_remove: str) -> bool:
        """Merge two entities. Keeps slug_keep, removes slug_remove.

        Aliases from removed entity are added to kept entity.
        Returns True if merge happened.
        """
        with self._lock:
            keep = self._entities.get(slug_keep)
            remove = self._entities.get(slug_remove)
            if not keep or not remove:
                return False

            # Merge aliases (include the removed entity's canonical name)
            keep_aliases = set(keep.get("aliases", []))
            keep_aliases.add(remove["canonical_name"])
            keep_aliases.update(remove.get("aliases", []))
            keep["aliases"] = sorted(keep_aliases)

            # Upgrade authority
            auth_rank = {"inferred": 0, "system": 1, "explicit": 2}
            if auth_rank.get(remove.get("authority", "inferred"), 0) > auth_rank.get(keep.get("authority", "inferred"), 0):
                keep["authority"] = remove["authority"]

            # Keep earliest first_seen
            if remove.get("first_seen", "z") < keep.get("first_seen", "z"):
                keep["first_seen"] = remove["first_seen"]

            # Sum access counts
            keep["access_count"] = keep.get("access_count", 0) + remove.get("access_count", 0)

            # Remove the merged entity
            del self._entities[slug_remove]

            # Rebuild indexes
            self._rebuild_indexes()
            self._dirty = True

            log.info("Merged entity '%s' into '%s'", slug_remove, slug_keep)
            return True

    def get_slug_redirect(self, old_slug: str) -> str | None:
        """Check if a slug was merged into another entity.

        After merge, the old slug's canonical name becomes an alias of the kept entity.
        This allows callers to find the new slug for a merged entity.
        """
        with self._lock:
            # Check alias index — merged entity's name is stored as alias
            for alias_lower, slug in self._alias_index.items():
                if _slugify(alias_lower) == old_slug:
                    return slug
            return None

    def upgrade_authority(self, slug: str, new_authority: str) -> bool:
        """Upgrade authority for an entity (never downgrades)."""
        with self._lock:
            ent = self._entities.get(slug)
            if not ent:
                return False
            auth_rank = {"inferred": 0, "system": 1, "explicit": 2}
            if auth_rank.get(new_authority, 0) > auth_rank.get(ent.get("authority", "inferred"), 0):
                ent["authority"] = new_authority
                self._dirty = True
                return True
            return False

    # ── Decay ──

    def mark_decayed(self, slug: str) -> None:
        """Mark an entity as decayed."""
        with self._lock:
            ent = self._entities.get(slug)
            if ent and ent.get("status") != "decayed":
                ent["status"] = "decayed"
                ent["decayed_at"] = _now_iso()
                self._dirty = True

    # ── Search ──

    def search(self, query: str, limit: int = 20) -> list[dict]:
        """Search entities by name/alias. Returns list of matches."""
        query_lower = query.lower()
        results = []

        with self._lock:
            for slug, ent in self._entities.items():
                if ent.get("status") == "decayed":
                    continue
                name = ent["canonical_name"]
                match = query_lower in name.lower()
                if not match:
                    match = any(query_lower in a.lower() for a in ent.get("aliases", []))
                if match:
                    results.append({
                        "slug": slug,
                        "name": name,
                        "entity_type": ent.get("entity_type", "thing"),
                        "authority": ent.get("authority", "inferred"),
                        "aliases": ent.get("aliases", []),
                        "access_count": ent.get("access_count", 0),
                    })

        _auth_rank = {"explicit": 2, "system": 1, "inferred": 0}
        results.sort(key=lambda x: (
            _auth_rank.get(x.get("authority", "inferred"), 0),
            x.get("access_count", 0),
        ), reverse=True)
        return results[:limit]

    def list_all(self, include_decayed: bool = False) -> list[dict]:
        """List all entities."""
        results = []
        with self._lock:
            for slug, ent in self._entities.items():
                if not include_decayed and ent.get("status") == "decayed":
                    continue
                results.append({"slug": slug, **ent})
        return results

    # ── Stats ──

    def stats(self) -> dict:
        """Entity statistics."""
        with self._lock:
            active = [e for e in self._entities.values() if e.get("status", "active") == "active"]
            decayed = [e for e in self._entities.values() if e.get("status") == "decayed"]
            total = len(self._entities)

        type_counts: dict[str, int] = {}
        for ent in active:
            etype = ent.get("entity_type", "thing")
            type_counts[etype] = type_counts.get(etype, 0) + 1

        authority_counts: dict[str, int] = {}
        for ent in active:
            auth = ent.get("authority", "inferred")
            authority_counts[auth] = authority_counts.get(auth, 0) + 1

        return {
            "total": total,
            "active": len(active),
            "decayed": len(decayed),
            "by_type": type_counts,
            "by_authority": authority_counts,
        }

    @property
    def entities(self) -> dict[str, dict[str, Any]]:
        """Snapshot view of entity data (for graph building).

        Returns a SHALLOW COPY of the outer dict so callers (graph.build,
        KBService, query_entity, etc.) can iterate .items() safely even
        while other threads are adding/merging entities. Inner per-entity
        dicts remain shared references — existing read-only access patterns
        (``registry.entities[slug].get("entity_type")``) still work; callers
        that mutate would need to go through registry methods anyway.
        """
        with self._lock:
            return dict(self._entities)

    @property
    def dirty(self) -> bool:
        with self._lock:
            return self._dirty

    @dirty.setter
    def dirty(self, value: bool) -> None:
        with self._lock:
            self._dirty = value
