"""
simulation.py
-------------
Support for ephemeral simulation contexts.
"""
from typing import Optional, List, Iterator
from contextlib import contextmanager
from hieromem_core.memory import Memory
from hieromem_core.persistence.memory_store import InMemoryStore
from hieromem_core.episode import Episode


class SimulationMemory(Memory):
    """
    An ephemeral memory context that overlays a parent memory.

    Writes are local and discarded on exit.
    Reads aggregate results from both local and parent memory.
    """

    def __init__(self, parent: Memory):
        # Initialize with a fresh in-memory store
        # This store will hold the "diff" or "simulation delta"
        super().__init__(
            store=InMemoryStore(),
            encoder=parent.encoder,
            encoder_registry=parent.encoder_registry,
            multi_modal_encoder=parent.multi_modal_encoder,
            symbolic_reasoner=parent.symbolic_reasoner,
            # Share configuration
            Nmax=parent.Nmax,
            lam=parent.lam,
            tau=parent.tau,
            theta=parent.theta,
            debug=parent.debug,
            enable_multi_tenant=parent._multi_tenant,
            default_namespace=parent.namespace,
        )
        self.parent = parent
        # Share realms registry
        self.realms = parent.realms

    def retrieve(
        self,
        query: str,
        k: int = 5,
        *,
        mode: str = "graph_hybrid",
        alpha_dense: float = 1.0,
        beta_symbolic: float = 50.0,
        min_confidence: float = 0.0,
        namespace: Optional[str] = None,
        realm: Optional[str] = None,
        agent: Optional[str] = None,
    ) -> List[Episode]:
        """
        Merged retrieval from local simulation and parent memory.
        """
        # 1. Parent retrieval
        # We fetch k items from parent to ensure coverage
        parent_eps = self.parent.retrieve(
            query=query,
            k=k,
            mode=mode,
            alpha_dense=alpha_dense,
            beta_symbolic=beta_symbolic,
            min_confidence=min_confidence,
            namespace=namespace,
            realm=realm,
            agent=agent,
        )

        # 2. Local retrieval (simulation delta)
        local_eps = super().retrieve(
            query=query,
            k=k,
            mode=mode,
            alpha_dense=alpha_dense,
            beta_symbolic=beta_symbolic,
            min_confidence=min_confidence,
            namespace=namespace,
            realm=realm,
            agent=agent,
        )

        # 3. Merge and deduplicate by ID
        # Prioritize local episodes if IDs collide (updates/overwrites)
        seen_ids = set()
        merged = []

        # Add local first (higher priority)
        for ep in local_eps:
            seen_ids.add(ep.id)
            merged.append(ep)

        for ep in parent_eps:
            if ep.id not in seen_ids:
                seen_ids.add(ep.id)
                merged.append(ep)

        # 4. Re-rank based on score (if present)
        # Note: 'score' is in meta, populated by retrievers
        merged.sort(key=lambda x: x.meta.get("score", 0.0), reverse=True)

        return merged[:k]


@contextmanager
def simulate(memory: Memory) -> Iterator[SimulationMemory]:
    """
    Context manager for creating a simulation session.

    Usage:
        with simulate(main_memory) as sim:
            sim.insert(...)  # Only affects sim
            sim.retrieve(...) # Sees main + sim
    """
    sim_memory = SimulationMemory(memory)
    try:
        yield sim_memory
    finally:
        # Cleanup if necessary (InMemoryStore clears itself when GC'd)
        pass
