app/retrieval/retrieval_trace.py

"""
RetrievalTrace — full pipeline provenance for every LETA query.

Design principles:
  - Zero retrieval-behavior change: recording is purely additive / read-only.
  - Non-failing: every recording method is wrapped; trace errors never surface to callers.
  - Thread-safe for reads after the retrieval thread finishes writing.
  - Two output formats:
      to_log_dict()  → compact single JSON line for CloudWatch
      to_debug_dict() → full detail for /debug/trace/{query_id}

Usage (inside retriever.search()):
    if trace:
        trace.record_faiss(chunks_with_scores)
        ...
    After every transformation:
        trace.snapshot_stage("after_crossencoder", chunks)
"""

from __future__ import annotations

import json
import logging
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# ─── Helpers ──────────────────────────────────────────────────────────────────

def _doc_id(rel_path: str) -> str:
    """Stable human-readable document ID from rel_path (filename without ext)."""
    if not rel_path:
        return "unknown"
    return Path(rel_path.replace("\\", "/")).stem


def _doc_type(rel_path: str) -> str:
    """Classify document category from path (mirrors _chunk_category in retriever)."""
    p = rel_path.lower().replace("\\", "/")
    if any(f in p for f in ("high court", "supreme court", "aar", "other app")):
        return "case_law"
    if any(f in p for f in ("notification", "notifications")):
        return "notification"
    if any(f in p for f in ("circulars", "circular", "icai", "brochures", "faqs")):
        return "circular"
    if any(f in p for f in ("act", "rules", "cgst", "igst", "utgst", "export")):
        return "statute"
    return "other"


def _safe(fn):
    """Decorator: silently swallow any exception inside trace recording methods."""
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except Exception as e:
            logger.debug(f"RetrievalTrace recording error (non-fatal): {e}")
    return wrapper


# ─── Per-chunk provenance record ──────────────────────────────────────────────

@dataclass
class ChunkRecord:
    chunk_id: str
    document_id: str
    document_type: str
    rel_path: str
    text_preview: str          # first 180 chars

    # Where this chunk first entered the candidate pool
    entered_at: str = "unknown"   # faiss | bm25 | tfidf | taxonomy_pin | quota_fill |
                                  # circular_bm25 | circular_faiss | statute_faiss |
                                  # circular_floor | layer6 | mandatory_inject |
                                  # citation_graph | self_critique

    # Which retrieval signals found this chunk
    retrieval_sources: list = field(default_factory=list)

    # Per-source scores / ranks
    faiss_rank:  Optional[int]   = None
    faiss_score: Optional[float] = None
    bm25_rank:   Optional[int]   = None
    bm25_score:  Optional[float] = None
    tfidf_rank:  Optional[int]   = None
    tfidf_score: Optional[float] = None
    rrf_score:   Optional[float] = None

    # Reranking scores
    cross_encoder_score:  Optional[float] = None
    legal_reranker_score: Optional[float] = None
    authority_boost:      Optional[float] = None
    doc_level_boost:      Optional[float] = None

    # Injection metadata (for non-FAISS/BM25 paths)
    injection_reason: Optional[str]  = None   # e.g. "circular_floor", "layer6_statute"
    injection_score:  Optional[float] = None

    # MMR
    mmr_kept: Optional[bool]  = None
    mmr_similarity: Optional[float] = None  # max Jaccard to already-selected chunks

    # Provision keys from metadata (P2.5 — used by gold matcher)
    provisions: list = field(default_factory=list)

    # Provision Anchoring metadata (P2.5)
    anchor_provision: Optional[str] = None    # which taxonomy provision key pinned this

    # Final outcome
    selected: bool = False
    selection_reasons: list = field(default_factory=list)
    elimination_reason: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "chunk_id":            self.chunk_id,
            "document_id":         self.document_id,
            "document_type":       self.document_type,
            "rel_path":            self.rel_path,
            "text_preview":        self.text_preview,
            "entered_at":          self.entered_at,
            "retrieval_sources":   self.retrieval_sources,
            "faiss_rank":          self.faiss_rank,
            "faiss_score":         self.faiss_score,
            "bm25_rank":           self.bm25_rank,
            "bm25_score":          self.bm25_score,
            "tfidf_rank":          self.tfidf_rank,
            "tfidf_score":         self.tfidf_score,
            "rrf_score":           self.rrf_score,
            "cross_encoder_score": self.cross_encoder_score,
            "legal_reranker_score":self.legal_reranker_score,
            "authority_boost":     self.authority_boost,
            "doc_level_boost":     self.doc_level_boost,
            "injection_reason":    self.injection_reason,
            "injection_score":     self.injection_score,
            "mmr_kept":            self.mmr_kept,
            "mmr_similarity":      self.mmr_similarity,
            "provisions":          self.provisions,
            "anchor_provision":    self.anchor_provision,
            "selected":            self.selected,
            "selection_reasons":   self.selection_reasons,
            "elimination_reason":  self.elimination_reason,
        }


