"""
cortex/core/agentic_retrieval.py — Agentic Retrieval Pipeline (Module 8 Phase 2).

Extends the Phase 1 RetrievalPipeline with LLM-powered multi-round retrieval.
Where Phase 1 does a single BM25+vector search, this pipeline can:

    1. Analyze the query → decide if it needs decomposition or refinement
    2. Run initial retrieval → assess quality of results
    3. If insufficient → refine the query and re-retrieve (up to max_rounds)
    4. Assemble cross-layer results (claims + episodes + procedures + cornerstones)
    5. Rerank using feedback scores + recency signals
    6. Enforce token budget across all rounds

This is the "smart" retrieval path — it costs LLM tokens but produces better
results for complex queries like "why did we decide to use SQLite?" or
"what's Andrew's preferred deployment strategy?"

Guard rails:
    - max_rounds (default 3) prevents infinite refinement loops
    - Each round is logged for debugging and cost tracking
    - Falls back to Phase 1 lightweight retrieval if LLM is unavailable
    - Token budget is enforced cumulatively across all rounds

Dependencies:
    - cortex.core.retrieval (Phase 1 pipeline, data types, scoring functions)
    - cortex.llm.provider.LLMProvider — query analysis + refinement
    - cortex.core.token_budget.TokenBudgetManager — budget enforcement

Data flow:
    RetrievalPipeline.retrieve(mode='agentic')
      → AgenticRetrievalPipeline.retrieve()
          → _analyze_query() (LLM: is the query complex enough for multi-round?)
          → _lightweight_retrieve() (Phase 1 hybrid search — round 1)
          → _assess_results() (LLM: are results sufficient?)
          → if insufficient: _refine_query() → re-retrieve (round 2+)
          → _rerank() (feedback scores + recency)
          → _apply_budget() (token budget trim)
      → RetrievalResult
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

from cortex.core.retrieval import (
    ContextBlock,
    RetrievalConstraints,
    RetrievedItem,
    RetrievalPipeline,
    RetrievalResult,
    _new_id,
)

if TYPE_CHECKING:
    from cortex.core.token_budget import TokenBudgetManager
    from cortex.llm.provider import LLMProvider
    from cortex.storage.adapter import StorageAdapter

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Agentic-specific data types
# ---------------------------------------------------------------------------

# Memory layers used for cross-layer assembly.
# Maps query constraint memory_types → display layer name
_MEMORY_LAYERS: Dict[str, List[str]] = {
    "claims":       ["claim"],
    "episodes":     ["event"],
    "procedures":   ["procedure"],
    "cornerstones": ["cornerstone"],
}

# All memory types we consider in cross-layer assembly (union of all layers)
_ALL_MEMORY_TYPES: List[str] = [
    t for types in _MEMORY_LAYERS.values() for t in types
]


@dataclass
class ReasoningDecision:
    """LLM reasoning output after analysing retrieval results."""
    should_refine: bool
    """True → run another retrieval round with refined_query."""

    refined_query: Optional[str] = None
    """New query to use in the next round (populated when should_refine=True)."""

    reasoning: str = ""
    """Human-readable explanation of the decision (for logs/explain)."""

    additional_filters: Optional[Dict[str, Any]] = None
    """Optional extra constraints for the next round (e.g. memory_types)."""


@dataclass
class AgenticRetrievalResult(RetrievalResult):
    """Extended result carrying agentic-specific diagnostics."""
    rounds_executed: int = 1
    """Number of retrieval rounds that were actually executed."""

    reasoning_trace: List[ReasoningDecision] = field(default_factory=list)
    """LLM decisions made at each intermediate round."""

    layer_counts: Dict[str, int] = field(default_factory=dict)
    """How many items came from each memory layer."""


# ---------------------------------------------------------------------------
# Reranking helpers (pure functions — easy to unit-test)
# ---------------------------------------------------------------------------

def rerank_with_feedback_and_recency(
    items: List[RetrievedItem],
    feedback_scores: Dict[str, float],
    item_ages_days: Dict[str, float],
    feedback_weight: float = 0.2,
    recency_weight: float = 0.1,
    recency_half_life_days: float = 30.0,
) -> List[RetrievedItem]:
    """
    Rerank retrieved items using:
      - Base RRF score (from retrieval)
      - Feedback helpfulness ratio (0..1 from FeedbackManager)
      - Recency bonus (exponential decay based on item age)

    All weights are additive on top of the base score.

    Args:
        items: Candidates from hybrid retrieval (sorted best-first).
        feedback_scores: {item_id: helpfulness_ratio} from FeedbackManager.
            Ratio in [0, 1] where 0 = always harmful, 1 = always helpful.
        item_ages_days: {item_id: age_in_days} (0 = created now).
        feedback_weight: Multiplier for the feedback contribution.
        recency_weight: Multiplier for the recency contribution.
        recency_half_life_days: Age (days) at which recency bonus is halved.

    Returns:
        New list of RetrievedItem sorted by combined score, best-first.
        Item.score is updated to the combined score.
    """
    import math

    reranked: List[RetrievedItem] = []
    for item in items:
        base = item.score

        # Feedback contribution: shift by (ratio - 0.5) → [-0.5, +0.5] range
        fb_ratio = feedback_scores.get(item.item_id, 0.5)  # default neutral
        feedback_bonus = feedback_weight * (fb_ratio - 0.5) * 2.0  # → [-fw, +fw]

        # Recency contribution: exponential decay e^(-age * ln2 / half_life)
        age = item_ages_days.get(item.item_id, recency_half_life_days)
        recency_bonus = recency_weight * math.exp(
            -age * math.log(2) / max(recency_half_life_days, 1.0)
        )

        combined = base + feedback_bonus + recency_bonus
        reranked.append(
            RetrievedItem(
                item_type=item.item_type,
                item_id=item.item_id,
                content=item.content,
                score=combined,
                provenance=item.provenance,
            )
        )

    reranked.sort(key=lambda x: x.score, reverse=True)
    return reranked


def deduplicate_items(items: List[RetrievedItem]) -> List[RetrievedItem]:
    """Remove duplicate item_ids, keeping the highest-scored copy."""
    seen: Dict[str, RetrievedItem] = {}
    for item in items:
        if item.item_id not in seen or item.score > seen[item.item_id].score:
            seen[item.item_id] = item
    # Preserve overall score order
    return sorted(seen.values(), key=lambda x: x.score, reverse=True)


# ---------------------------------------------------------------------------
# LLM reasoning prompt
# ---------------------------------------------------------------------------

_REFINEMENT_SYSTEM_PROMPT = """\
You are a memory retrieval reasoning engine.
Your job is to decide whether the retrieved memory results adequately answer the query,
or whether a refined search query would produce better results.

