"""In-memory indexed access over the triple list.

This module provides O(1) lookup paths that replace the O(n) linear
scans over ``self._triples`` in ``KnowledgeGraph``. The index is
rebuilt from the canonical JSON on every ``build()`` call and
maintained incrementally as triples are added or mutated.

Design decisions:

  - **In-memory dicts, not SQLite tables.** At our current scale
    (~25K triples) and target ceiling (~1M triples), Python dicts
    with integer keys are faster than SQLite round-trips for the
    query patterns we actually use. The JSON files remain the
    canonical vault-readable source of truth (Rule 4). If we ever
    need cross-process concurrent reads (e.g. a second server
    instance), SQLite becomes worth it; until then, dicts are ~10x
    faster for single-process indexed lookups.

  - **Six indexes** covering the five hot-path query patterns
    identified by profiling the 35 linear scans in graph.py:

      1. ``by_subject[slug] -> list[int]`` — triple indices where
         the entity is the subject. Used by query_entity.
      2. ``by_object[slug] -> list[int]`` — triple indices where
         the entity is the object. Used by query_entity.
      3. ``by_layer[layer] -> set[int]`` — triple indices per
         system_layer. Used by triples_by_layer.
      4. ``by_episode[episode_id] -> list[int]`` — triple indices
         sharing an episode. Used by triples_by_episode.
      5. ``by_prediction[outcome] -> set[int]`` — triple indices
         per prediction_outcome. Used by list_pending.
      6. ``active_set -> set[int]`` — indices of triples with
         status="active". Used as a pre-filter on nearly every scan.

  - **Index values are list indices into ``self._triples``**, not
    triple IDs. This avoids a secondary ID->index lookup. When a
    triple is appended, its index is ``len(self._triples) - 1``.

  - **Rebuild is O(n).** Incremental maintenance is O(1) per
    add/mutate. The rebuild runs once on startup (during ``build()``)
    and is fast: ~50ms for 100K triples on a modern laptop.

  - **No locking.** KnowledgeGraph is single-writer (one process,
    one async task mutates at a time). The index is always consistent
    with the triple list because every mutation path (add_triple,
    status change, epistemic update) calls the appropriate index
    maintenance method before returning.

Usage from KnowledgeGraph:

    # During build():
    self._index = TripleIndex()
    self._index.rebuild(self._triples)

    # During add_triple():
    self._index.on_add(idx, triple)

    # During status/layer/prediction mutations:
    self._index.on_mutate(idx, old_triple_snapshot, new_triple)

    # Query:
    for idx in self._index.active_with_subject(slug):
        t = self._triples[idx]
        ...
"""

from __future__ import annotations

import logging
from collections import defaultdict

log = logging.getLogger(__name__)


class TripleIndex:
    """In-memory indexed access over the triple list.

    All index values are integer indices into the parent's ``_triples``
    list. The caller owns the list; we only store positions.
    """

    def __init__(self) -> None:
        # slug -> list of triple indices (preserves insertion order)
        self.by_subject: dict[str, list[int]] = defaultdict(list)
        self.by_object: dict[str, list[int]] = defaultdict(list)
        # layer -> set of triple indices
        self.by_layer: dict[str, set[int]] = defaultdict(set)
        # episode_id -> list of triple indices (order matters for chronological)
        self.by_episode: dict[str, list[int]] = defaultdict(list)
        # prediction_outcome -> set of triple indices
        self.by_prediction: dict[str, set[int]] = defaultdict(set)
        # set of indices where status == "active"
        self.active_set: set[int] = set()
        # Triple id -> index (for O(1) id lookups)
        self.id_to_idx: dict[int, int] = {}

    def rebuild(self, triples: list[dict]) -> None:
        """Full rebuild from the canonical triple list. O(n)."""
        self.by_subject.clear()
        self.by_object.clear()
        self.by_layer.clear()
        self.by_episode.clear()
        self.by_prediction.clear()
        self.active_set.clear()
        self.id_to_idx.clear()

        for idx, t in enumerate(triples):
            self._index_one(idx, t)

    def on_add(self, idx: int, triple: dict) -> None:
        """Incrementally index a newly-appended triple. O(1)."""
        self._index_one(idx, triple)

    def on_mutate(
        self,
        idx: int,
        old: dict,
        new: dict,
    ) -> None:
        """Update indexes after an in-place mutation on a triple.

        Only the fields that changed need re-indexing. The caller
        snapshots the old values before mutating, then calls this
        with both old and new dicts. We remove the old index entries
        and add the new ones. O(1) per changed field.
        """
        # Status changed
        old_status = old.get("status")
        new_status = new.get("status")
        if old_status != new_status:
            if old_status == "active":
                self.active_set.discard(idx)
            if new_status == "active":
                self.active_set.add(idx)

        # Layer changed
        old_layer = old.get("system_layer", "neocortical")
        new_layer = new.get("system_layer", "neocortical")
        if old_layer != new_layer:
            self.by_layer.get(old_layer, set()).discard(idx)
            self.by_layer[new_layer].add(idx)

        # Prediction outcome changed
        old_pred = old.get("prediction_outcome")
        new_pred = new.get("prediction_outcome")
        if old_pred != new_pred:
            if old_pred:
                self.by_prediction.get(old_pred, set()).discard(idx)
            if new_pred:
                self.by_prediction[new_pred].add(idx)

    # ── Query helpers ──────────────────────────────────────────────────

    def active_with_subject(self, slug: str) -> list[int]:
        """Return indices of active triples where subject == slug."""
        candidates = self.by_subject.get(slug)
        if not candidates:
            return []
        return [i for i in candidates if i in self.active_set]

    def active_with_object(self, slug: str) -> list[int]:
        """Return indices of active triples where object == slug."""
        candidates = self.by_object.get(slug)
        if not candidates:
            return []
        return [i for i in candidates if i in self.active_set]

    def active_in_layer(self, layer: str) -> list[int]:
        """Return indices of active triples in a system layer."""
        layer_set = self.by_layer.get(layer)
        if not layer_set:
            return []
        return [i for i in layer_set if i in self.active_set]

    def with_episode(self, episode_id: str) -> list[int]:
        """Return indices of triples sharing an episode_id."""
        return list(self.by_episode.get(episode_id, []))

    def with_prediction_outcome(self, outcome: str) -> list[int]:
        """Return indices of active triples with a given prediction_outcome."""
        candidates = self.by_prediction.get(outcome)
        if not candidates:
            return []
        return [i for i in candidates if i in self.active_set]

    def by_id(self, triple_id: int) -> int | None:
        """Return the list index for a given triple ID, or None."""
        return self.id_to_idx.get(triple_id)

    # ── Internal ───────────────────────────────────────────────────────

    def _index_one(self, idx: int, t: dict) -> None:
        """Index a single triple at position ``idx``."""
        subj = t.get("subject")
        if subj:
            self.by_subject[subj].append(idx)

        obj = t.get("object")
        if obj:
            self.by_object[obj].append(idx)

        layer = t.get("system_layer", "neocortical")
        self.by_layer[layer].add(idx)

        ep = t.get("episode_id")
        if ep:
            self.by_episode[ep].append(idx)

        pred_outcome = t.get("prediction_outcome")
        if pred_outcome:
            self.by_prediction[pred_outcome].add(idx)

        if t.get("status") == "active":
            self.active_set.add(idx)

        tid = t.get("id")
        if tid is not None:
            self.id_to_idx[tid] = idx
