"""
Entity Index for Hybrid Retrieval.

Maps discrete entities (nodes in the graph) to the episodes that contain them.
This allows for O(1) retrieval of candidate episodes given a query entity,
bypassing vector similarity bottlenecks.
"""

from __future__ import annotations

from collections import defaultdict
from typing import Iterable

from hieromem_core.episode import Episode


class EntityIndex:
    """In-memory inverted index mapping entities to episode IDs."""

    def __init__(self) -> None:
        # Map: entity_name -> set[episode_id]
        self._index: dict[str, set[str]] = defaultdict(set)
        # Map: episode_id -> set[entity_name] (for fast removal)
        self._reverse_index: dict[str, set[str]] = defaultdict(set)

    def add(self, episode: Episode) -> None:
        """Index all entities found in the episode's graph."""
        if not episode.G or episode.G.number_of_nodes() == 0:
            return

        graph = episode.G
        episode_id = episode.id
        entities = set()

        for node in graph.nodes():
            entity = str(node).strip()
            if entity:
                self._index[entity].add(episode_id)
                entities.add(entity)

        if entities:
            self._reverse_index[episode_id].update(entities)

    def remove(self, episode_id: str) -> None:
        """Remove an episode from the index."""
        if episode_id not in self._reverse_index:
            return

        entities = self._reverse_index.pop(episode_id)
        for entity in entities:
            if entity in self._index:
                self._index[entity].discard(episode_id)
                if not self._index[entity]:
                    del self._index[entity]

    def lookup(self, entity: str) -> set[str]:
        """Return IDs of episodes containing the exact entity."""
        return self._index.get(str(entity), set())

    def lookup_many(self, entities: Iterable[str]) -> set[str]:
        """Return union of episode IDs containing any of the entities."""
        result = set()
        for entity in entities:
            result.update(self.lookup(entity))
        return result

    def clear(self) -> None:
        """Clear the index."""
        self._index.clear()
        self._reverse_index.clear()