# ─── Main trace object ────────────────────────────────────────────────────────

class RetrievalTrace:
    """
    Accumulates full retrieval provenance for one LETA query.
    Created at /ask entry, passed through retriever.search() and supplement_and_rerank(),
    finalized before CloudWatch logging.

    All recording methods are @_safe — no exception can escape into the retrieval path.
    """

    def __init__(self, query_id: str, query: str):
        self.query_id  = query_id
        self.query     = query
        self.created_at = time.time()

        # Preprocessing (populated from app.py before calling retriever)
        self.preprocessing: dict = {}

        # Per-chunk records — the core provenance store
        self._chunks: dict[str, ChunkRecord] = {}

        # Stage snapshots: ordered list of (chunk_id, score) at each point
        # Key = stage name, value = list of {chunk_id, score, rank}
        self._stages: dict[str, list] = {}

        # Document-level aggregation (populated at finalize())
        self.documents: dict = {}

        # Counts at each stage
        self.stage_counts: dict = {}

        # Validation / coverage
        self.validation: dict = {}

        # Final context + answer
        self.final_context: dict = {}
        self.answer: dict = {}

        # Internal: set of seen chunk_ids to avoid double-recording entered_at
        self._entered: set = set()

    # ── Internal helpers ───────────────────────────────────────────────────────

    def _get_or_create(self, chunk: dict) -> Optional[ChunkRecord]:
        """Return the ChunkRecord for this chunk, creating it if needed."""
        cid = chunk.get("chunk_id")
        if not cid:
            return None
        if cid not in self._chunks:
            rel = (chunk.get("rel_path") or
                   chunk.get("metadata", {}).get("rel_path", "") or
                   chunk.get("source", ""))
            _meta = chunk.get("metadata", {})
            # P2.5: capture provision keys so gold matcher can identify statute chunks
            # even when the text_preview doesn't explicitly say "Section N"
            _provs = list(set(
                _meta.get("provisions", []) + _meta.get("citations", [])
            ))
            self._chunks[cid] = ChunkRecord(
                chunk_id     = cid,
                document_id  = _doc_id(rel),
                document_type= _doc_type(rel),
                rel_path     = rel,
                text_preview = (chunk.get("text") or "")[:180],
                provisions   = _provs,
                anchor_provision = chunk.get("_anchor_provision"),
            )
        else:
            # Update anchor_provision if newly set (e.g. chunk re-encountered as pinned)
            _rec = self._chunks[cid]
            if not _rec.anchor_provision and chunk.get("_anchor_provision"):
                _rec.anchor_provision = chunk["_anchor_provision"]
        return self._chunks[cid]

    # ── Preprocessing ──────────────────────────────────────────────────────────

    @_safe
    def record_preprocessing(self, *,
                              original_query: str,
                              refined_query: str = "",
                              sub_queries: list = None,
                              hyde_doc: str = "",
                              topic: str = "",
                              subtopic: str = "",
                              detected_refs: list = None,
                              taxonomy: dict = None,
                              domain_route: list = None,
                              complexity_score: float = 0.0,
                              response_mode: str = ""):
        self.preprocessing = {
            "original_query":  original_query,
            "refined_query":   refined_query or original_query,
            "sub_queries":     sub_queries or [],
            "hyde_doc":        hyde_doc[:400] if hyde_doc else "",
            "topic":           topic,
            "subtopic":        subtopic,
            "detected_refs":   detected_refs or [],
            "taxonomy": {
                "topics":      (taxonomy or {}).get("topics", []),
                "sections":    (taxonomy or {}).get("sections", []),
                "rules":       (taxonomy or {}).get("rules", []),
                "circulars":   (taxonomy or {}).get("circulars", []),
                "confidence":  (taxonomy or {}).get("confidence", 0),
                "expected_cats": list((taxonomy or {}).get("expected_cats", set())),
            },
            "domain_route":    domain_route or [],
            "complexity_score": round(complexity_score, 3),
            "response_mode":   response_mode,
        }

    # ── Stage recording ────────────────────────────────────────────────────────

    @_safe
    def record_faiss(self, chunks: list, scores: list):
        """Record FAISS results. chunks[i] ↔ scores[i] (cosine similarities)."""
        snap = []
        for rank, (ch, sim) in enumerate(zip(chunks, scores)):
            rec = self._get_or_create(ch)
            if rec is None:
                continue
            if rec.chunk_id not in self._entered:
                rec.entered_at = "faiss"
                self._entered.add(rec.chunk_id)
            if "faiss" not in rec.retrieval_sources:
                rec.retrieval_sources.append("faiss")
            # Only record if this is the best FAISS rank seen for this chunk
            if rec.faiss_rank is None or rank < rec.faiss_rank:
                rec.faiss_rank  = rank
                rec.faiss_score = round(float(sim), 5)
            snap.append({"chunk_id": rec.chunk_id, "rank": rank, "score": round(float(sim), 5)})
        self._stages["faiss"] = snap
        self.stage_counts["faiss"] = len(snap)

    @_safe
    def record_bm25(self, chunks: list, scores: list):
        """Record BM25 results. chunks[i] ↔ scores[i]."""
        snap = []
        for rank, (ch, score) in enumerate(zip(chunks, scores)):
            rec = self._get_or_create(ch)
            if rec is None:
                continue
            if rec.chunk_id not in self._entered:
                rec.entered_at = "bm25"
                self._entered.add(rec.chunk_id)
            if "bm25" not in rec.retrieval_sources:
                rec.retrieval_sources.append("bm25")
            if rec.bm25_rank is None or rank < rec.bm25_rank:
                rec.bm25_rank  = rank
                rec.bm25_score = round(float(score), 5)
            snap.append({"chunk_id": rec.chunk_id, "rank": rank, "score": round(float(score), 5)})
        self._stages["bm25"] = snap
        self.stage_counts["bm25"] = len(snap)

    @_safe
    def record_tfidf(self, chunks: list, scores: list):
        """Record TF-IDF results."""
        snap = []
        for rank, (ch, score) in enumerate(zip(chunks, scores)):
            rec = self._get_or_create(ch)
            if rec is None:
                continue
            if rec.chunk_id not in self._entered:
                rec.entered_at = "tfidf"
                self._entered.add(rec.chunk_id)
            if "tfidf" not in rec.retrieval_sources:
                rec.retrieval_sources.append("tfidf")
            if rec.tfidf_rank is None or rank < rec.tfidf_rank:
                rec.tfidf_rank  = rank
                rec.tfidf_score = round(float(score), 5)
            snap.append({"chunk_id": rec.chunk_id, "rank": rank, "score": round(float(score), 5)})
        self._stages["tfidf"] = snap
        self.stage_counts["tfidf"] = len(snap)

    @_safe
    def record_rrf(self, chunks: list):
        """Record RRF-merged pool. Each chunk should have _rrf_score set."""
        snap = []
        for rank, ch in enumerate(chunks):
            rec = self._get_or_create(ch)
            if rec is None:
                continue
            rrf_val = ch.get("_rrf_score")
            if rrf_val is not None:
                rec.rrf_score = round(float(rrf_val), 6)
            snap.append({"chunk_id": rec.chunk_id, "rank": rank, "score": rec.rrf_score})
        self._stages["rrf"] = snap
        self.stage_counts["rrf"] = len(snap)

    @_safe
    def record_injected(self, chunk: dict, reason: str, score: float = None):
        """Record a chunk injected by any non-FAISS/BM25 mechanism."""
        rec = self._get_or_create(chunk)
        if rec is None:
            return
        if rec.chunk_id not in self._entered:
            rec.entered_at = reason
            self._entered.add(rec.chunk_id)
        rec.injection_reason = reason
        if score is not None:
            rec.injection_score = round(float(score), 5)
        if reason not in rec.retrieval_sources:
            rec.retrieval_sources.append(reason)

    @_safe
    def snapshot_stage(self, name: str, chunks: list, score_field: str = "_final_legal_score"):
        """Capture an ordered snapshot of (chunk_id, score) at a named stage.

        Used after CrossEncoder, LegalReranker, doc-boost, MMR, circular-floor, Layer 6.
        Also updates the per-chunk record with the score from this stage if it's a
        reranking stage.
        """
        snap = []
        for rank, ch in enumerate(chunks):
            cid = ch.get("chunk_id")
            if not cid:
                continue
            score = ch.get(score_field)
            snap.append({
                "chunk_id": cid,
                "rank":     rank,
                "score":    round(float(score), 5) if score is not None else None,
            })
            # Update per-chunk record with stage-specific scores
            rec = self._chunks.get(cid)
            if rec is None:
                rec = self._get_or_create(ch)
            if rec:
                if name == "after_crossencoder":
                    rec.cross_encoder_score = round(float(score), 5) if score else None
                    # Also capture authority_boost if stored on chunk
                    ab = ch.get("_authority_boost")
                    if ab is not None:
                        rec.authority_boost = round(float(ab), 4)
                elif name in ("after_legalreranker", "after_doc_boost", "after_circular_floor",
                              "after_layer6", "after_mandatory", "final"):
                    rec.legal_reranker_score = round(float(score), 5) if score else None

        self._stages[name] = snap
        self.stage_counts[name] = len(snap)

    @_safe
    def record_crossencoder_scores(self, chunks: list):
        """After _cascade_rerank: read _rerank_score from every chunk."""
        snap = []
        for rank, ch in enumerate(chunks):
            cid = ch.get("chunk_id")
            if not cid:
                continue
            rec = self._chunks.get(cid) or self._get_or_create(ch)
            if rec:
                rs = ch.get("_rerank_score")
                if rs is not None:
                    rec.cross_encoder_score = round(float(rs), 5)
                ab = ch.get("_authority_boost")
                if ab is not None:
                    rec.authority_boost = round(float(ab), 4)
            score = ch.get("_rerank_score")
            snap.append({"chunk_id": cid, "rank": rank,
                         "score": round(float(score), 5) if score else None})
        self._stages["after_crossencoder"] = snap
        self.stage_counts["after_crossencoder"] = len(snap)

    @_safe
    def record_legalreranker_scores(self, chunks: list):
        """After LegalReranker: read _final_legal_score from every chunk."""
        snap = []
        for rank, ch in enumerate(chunks):
            cid = ch.get("chunk_id")
            if not cid:
                continue
            rec = self._chunks.get(cid) or self._get_or_create(ch)
            if rec:
                fs = ch.get("_final_legal_score")
                if fs is not None:
                    rec.legal_reranker_score = round(float(fs), 5)
                db = ch.get("_doc_hit_count")
                if db is not None and db > 1:
                    import math
                    rec.doc_level_boost = round(math.log(db) * 0.02, 5)
            score = ch.get("_final_legal_score")
            snap.append({"chunk_id": cid, "rank": rank,
                         "score": round(float(score), 5) if score else None})
        self._stages["after_legalreranker"] = snap
        self.stage_counts["after_legalreranker"] = len(snap)

    @_safe
    def record_mmr(self, pre_mmr: list, post_mmr: list):
        """Compare pre/post MMR lists to record what was kept vs. removed."""
        kept_ids  = {c.get("chunk_id") for c in post_mmr if c.get("chunk_id")}
        snap_pre  = []
        snap_post = []

        for rank, ch in enumerate(pre_mmr):
            cid = ch.get("chunk_id")
            if not cid:
                continue
            snap_pre.append({"chunk_id": cid, "rank": rank,
                              "score": round(float(ch.get("_final_legal_score", 0)), 5)})
            rec = self._chunks.get(cid) or self._get_or_create(ch)
            if rec:
                if cid in kept_ids:
                    rec.mmr_kept = True
                else:
                    rec.mmr_kept = False
                    rec.elimination_reason = "mmr_removed"

        for rank, ch in enumerate(post_mmr):
            cid = ch.get("chunk_id")
            if not cid:
                continue
            snap_post.append({"chunk_id": cid, "rank": rank,
                               "score": round(float(ch.get("_final_legal_score", 0)), 5)})

        self._stages["pre_mmr"]  = snap_pre
        self._stages["post_mmr"] = snap_post
        self.stage_counts["pre_mmr"]  = len(snap_pre)
        self.stage_counts["post_mmr"] = len(snap_post)
        self.stage_counts["mmr_removed"] = len(snap_pre) - len(snap_post)

    # ── Validation ────────────────────────────────────────────────────────────

    @_safe
    def record_validation(self, *,
                           expected_cats: set = None,
                           present_cats: set = None,
                           missing_cats: set = None,
                           layer6_injections: list = None,
                           mandatory_coverage_pct: float = 100.0,
                           mandatory_missing: list = None):
        self.validation = {
            "expected_categories":    sorted(expected_cats or []),
            "present_categories":     sorted(present_cats or []),
            "missing_categories":     sorted(missing_cats or []),
            "layer6_injections":      layer6_injections or [],
            "mandatory_coverage_pct": mandatory_coverage_pct,
            "mandatory_missing":      mandatory_missing or [],
        }

    # ── Finalization ──────────────────────────────────────────────────────────

    @_safe
    def finalize(self, final_chunks: list, answer_meta: dict = None):
        """
        Called after search() returns. Marks which chunks are selected,
        builds document aggregation, records answer metadata.
        """
        selected_ids = {c.get("chunk_id") for c in final_chunks if c.get("chunk_id")}

        # Mark selected
        for ch in final_chunks:
            cid = ch.get("chunk_id")
            if not cid:
                continue
            rec = self._chunks.get(cid) or self._get_or_create(ch)
            if rec:
                rec.selected = True
                # Build selection_reasons from chunk metadata flags
                reasons = []
                if ch.get("_pinned_by_ref"):       reasons.append("taxonomy_pin")
                if ch.get("_targeted_fill"):        reasons.append("quota_fill")
                if ch.get("_targeted_circ_bm25"):   reasons.append("circular_bm25")
                if ch.get("_targeted_circ_faiss"):  reasons.append("circular_faiss")
                if ch.get("_coverage_fill"):        reasons.append("layer6_coverage")
                if ch.get("_mandatory_inject"):     reasons.append("mandatory_inject")
                if ch.get("_citation_graph_expand"):reasons.append("citation_graph")
                if ch.get("_self_critique_inject"): reasons.append("self_critique")
                # Infer FAISS / BM25 from existing record
                if rec.faiss_rank is not None:  reasons.append("faiss")
                if rec.bm25_rank  is not None:  reasons.append("bm25")
                if rec.tfidf_rank is not None:  reasons.append("tfidf")
                # Deduplicate preserving order
                seen = set()
                rec.selection_reasons = [r for r in reasons if not (r in seen or seen.add(r))]
                # Final score
                fs = ch.get("_final_legal_score") or ch.get("_rerank_score")
                if fs is not None:
                    rec.legal_reranker_score = round(float(fs), 5)

        # Mark non-selected as eliminated (if not already MMR-removed)
        for cid, rec in self._chunks.items():
            if cid not in selected_ids and not rec.selected:
                if rec.elimination_reason is None:
                    rec.elimination_reason = "score_cutoff"

        # Build final stage snapshot
        snap = []
        for rank, ch in enumerate(final_chunks):
            cid = ch.get("chunk_id")
            if not cid:
                continue
            score = ch.get("_final_legal_score") or ch.get("_rerank_score")
            snap.append({"chunk_id": cid, "rank": rank,
                          "score": round(float(score), 5) if score else None,
                          "document_id": _doc_id(ch.get("rel_path") or ""),
                          "document_type": _doc_type(ch.get("rel_path") or "")})
        self._stages["final"] = snap

        # Document aggregation
        doc_data: dict = {}
        for rank, ch in enumerate(final_chunks):
            rel  = ch.get("rel_path") or ""
            did  = _doc_id(rel)
            score = float(ch.get("_final_legal_score") or ch.get("_rerank_score") or 0)
            srcs = self._chunks.get(ch.get("chunk_id"), ChunkRecord("","","","","")).retrieval_sources
            if did not in doc_data:
                doc_data[did] = {
                    "document_id":       did,
                    "document_type":     _doc_type(rel),
                    "rel_path":          rel,
                    "chunks_in_final":   0,
                    "best_chunk_score":  0.0,
                    "mean_chunk_score":  0.0,
                    "_score_sum":        0.0,
                    "retrieval_sources": set(),
                    "final_rank":        rank,
                }
            d = doc_data[did]
            d["chunks_in_final"]  += 1
            d["_score_sum"]       += score
            d["mean_chunk_score"]  = round(d["_score_sum"] / d["chunks_in_final"], 5)
            d["best_chunk_score"]  = round(max(d["best_chunk_score"], score), 5)
            d["retrieval_sources"].update(srcs)

        # Convert sets to lists; remove internal _score_sum
        for did, d in doc_data.items():
            d["retrieval_sources"] = sorted(d["retrieval_sources"])
            del d["_score_sum"]
        self.documents = doc_data

        # Final context metadata
        total_tokens = sum(len((c.get("text") or "").split()) for c in final_chunks)
        self.final_context = {
            "chunk_count":    len(final_chunks),
            "document_count": len(doc_data),
            "approx_tokens":  total_tokens,
            "chunk_ids":      [c.get("chunk_id") for c in final_chunks],
        }

        # Answer metadata (if provided)
        if answer_meta:
            self.answer = answer_meta

        self.stage_counts["final"] = len(final_chunks)

    # ── Serialization ─────────────────────────────────────────────────────────

    def to_log_dict(self) -> dict:
        """
        Compact CloudWatch log dict.
        Logged as a single JSON line with trace_type="retrieval_trace" for
        easy filtering in CloudWatch Logs Insights:
          filter trace_type = "retrieval_trace" | filter query_id = "LETA-..."
        """
        elapsed_ms = round((time.time() - self.created_at) * 1000)
        pre = self.preprocessing

        # Summarize what each stage produced
        stage_summary = {k: v for k, v in self.stage_counts.items()}

        # Per-stage top-3 chunk_ids + scores for quick scanning
        top_per_stage = {}
        for stage_name in ("faiss", "bm25", "rrf", "after_crossencoder",
                           "after_legalreranker", "post_mmr", "final"):
            snap = self._stages.get(stage_name, [])
            top_per_stage[stage_name] = [
                {"chunk_id": s["chunk_id"][:20], "doc": _doc_id(
                    (self._chunks.get(s["chunk_id"]) or ChunkRecord("","","","","")).rel_path
                ), "score": s["score"]}
                for s in snap[:3]
            ]

        # Selected chunk summary
        selected_chunks = [r.to_dict() for r in self._chunks.values() if r.selected]
        selected_summary = [
            {
                "chunk_id":        c["chunk_id"][:20],
                "document_id":     c["document_id"],
                "document_type":   c["document_type"],
                "faiss_rank":      c["faiss_rank"],
                "bm25_rank":       c["bm25_rank"],
                "rrf_score":       c["rrf_score"],
                "cross_encoder":   c["cross_encoder_score"],
                "legal_reranker":  c["legal_reranker_score"],
                "selected_by":     c["selection_reasons"],
            }
            for c in selected_chunks
        ]

        return {
            "trace_type":     "retrieval_trace",
            "query_id":       self.query_id,
            "query":          self.query[:200],
            "elapsed_ms":     elapsed_ms,
            "topic":          pre.get("topic"),
            "complexity":     pre.get("complexity_score"),
            "response_mode":  pre.get("response_mode"),
            "taxonomy_topics":pre.get("taxonomy", {}).get("topics", []),
            "taxonomy_conf":  pre.get("taxonomy", {}).get("confidence", 0),
            "domain_route":   pre.get("domain_route", [])[:4],
            "detected_refs":  pre.get("detected_refs", [])[:6],
            "stage_counts":   stage_summary,
            "top_per_stage":  top_per_stage,
            "selected_chunks":selected_summary,
            "documents_in_final": [
                {"doc_id": d["document_id"], "type": d["document_type"],
                 "chunks": d["chunks_in_final"], "best_score": d["best_chunk_score"]}
                for d in sorted(self.documents.values(),
                                key=lambda x: x["best_chunk_score"], reverse=True)
            ],
            "validation":      self.validation,
            "answer_model":    self.answer.get("model"),
            "answer_latency":  self.answer.get("latency_ms"),
            "cache_hit":       self.answer.get("cache_hit", False),
        }

    def to_debug_dict(self) -> dict:
        """Full detail for /debug/trace/{query_id}."""
        return {
            "query_id":      self.query_id,
            "query":         self.query,
            "created_at":    self.created_at,
            "elapsed_ms":    round((time.time() - self.created_at) * 1000),
            "preprocessing": self.preprocessing,
            "stage_counts":  self.stage_counts,
            "stages":        self._stages,
            "all_chunks":    [r.to_dict() for r in self._chunks.values()],
            "documents":     list(self.documents.values()),
            "validation":    self.validation,
            "final_context": self.final_context,
            "answer":        self.answer,
        }


