"""Consolidation.

Raw episodic memory is redundant. Consolidation folds repeated observations
about the same subject into a single durable summary while recording, in the
ledger, exactly which episodes it was derived from. The originals are never
deleted — they remain the ground truth.

The grouping key here is deliberately simple (shared top tags / noun overlap)
so the code stays dependency-free and explainable. Real systems can substitute
an embedding-clustering step behind the same interface.
"""

from __future__ import annotations

from collections import Counter
from typing import Callable, Dict, List, Optional

from .ledger import ProvenanceLedger
from .mem import Entry, MemoryStore, _tokenize


class Consolidator:
    """Group similar entries and write a consolidated summary per group.

    Parameters
    ----------
    key_fn:
        Maps an Entry to a grouping key string. Default groups by the most
        common tag(s); entries with no tags fall back to their top token.
    """

    def __init__(
        self,
        key_fn: Optional[Callable[[Entry], str]] = None,
        topk: int = 3,
    ) -> None:
        self.key_fn = key_fn or self._default_key
        self.topk = topk

    @staticmethod
    def _default_key(entry: Entry) -> str:
        if entry.tags:
            return " ".join(sorted(entry.tags)[:2])
        toks = _tokenize(entry.text)
        if not toks:
            return "_empty"
        return toks[0]

    def consolidate(
        self,
        store: MemoryStore,
        ledger: Optional[ProvenanceLedger] = None,
        min_group: int = 1,
    ) -> List[Entry]:
        """Return a new list of summary :class:`Entry` objects.

        Each summary entry carries the id prefix ``consolidated-`` and its
        metadata lists ``derived_from`` so the ledger can trace it.
        """
        groups: Dict[str, List[Entry]] = {}
        for entry in store:
            groups.setdefault(self.key_fn(entry), []).append(entry)

        summaries: List[Entry] = []
        for key, members in groups.items():
            if len(members) < min_group:
                continue
            combined = Counter()
            for m in members:
                combined.update(_tokenize(m.text))
            top_tokens = [tok for tok, _ in combined.most_common(self.topk)]
            summary_text = (
                f"Consolidated memory for {key!r}: "
                + ", ".join(f"'{t}'" for t in top_tokens)
                + f" (synthesized from {len(members)} episodes)."
            )
            summary = Entry(
                text=summary_text,
                source="consolidation",
                tags=["consolidated"],
                metadata={"derived_from": [m.id for m in members], "group_key": key},
                id=f"consolidated-{key[:40]}",
            )
            summaries.append(summary)
            if ledger is not None:
                ledger.record_consolidation(summary.id, [m.id for m in members], reason="duplicate compression")
        return summaries