Respond ONLY with a JSON object. No extra text.

Schema:
{
  "should_refine": <true|false>,
  "refined_query": <string or null>,
  "reasoning": <string — 1-2 sentences>,
  "additional_filters": <object or null>
}

Rules:
- Set should_refine=true ONLY if the results are clearly missing important information
  that a refined query would find.
- If results already cover the query well, set should_refine=false.
- refined_query must be meaningfully different from the original — don't just rephrase.
- additional_filters can include: {"memory_types": ["claim", "event", "procedure"]}
  Use it to focus the next round on a specific memory layer if that would help.
- Be conservative: prefer stopping over unnecessary extra rounds.
"""

_REFINEMENT_USER_TEMPLATE = """\
Original query: {query}

Retrieved results ({n_results} items):
{results_preview}

Round: {round_num} of {max_rounds}

Decide: should we refine the search or stop here?
"""


def _format_results_preview(items: List[RetrievedItem], max_items: int = 5) -> str:
    """Format top items for LLM consumption (truncated)."""
    preview_items = items[:max_items]
    lines = []
    for i, item in enumerate(preview_items, 1):
        content_snippet = item.content[:200].replace("\n", " ")
        lines.append(f"{i}. [{item.item_type}] {content_snippet}…")
    if not lines:
        return "(no results)"
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Main agentic pipeline
# ---------------------------------------------------------------------------

class AgenticRetrievalPipeline(RetrievalPipeline):
    """
    Phase 2 agentic retrieval pipeline.

    Extends RetrievalPipeline with:
      - Multi-step retrieval driven by LLM reasoning
      - Cross-memory-layer assembly (claims + episodes + procedures + cornerstones)
      - Reranking using feedback helpfulness scores and item recency
      - Configurable max rounds (default 3)
      - Full TokenBudgetManager integration

    Usage:
        pipeline = AgenticRetrievalPipeline(
            storage=storage,
            config=config,
            token_budget=token_budget_manager,
            llm=llm_provider,
            max_rounds=3,
        )
        result = await pipeline.retrieve(query, mode="agentic")

    If ``llm`` is None, agentic mode falls back to a single-round cross-layer
    retrieval without refinement (no LLM calls).
    """

    def __init__(
        self,
        storage: "StorageAdapter",
        config: Any,
        token_budget: Optional["TokenBudgetManager"] = None,
        cornerstone_guardian: Any = None,
        llm: Any = None,
        *,
        max_rounds: int = 3,
        feedback_manager: Any = None,     # cortex.core.feedback.FeedbackManager (optional)
        vector_weight: Optional[float] = None,
        bm25_weight: Optional[float] = None,
        rrf_k: int = 60,
        candidate_multiplier: int = 3,
        rerank_feedback_weight: float = 0.2,
        rerank_recency_weight: float = 0.1,
        recency_half_life_days: float = 30.0,
    ) -> None:
        # Pull retrieval weights from config, falling back to hardcoded defaults
        if vector_weight is None:
            vector_weight = getattr(config, "retrieval_vector_weight", 0.7)
        if bm25_weight is None:
            bm25_weight = getattr(config, "retrieval_bm25_weight", 0.3)
        super().__init__(
            storage=storage,
            config=config,
            token_budget=token_budget,
            cornerstone_guardian=cornerstone_guardian,
            llm=llm,
            vector_weight=vector_weight,
            bm25_weight=bm25_weight,
            rrf_k=rrf_k,
            candidate_multiplier=candidate_multiplier,
        )
        self.max_rounds = max(1, max_rounds)
        self.feedback_manager = feedback_manager
        self.rerank_feedback_weight = rerank_feedback_weight
        self.rerank_recency_weight = rerank_recency_weight
        self.recency_half_life_days = recency_half_life_days

    # ------------------------------------------------------------------
    # Override: retrieve() — supports 'agentic' mode
    # ------------------------------------------------------------------

    async def retrieve(
        self,
        query: str,
        mode: str = "auto",
        constraints: Optional[RetrievalConstraints] = None,
        owner_id: str = "default",
    ) -> RetrievalResult:
        """
        Extended retrieve().

        New mode 'agentic':
          1. Cross-layer candidate generation (all memory types in parallel)
          2. RRF fusion
          3. LLM reasoning: should we refine? (up to max_rounds rounds)
          4. Rerank with feedback + recency
          5. Token-budgeted assembly
          6. Log retrieval

        All other modes are handled by the parent Phase 1 pipeline.
        """
        t0 = time.monotonic()
        retrieval_id = _new_id()
        constraints = constraints or RetrievalConstraints()

        # Resolve 'auto' mode — upgrade to agentic if LLM is available and
        # the query looks complex (semantic cues present).
        if mode == "auto":
            mode = self._route_query(query, constraints)

        if mode != "agentic":
            return await super().retrieve(query, mode=mode,
                                          constraints=constraints, owner_id=owner_id)

        # ---- Agentic path -------------------------------------------------
        result = await self._agentic_retrieve(
            query=query,
            constraints=constraints,
            owner_id=owner_id,
            retrieval_id=retrieval_id,
            t0=t0,
        )
        return result

    # ------------------------------------------------------------------
    # Override: assemble_context() — uses agentic retrieval
    # ------------------------------------------------------------------

    async def assemble_context(
        self,
        query: str,
        owner_id: str = "default",
        mode: str = "agentic",
    ) -> ContextBlock:
        """
        Assemble memory context for prompt injection.

        In agentic mode, retrieval is cross-layer and multi-round.
        Respects TokenBudgetManager allocations per memory layer.
        """
        retrieval_id = _new_id()

        total_budget = getattr(self.config, "retrieval_token_budget", 3000)
        if self.token_budget is not None:
            try:
                budget_obj = self.token_budget.create_budget(total=total_budget)
                available = min(total_budget, budget_obj.remaining())
            except Exception:
                available = total_budget
        else:
            available = total_budget

        result = await self.retrieve(
            query,
            mode=mode,
            constraints=RetrievalConstraints(token_budget=available),
            owner_id=owner_id,
        )

        parts: List[Tuple[str, str, int]] = []
        total_tokens = 0

        for item in result.items:
            item_tokens = self._count_tokens(item.content)
            if total_tokens + item_tokens > available:
                break
            parts.append((item.item_type, item.content, item_tokens))
            total_tokens += item_tokens

        return ContextBlock(
            parts=parts,
            total_tokens=total_tokens,
            budget_remaining=available - total_tokens,
            retrieval_id=retrieval_id,
        )

    # ------------------------------------------------------------------
    # Override: _route_query — upgrade complex queries to agentic
    # ------------------------------------------------------------------

    def _route_query(
        self,
        query: str,
        constraints: RetrievalConstraints,
    ) -> str:
        """
        Route 'auto' queries.

        Upgrades to 'agentic' when:
          - LLM is available
          - Query has semantic complexity cues (explain, why, how, compare, …)
          - Query is longer than 6 words

        Otherwise falls back to parent lightweight routing.
        """
        if self.llm is None:
            return "lightweight"

        words = query.lower().split()
        semantic_cues = {
            "explain", "why", "how", "what", "when", "who", "describe",
            "compare", "summarize", "relate", "find", "list", "show",
        }
        has_cues = bool(semantic_cues & set(words))
        is_complex = len(words) > 6 or has_cues

        return "agentic" if is_complex else "lightweight"

    # ------------------------------------------------------------------
    # Agentic core logic
    # ------------------------------------------------------------------

    async def _agentic_retrieve(
        self,
        *,
        query: str,
        constraints: RetrievalConstraints,
        owner_id: str,
        retrieval_id: str,
        t0: float,
    ) -> AgenticRetrievalResult:
        """
        Full agentic retrieval loop:
          Round 1: cross-layer retrieval with original query
          Rounds 2..N: LLM refinement → re-retrieve → merge

        Returns AgenticRetrievalResult with merged, reranked, budget-trimmed items.
        """
        current_query = query
        all_items: List[RetrievedItem] = []
        reasoning_trace: List[ReasoningDecision] = []
        rounds_executed = 0
        seen_queries: set = {query}

        for round_num in range(1, self.max_rounds + 1):
            rounds_executed = round_num

            # Check remaining token budget before each round
            budget_remaining = self._get_budget_remaining(constraints)
            if budget_remaining <= 0:
                logger.info("Agentic retrieval: token budget exhausted, stopping at round %d",
                            round_num)
                break

            logger.debug("Agentic round %d/%d: query=%r", round_num, self.max_rounds,
                         current_query)

            # Cross-layer candidate retrieval
            round_items = await self._cross_layer_retrieve(
                query=current_query,
                constraints=constraints,
                owner_id=owner_id,
            )
            all_items.extend(round_items)

            # If this is the last allowed round, skip LLM reasoning
            if round_num >= self.max_rounds:
                break

            # LLM reasoning: should we refine?
            decision = await self._reason_about_refinement(
                original_query=query,
                current_query=current_query,
                items=round_items,
                round_num=round_num,
                max_rounds=self.max_rounds,
            )
            reasoning_trace.append(decision)

            if not decision.should_refine:
                logger.debug("Agentic round %d: LLM decided to stop.", round_num)
                break

            # Prevent infinite loops: don't repeat a query we've already tried
            next_query = decision.refined_query or current_query
            if next_query in seen_queries:
                logger.debug("Agentic: refined_query %r already seen — stopping.", next_query)
                break
            seen_queries.add(next_query)

            # Apply any additional filters from LLM (e.g. restrict memory_types)
            if decision.additional_filters:
                af = decision.additional_filters
                if "memory_types" in af and isinstance(af["memory_types"], list):
                    constraints = RetrievalConstraints(
                        memory_types=af["memory_types"],
                        time_range=constraints.time_range,
                        entity_ids=constraints.entity_ids,
                        sensitivity_max=constraints.sensitivity_max,
                        token_budget=constraints.token_budget,
                        top_k=constraints.top_k,
                    )

            current_query = next_query

        # Deduplicate across all rounds
        merged = deduplicate_items(all_items)

        # Rerank with feedback + recency
        merged = await self._rerank(merged, owner_id=owner_id)

        # Count items per layer
        layer_counts = _count_by_layer(merged)

        # Apply token budget
        tokens_used, trimmed = await self._apply_budget(merged, constraints)

        latency_ms = int((time.monotonic() - t0) * 1000)

        # Log retrieval
        await self._log_retrieval(
            retrieval_id=retrieval_id,
            owner_id=owner_id,
            query=query,
            mode="agentic",
            constraints=constraints,
            items=trimmed,
            tokens_used=tokens_used,
            latency_ms=latency_ms,
        )

        return AgenticRetrievalResult(
            items=trimmed,
            retrieval_id=retrieval_id,
            mode="agentic",
            latency_ms=latency_ms,
            tokens_used=tokens_used,
            rounds_executed=rounds_executed,
            reasoning_trace=reasoning_trace,
            layer_counts=layer_counts,
        )

    # ------------------------------------------------------------------
    # Cross-layer retrieval
    # ------------------------------------------------------------------

    async def _cross_layer_retrieve(
        self,
        *,
        query: str,
        constraints: RetrievalConstraints,
        owner_id: str,
    ) -> List[RetrievedItem]:
        """
        Run lightweight retrieval across all memory layers in parallel.

        Each layer gets its own constraints (filtered memory_types), and
        results are merged via RRF to produce a single ranked list.

        If constraints already specify memory_types, only those layers run.
        """
        target_layers = constraints.memory_types or _ALL_MEMORY_TYPES

        # Build per-layer constraints
        tasks: List[asyncio.Task] = []
        layer_names: List[str] = []

        for memory_type in target_layers:
            layer_constraints = RetrievalConstraints(
                memory_types=[memory_type],
                time_range=constraints.time_range,
                entity_ids=constraints.entity_ids,
                sensitivity_max=constraints.sensitivity_max,
                token_budget=None,         # budget applied after merge
                top_k=constraints.top_k,
            )
            task = asyncio.create_task(
                self._lightweight_retrieve(query, layer_constraints, owner_id),
                name=f"layer_{memory_type}",
            )
            tasks.append(task)
            layer_names.append(memory_type)

        results_per_layer: List[List[RetrievedItem]] = await asyncio.gather(
            *tasks, return_exceptions=True
        )

        # Collect items from all layers (skip failed layers gracefully)
        all_layer_items: List[RetrievedItem] = []
        for layer_name, layer_result in zip(layer_names, results_per_layer):
            if isinstance(layer_result, Exception):
                logger.warning("Cross-layer retrieval failed for layer=%s: %s",
                               layer_name, layer_result)
                continue
            all_layer_items.extend(layer_result)

        return all_layer_items

    # ------------------------------------------------------------------
    # LLM reasoning
    # ------------------------------------------------------------------

    async def _reason_about_refinement(
        self,
        *,
        original_query: str,
        current_query: str,
        items: List[RetrievedItem],
        round_num: int,
        max_rounds: int,
    ) -> ReasoningDecision:
        """
        Ask the LLM whether to refine the search query.

        Returns a ReasoningDecision. If LLM is unavailable or errors,
        returns a conservative "stop" decision so retrieval always terminates.
        """
        if self.llm is None:
            return ReasoningDecision(
                should_refine=False,
                reasoning="LLM unavailable — single-round retrieval only.",
            )

        results_preview = _format_results_preview(items, max_items=5)
        user_msg = _REFINEMENT_USER_TEMPLATE.format(
            query=current_query,
            n_results=len(items),
            results_preview=results_preview,
            round_num=round_num,
            max_rounds=max_rounds,
        )

        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "agentic_refinement")
            response = await self.llm.complete(
                system=_REFINEMENT_SYSTEM_PROMPT,
                user=user_msg,
                response_format="json",
                max_tokens=256,
                temperature=0.0,
                **({"model": _routed_model} if _routed_model else {}),
            )
            raw_text = (
                response.text
                if hasattr(response, "text")
                else str(response)
            )
            data = json.loads(raw_text)

            return ReasoningDecision(
                should_refine=bool(data.get("should_refine", False)),
                refined_query=data.get("refined_query") or None,
                reasoning=data.get("reasoning", ""),
                additional_filters=data.get("additional_filters") or None,
            )

        except json.JSONDecodeError as exc:
            logger.warning("Agentic reasoning: LLM returned non-JSON response: %s", exc)
            return ReasoningDecision(should_refine=False, reasoning="Parse error — stopping.")

        except Exception as exc:
            logger.warning("Agentic reasoning: LLM call failed: %s", exc)
            return ReasoningDecision(should_refine=False, reasoning=f"LLM error: {exc}")

    # ------------------------------------------------------------------
    # Reranking
    # ------------------------------------------------------------------

    async def _rerank(
        self,
        items: List[RetrievedItem],
        owner_id: str,
    ) -> List[RetrievedItem]:
        """
        Rerank items using feedback helpfulness scores and item recency.

        Fetches scores from FeedbackManager (if available) and item ages from
        storage (if available). Falls back gracefully if either is unavailable.
        """
        if not items:
            return items

        item_ids = [item.item_id for item in items]

        # Gather feedback scores
        feedback_scores = await self._fetch_feedback_scores(item_ids, owner_id)

        # Gather item ages
        item_ages = await self._fetch_item_ages(item_ids, owner_id)

        return rerank_with_feedback_and_recency(
            items=items,
            feedback_scores=feedback_scores,
            item_ages_days=item_ages,
            feedback_weight=self.rerank_feedback_weight,
            recency_weight=self.rerank_recency_weight,
            recency_half_life_days=self.recency_half_life_days,
        )

    async def _fetch_feedback_scores(
        self,
        item_ids: List[str],
        owner_id: str,
    ) -> Dict[str, float]:
        """
        Fetch helpfulness ratios from FeedbackManager in parallel.

        Returns dict of {item_id: score}. Missing items default to 0.5 (neutral).
        """
        if self.feedback_manager is None or not item_ids:
            return {}

        async def get_one(item_id: str) -> Tuple[str, float]:
            try:
                score = await self.feedback_manager.get_score(
                    target_id=item_id, owner_id=owner_id
                )
                # get_score may return a HelpfulnessScore object or a float
                if hasattr(score, "helpfulness_ratio"):
                    return item_id, float(score.helpfulness_ratio)
                return item_id, float(score) if score is not None else 0.5
            except Exception as exc:
                logger.debug("Feedback score fetch failed for %s: %s", item_id, exc)
                return item_id, 0.5

        pairs = await asyncio.gather(*[get_one(iid) for iid in item_ids])
        return dict(pairs)

    async def _fetch_item_ages(
        self,
        item_ids: List[str],
        owner_id: str,
    ) -> Dict[str, float]:
        """
        Fetch item ages in days from storage.

        Returns dict of {item_id: age_days}. Missing items default to half-life
        (neutral recency weight).
        """
        if not item_ids:
            return {}

        if not hasattr(self.storage, "get_item_ages"):
            # Storage doesn't support age lookup — return neutral ages
            return {iid: self.recency_half_life_days for iid in item_ids}

        try:
            ages = await self.storage.get_item_ages(item_ids, owner_id=owner_id)
            if isinstance(ages, dict):
                return ages
        except Exception as exc:
            logger.debug("Item age fetch failed: %s", exc)

        return {iid: self.recency_half_life_days for iid in item_ids}

    # ------------------------------------------------------------------
    # Budget helper
    # ------------------------------------------------------------------

    def _get_budget_remaining(self, constraints: RetrievalConstraints) -> int:
        """
        Get remaining token budget for the current retrieval session.

        Checks constraints.token_budget first, then config, then TokenBudgetManager.
        Returns a large number (no limit) if no budget is configured.
        """
        if constraints.token_budget is not None:
            return constraints.token_budget

        config_budget = getattr(self.config, "retrieval_token_budget", None)
        if config_budget is not None:
            return config_budget

        if self.token_budget is not None:
            try:
                b = self.token_budget.create_budget()
                return b.remaining()
            except Exception as e:
                logger.warning("cortex_core_agentic_retrieval: unhandled exception: %s", e)

        return 999_999  # no effective limit


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _count_by_layer(items: List[RetrievedItem]) -> Dict[str, int]:
    """Count how many items belong to each memory layer/type."""
    counts: Dict[str, int] = {}
    for item in items:
        counts[item.item_type] = counts.get(item.item_type, 0) + 1
    return counts