# ─── In-memory ring buffer (debug endpoint storage) ───────────────────────────
# Stores the last 500 traces. Access via get_trace(query_id) and list_traces().

_RING_MAX = 500
_ring: deque = deque(maxlen=_RING_MAX)             # ordered list of query_ids (newest last)
_store: dict  = {}                                   # query_id → trace.to_debug_dict()
_store_lock   = threading.Lock()


def store_trace(trace: RetrievalTrace) -> None:
    """Persist trace to in-memory ring buffer. Called once per query."""
    try:
        with _store_lock:
            if len(_ring) == _RING_MAX:
                oldest = _ring[0]
                _store.pop(oldest, None)
            _ring.append(trace.query_id)
            _store[trace.query_id] = trace.to_debug_dict()
    except Exception as e:
        logger.debug(f"store_trace failed (non-fatal): {e}")


def get_trace(query_id: str) -> Optional[dict]:
    with _store_lock:
        return _store.get(query_id)


def list_traces(limit: int = 50) -> list:
    """Return summary of most recent traces (newest first)."""
    with _store_lock:
        ids = list(reversed(list(_ring)))[:limit]
        result = []
        for qid in ids:
            d = _store.get(qid, {})
            result.append({
                "query_id":    qid,
                "query":       d.get("query", "")[:120],
                "elapsed_ms":  d.get("elapsed_ms"),
                "topic":       d.get("preprocessing", {}).get("topic"),
                "final_count": d.get("stage_counts", {}).get("final", 0),
                "cache_hit":   d.get("answer", {}).get("cache_hit", False),
            })
        return result