"""
sequence.py
-----------
Replay functionality for ordered memory retrieval.
"""
from typing import Iterator, Optional
from datetime import datetime
from hieromem_core.memory import Memory
from hieromem_core.episode import Episode


def replay(
    memory: Memory,
    start: Optional[datetime] = None,
    end: Optional[datetime] = None,
    namespace: Optional[str] = None
) -> Iterator[Episode]:
    """
    Yield episodes in chronological order.

    Args:
        memory: The memory instance to replay from.
        start: Start datetime (inclusive).
        end: End datetime (inclusive).
        namespace: Optional namespace to filter by.

    Yields:
        Episodes sorted by creation time (tocc).
    """
    # 1. Fetch episodes
    # Access internal list logic or use store?
    # Memory keeps episodes in self.episodes (if single tenant)
    # or self._episodes_by_namespace

    all_episodes = []

    if memory._multi_tenant:
        if namespace:
            all_episodes = list(memory._episodes_by_namespace[namespace])
        else:
            # Aggregate all namespaces?
            for ns_eps in memory._episodes_by_namespace.values():
                all_episodes.extend(ns_eps)
    else:
        all_episodes = list(memory.episodes)

    # 2. Sort by time (just to be safe, though usually insertion order = time order)
    all_episodes.sort(key=lambda x: x.tocc)

    # 3. Filter range
    for ep in all_episodes:
        if start and ep.tocc < start:
            continue
        if end and ep.tocc > end:
            continue
        yield ep
