"""
FailureTaxonomist — Tier 3: Diagnose WHY an agent chain failed.
===============================================================

Converts a raw risk score into an actionable failure mode using
the signal families from exp39-44:

  Behavioral  (exp39): step count, completion, hedging, answer gap
  Retrieval   (exp43): question↔observation cosine similarity
  Structural  (exp40): thought variance, answer alignment

Failure modes (non-exclusive — a chain can have multiple):

  PREMATURE_STOP      — finished after very few searches, likely guessed
  EXCESSIVE_SEARCH    — many more searches than typical; agent was struggling
  RETRIEVAL_FAILURE   — observations are irrelevant to the question
  CONFLICTING_EVIDENCE — high search diversity + high thought variance
  INSUFFICIENT_EVIDENCE — low mean relevance AND low n_relevant
  ANSWER_UNSUPPORTED  — answer words absent from reasoning chain
  LOW_RISK            — no significant failure signal detected

Usage:
    from qppg_service import FailureTaxonomist

    tx = FailureTaxonomist()
    result = tx.classify(question, steps, final_answer, finished)

    print(result.primary_mode)      # "EXCESSIVE_SEARCH"
    print(result.explanation)       # human-readable string
    print(result.signals)           # dict of raw signal values
    for mode in result.all_modes:   # non-exclusive list
        print(mode)
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import numpy as np

# ── Thresholds derived from exp39/43 empirical distributions ─────────────────

# Correct chains: avg 2.34 steps; incorrect: avg 4.75 — use 4 as threshold
_STEP_HIGH_THRESHOLD    = 4       # > this → likely struggling
_STEP_LOW_THRESHOLD     = 1       # <= this + not finished → premature stop

# Retrieval relevance thresholds (exp43)
_GOOD_MEAN_SIM          = 0.50    # mean cosine sim(question, obs) for good retrieval
_POOR_MEAN_SIM          = 0.35    # below this → retrieval failure
_GOOD_MIN_SIM           = 0.35    # min_sim for correct chains avg 0.456
_POOR_MIN_SIM           = 0.22    # below this → at least one useless search

# Thought variance (exp39 SC3)
_HIGH_THOUGHT_VAR       = 0.006   # above → conflicting reasoning

# Uncertainty words density (exp39 SC5)
_HIGH_UNCERTAINTY       = 0.5     # per 100 words

# Answer gap (exp39 SC6) — word-overlap version
_HIGH_ANSWER_GAP        = 0.70    # answer words absent from reasoning

# Answer-support embedding threshold (v0.54 improved ANSWER_UNSUPPORTED)
# Below this cosine sim between final_answer and concatenated reasoning → UNSUPPORTED
_ANSWER_SIM_LOW         = 0.30

# Observation contradiction: if second-half obs similarity to Q drops below first-half
# by this amount → agent pivoted away from the question (CONFLICTING_EVIDENCE signal)
_OBS_CONTRADICTION_DROP = 0.10

# Multi-hop question keywords: if these appear + few steps → stronger PREMATURE_STOP signal
_MULTI_HOP_WORDS        = frozenset([
    "after", "because", "both", "compare", "connection", "effect", "following",
    "how did", "impact", "relation", "role", "what caused", "which of",
])


@dataclass
class FailureMode:
    """Diagnosis result for a single agent chain."""

    primary_mode: str
    """The single most likely failure mode."""

    all_modes: list[str]
    """All detected failure modes (non-exclusive)."""

    confidence: float
    """Confidence in primary_mode [0,1]. Higher = stronger evidence."""

    explanation: str
    """Human-readable explanation of the failure."""

    signals: dict = field(default_factory=dict)
    """Raw signal values used to make the decision."""

    recommendation: str = ""
    """Actionable recommendation for the operator / self-healer."""

    @property
    def simplified_mode(self) -> str:
        """
        Production-safe mode for decision-making logic.

        Returns the primary_mode when it is one of the three validated modes:
          - ``"RETRIEVAL_FAILURE"`` — F1=0.701 (exp_failure_taxonomy_llm 2026-03-20)
          - ``"EXCESSIVE_SEARCH"``  — +20pp pass@1 (exp_selfhealer_ab 2026-03-20)
          - ``"LOW_RISK"``          — no-failure sentinel (always production-safe)

        Returns ``"OTHER"`` for all experimental modes that lack validation.
        Use this instead of primary_mode for any automated repair / alerting logic.
        primary_mode is available for diagnostic/logging purposes only.
        """
        _VALIDATED = {"RETRIEVAL_FAILURE", "EXCESSIVE_SEARCH", "LOW_RISK"}
        return self.primary_mode if self.primary_mode in _VALIDATED else "OTHER"

    def to_dict(self) -> dict:
        return {
            "primary_mode":   self.primary_mode,
            "simplified_mode": self.simplified_mode,
            "all_modes":      self.all_modes,
            "confidence":     round(self.confidence, 3),
            "explanation":    self.explanation,
            "recommendation": self.recommendation,
            "signals":        {k: round(v, 4) if isinstance(v, float) else v
                               for k, v in self.signals.items()},
        }


class FailureTaxonomist:
    """
    Rule-based failure mode classifier using exp39-44 behavioral + retrieval signals.

    No training required.  Classification is deterministic given the same chain.

    Parameters
    ----------
    embedder_model : str
        SentenceTransformer model for retrieval quality computation.
        Default: "all-MiniLM-L6-v2"

    LLM Agreement Study (exp_failure_taxonomy_llm, 2026-03-20)
    -----------------------------------------------------------
    Compared rule-based primary_mode against Haiku LLM annotation on 143 HP
    wrong chains (treated as oracle labels).  Per-mode F1:

      RETRIEVAL_FAILURE   : 0.701  ✓ (high agreement)
      LOW_RISK            : 0.396  (moderate)
      EXCESSIVE_SEARCH    : 0.000  (no agreement)
      ANSWER_UNSUPPORTED  : 0.046
      INSUFFICIENT_EVIDENCE: 0.000
      CONFLICTING_EVIDENCE: 0.000
      PREMATURE_STOP      : 0.000

    Macro F1 = 0.163, agreement_rate = 0.350.

    Conclusion: Only RETRIEVAL_FAILURE has validated LLM agreement.  Other modes
    are heuristic-only and may disagree substantially with human or LLM assessment.
    Treat non-RETRIEVAL_FAILURE labels as exploratory signals, not reliable diagnoses.

    Root cause of low agreement: rule thresholds over-assign LOW_RISK (88/143 rule
    labels) while LLM finds ANSWER_UNSUPPORTED (42) and INSUFFICIENT_EVIDENCE (39)
    as primary modes. Consider lowering thresholds for deployment diagnostics.
    """

    def __init__(self, embedder_model: str = "all-MiniLM-L6-v2",
                 experimental: bool = False):
        self._embedder_model   = embedder_model
        self._embedder         = None
        self._experimental     = experimental
        self._domain_thresholds: dict = {}
        # {domain: {poor_mean_sim, poor_min_sim, good_mean_sim, step_high, n_chains}}
        self._domain_fit_ts: float = 0.0

    def _get_embedder(self):
        if self._embedder is None:
            from sentence_transformers import SentenceTransformer
            self._embedder = SentenceTransformer(self._embedder_model)
        return self._embedder

    def fit_domain_thresholds(
        self,
        labeled_chains: list[dict],
        min_chains_per_domain: int = 50,
    ) -> dict:
        """
        Fit domain-specific RETRIEVAL_FAILURE thresholds from labeled chains.

        Called automatically by QppgMonitor every 50 new labeled chains (if
        `enable_domain_threshold_adaptation=True` on the monitor).  Can also
        be called manually.

        Parameters
        ----------
        labeled_chains : list of dicts
            Each dict must have:
              - ``domain`` : str  (e.g. "hotpotqa", "trivia", "medical")
              - ``steps``  : list of step dicts
              - ``correct``: bool  (True = agent answered correctly)
            Chains without a ``domain`` key are ignored.
        min_chains_per_domain : int
            Minimum chains in a domain before fitting a threshold (default 50).
            Domains below this threshold keep the global defaults.

        Returns
        -------
        dict  — summary of fitted domains and their thresholds.

        Notes
        -----
        Threshold derivation (validated on HP + NQ, exp43):
          ``poor_mean_sim`` ← 35th percentile of mean_sim for WRONG chains in domain
          ``poor_min_sim``  ← 25th percentile of min_sim  for WRONG chains in domain
          ``good_mean_sim`` ← 50th percentile of mean_sim for CORRECT chains in domain
          ``step_high``     ← 75th percentile of n_search_steps for WRONG chains

        Reverts to global defaults for domains with < min_chains_per_domain labeled chains.
        Only RETRIEVAL_FAILURE threshold is fitted; other thresholds stay global.
        """
        import time as _time

        embedder = self._get_embedder()

        # Group by domain
        by_domain: dict = {}
        for c in labeled_chains:
            domain = c.get("domain", "")
            if not domain:
                continue
            correct = bool(c.get("correct", True))
            steps   = c.get("steps", [])
            by_domain.setdefault(domain, {"correct": [], "wrong": []})
            by_domain[domain]["correct" if correct else "wrong"].append(steps)

        summary = {}
        for domain, groups in by_domain.items():
            wrong_steps   = groups["wrong"]
            correct_steps = groups["correct"]
            n_total = len(wrong_steps) + len(correct_steps)

            if n_total < min_chains_per_domain:
                summary[domain] = {"skipped": True, "n": n_total,
                                   "reason": f"< {min_chains_per_domain} chains"}
                continue

            def _mean_sims(steps_list: list) -> list:
                sims = []
                for steps in steps_list:
                    obs = [s.get("observation","").strip()[:300]
                           for s in steps if s.get("observation","").strip()]
                    q   = " ".join(
                        s.get("action_arg","") for s in steps
                        if s.get("action_type") == "Search"
                    )[:200]
                    if not obs or not q:
                        continue
                    q_emb  = embedder.encode([q], normalize_embeddings=True,
                                             show_progress_bar=False)[0]
                    o_embs = embedder.encode(obs, normalize_embeddings=True,
                                             show_progress_bar=False)
                    step_sims = (o_embs @ q_emb).tolist()
                    sims.append((float(np.mean(step_sims)), float(np.min(step_sims))))
                return sims  # list of (mean_sim, min_sim)

            wrong_sims   = _mean_sims(wrong_steps)
            correct_sims = _mean_sims(correct_steps)

            if not wrong_sims:
                summary[domain] = {"skipped": True, "n": n_total,
                                   "reason": "no wrong chains with observations"}
                continue

            wrong_means = [s[0] for s in wrong_sims]
            wrong_mins  = [s[1] for s in wrong_sims]
            wrong_counts = [sum(1 for s in steps if s.get("action_type") == "Search")
                            for steps in wrong_steps]
            correct_means = [s[0] for s in correct_sims] if correct_sims else wrong_means

            poor_mean = float(np.percentile(wrong_means,  35))
            poor_min  = float(np.percentile(wrong_mins,   25))
            good_mean = float(np.percentile(correct_means, 50)) if correct_means else _GOOD_MEAN_SIM
            step_high = int(np.percentile(wrong_counts, 75)) if wrong_counts else _STEP_HIGH_THRESHOLD

            # Sanity clamp: keep thresholds in [0.10, 0.60]
            poor_mean = max(0.10, min(0.60, poor_mean))
            poor_min  = max(0.08, min(0.50, poor_min))
            good_mean = max(poor_mean + 0.05, min(0.80, good_mean))
            step_high = max(2, min(12, step_high))

            self._domain_thresholds[domain] = {
                "poor_mean_sim": poor_mean,
                "poor_min_sim":  poor_min,
                "good_mean_sim": good_mean,
                "step_high":     step_high,
                "n_chains":      n_total,
                "n_wrong":       len(wrong_steps),
            }
            summary[domain] = {"fitted": True, **self._domain_thresholds[domain]}

        self._domain_fit_ts = _time.time()
        return summary

    def domain_threshold_summary(self) -> dict:
        """Return currently fitted domain thresholds + global defaults."""
        return {
            "global_defaults": {
                "poor_mean_sim": _POOR_MEAN_SIM,
                "poor_min_sim":  _POOR_MIN_SIM,
                "good_mean_sim": _GOOD_MEAN_SIM,
                "step_high":     _STEP_HIGH_THRESHOLD,
            },
            "fitted_domains": dict(self._domain_thresholds),
            "n_fitted_domains": len(self._domain_thresholds),
            "last_fit_ts": self._domain_fit_ts,
        }

    def classify(
        self,
        question: str,
        steps: list[dict],
        final_answer: str = "",
        finished: bool = True,
        domain: str = "",
    ) -> FailureMode:
        """
        Classify the failure mode of a completed agent chain.

        Parameters
        ----------
        question     : original question
        steps        : agent steps (thought/action_type/action_arg/observation)
        final_answer : agent's final answer
        finished     : True if agent called Finish[], False if fallback

        Returns
        -------
        FailureMode with primary_mode, all_modes, explanation, recommendation
        """
        import re

        # ── Domain-specific thresholds (fallback to global defaults) ──────────
        _dt = self._domain_thresholds.get(domain, {}) if domain else {}
        _poor_mean_sim  = _dt.get("poor_mean_sim", _POOR_MEAN_SIM)
        _poor_min_sim   = _dt.get("poor_min_sim",  _POOR_MIN_SIM)
        _good_mean_sim  = _dt.get("good_mean_sim", _GOOD_MEAN_SIM)
        _step_high      = _dt.get("step_high",     _STEP_HIGH_THRESHOLD)

        # ── Raw signals ───────────────────────────────────────────────────────
        search_steps = [s for s in steps if s.get("action_type") == "Search"]
        n_steps      = len(search_steps)

        # Thought variance (SC3)
        embedder = self._get_embedder()
        thoughts = [s.get("thought", "").strip() for s in steps if s.get("thought","").strip()]
        thought_var = 0.0
        if len(thoughts) >= 2:
            embs = embedder.encode(thoughts, normalize_embeddings=True, show_progress_bar=False)
            thought_var = float(embs.var(axis=0).mean())

        # Uncertainty words density (SC5)
        _UW = ["not sure","unclear","might","may not","possibly","perhaps","however",
               "although","uncertain","seems","appears","probably","could be","i think",
               "i believe","not certain","unsure","i'm not","don't know","confusing",
               "it's possible","not clear","ambiguous"]
        all_t = " ".join(s.get("thought","") for s in steps).lower()
        wc = max(1, len(all_t.split()))
        uncertainty = sum(all_t.count(w) for w in _UW) / wc * 100

        # Answer gap (SC6) — compare final answer against full chain text (thoughts + observations)
        # v0.54: include observations so answers that appear in retrieved text aren't penalised
        _STOP = {"the","a","an","is","it","in","of","to","and","or","that","this",
                 "was","are","for","with","on","at","by","be","has","had","i"}
        all_chain_text = " ".join(
            (s.get("thought","") + " " + s.get("observation",""))
            for s in steps
        ).lower()
        fa_words  = set(re.findall(r'\w+', final_answer.lower())) - _STOP
        chain_words = set(re.findall(r'\w+', all_chain_text)) - _STOP
        th_words  = set(re.findall(r'\w+', all_t)) - _STOP      # kept for backward-compat signals
        answer_gap = (1.0 - len(fa_words & chain_words) / len(fa_words)) if fa_words else 0.5

        # Retrieval quality (exp43)
        obs_list = [s.get("observation","").strip()[:300]
                    for s in steps if s.get("observation","").strip()]
        mean_sim = max_sim = min_sim = 0.0
        obs_contradiction = 0.0   # drop in relevance between early and late observations
        answer_sim = 0.0          # embedding similarity: final_answer ↔ reasoning chain
        if obs_list:
            q_emb    = embedder.encode([question], normalize_embeddings=True,
                                        show_progress_bar=False)[0]
            obs_embs = embedder.encode(obs_list, normalize_embeddings=True,
                                        show_progress_bar=False)
            sims     = (obs_embs @ q_emb).tolist()
            mean_sim = float(np.mean(sims))
            max_sim  = float(np.max(sims))
            min_sim  = float(np.min(sims))
            # Observation contradiction: compare first-half vs second-half similarity to Q
            if len(sims) >= 4:
                mid = len(sims) // 2
                early_mean = float(np.mean(sims[:mid]))
                late_mean  = float(np.mean(sims[mid:]))
                obs_contradiction = max(0.0, early_mean - late_mean)

        # Answer-support: embedding similarity between final_answer and full chain text
        # v0.54: use all_chain_text (thoughts + observations) so retrieved evidence counts
        if final_answer and all_chain_text:
            ans_emb  = embedder.encode([final_answer[:500]], normalize_embeddings=True,
                                        show_progress_bar=False)[0]
            reas_emb = embedder.encode([all_chain_text[:1000]], normalize_embeddings=True,
                                        show_progress_bar=False)[0]
            answer_sim = float(ans_emb @ reas_emb)

        # ── Failure mode detection ─────────────────────────────────────────────
        modes = []
        evidences = {}

        # PREMATURE_STOP: very few searches + didn't finish properly OR finished too quickly
        # Improved (v0.54): detect multi-hop questions with few searches more aggressively
        q_lower = question.lower()
        is_multi_hop = any(kw in q_lower for kw in _MULTI_HOP_WORDS)
        if n_steps <= _STEP_LOW_THRESHOLD and not finished:
            modes.append("PREMATURE_STOP")
            evidences["premature_stop"] = f"only {n_steps} search(es), no clean finish"
        elif n_steps == 0 and finished:
            modes.append("PREMATURE_STOP")
            evidences["premature_stop"] = "finished with zero searches — likely guessed"
        elif n_steps == 1 and finished and is_multi_hop:
            # Multi-hop question answered after a single search is suspicious
            modes.append("PREMATURE_STOP")
            evidences["premature_stop"] = (
                f"multi-hop question answered after only 1 search — likely under-retrieved"
            )

        # EXCESSIVE_SEARCH: too many steps
        if n_steps >= _step_high:
            modes.append("EXCESSIVE_SEARCH")
            evidences["excessive_search"] = f"{n_steps} search steps (threshold: {_step_high})"

        # RETRIEVAL_FAILURE: observations are irrelevant to the question
        if obs_list and mean_sim < _poor_mean_sim:
            modes.append("RETRIEVAL_FAILURE")
            evidences["retrieval_failure"] = (f"mean_sim={mean_sim:.3f} < {_poor_mean_sim:.3f} threshold; "
                                              f"search results appear off-topic")
        elif obs_list and min_sim < _poor_min_sim:
            modes.append("RETRIEVAL_FAILURE")
            evidences["retrieval_failure"] = (f"min_sim={min_sim:.3f} — at least one "
                                              f"observation irrelevant to question")

        # INSUFFICIENT_EVIDENCE: searched but results weakly relevant
        # Improved (v0.54): also fire when high uncertainty density accompanies weak retrieval
        if (obs_list and _poor_mean_sim <= mean_sim < _good_mean_sim
                and n_steps >= 2 and "RETRIEVAL_FAILURE" not in modes):
            modes.append("INSUFFICIENT_EVIDENCE")
            unc_note = f", uncertainty_density={uncertainty:.2f}%" if uncertainty > _HIGH_UNCERTAINTY else ""
            evidences["insufficient_evidence"] = (
                f"mean_sim={mean_sim:.3f} — evidence found but weakly relevant; "
                f"agent may need more or better searches{unc_note}"
            )
        elif (uncertainty > _HIGH_UNCERTAINTY * 2 and n_steps >= 2
              and "RETRIEVAL_FAILURE" not in modes and "CONFLICTING_EVIDENCE" not in modes):
            # High hedge density without retrieval failure = likely insufficient evidence
            modes.append("INSUFFICIENT_EVIDENCE")
            evidences["insufficient_evidence"] = (
                f"uncertainty_density={uncertainty:.2f}% — agent expresses high uncertainty "
                f"despite {n_steps} searches, suggesting evidence was insufficient"
            )

        # CONFLICTING_EVIDENCE: high thought variance AND (high search diversity OR obs contradiction)
        # Improved (v0.54): add observation contradiction signal (late obs less relevant than early)
        search_queries = [s.get("action_arg","") for s in search_steps]
        query_diversity = 0.0
        if len(search_queries) >= 2:
            q_embs = embedder.encode(search_queries, normalize_embeddings=True, show_progress_bar=False)
            q_sim  = q_embs @ q_embs.T
            n = len(search_queries)
            dists = [1.0 - q_sim[i,j] for i in range(n) for j in range(i+1,n)]
            query_diversity = float(np.mean(dists)) if dists else 0.0

        if thought_var > _HIGH_THOUGHT_VAR and query_diversity > 0.5:
            modes.append("CONFLICTING_EVIDENCE")
            evidences["conflicting_evidence"] = (
                f"thought_var={thought_var:.4f} (high) AND query_diversity={query_diversity:.3f} "
                f"— agent repeatedly pivoting, evidence may be contradictory")
        elif thought_var > _HIGH_THOUGHT_VAR and obs_contradiction > _OBS_CONTRADICTION_DROP:
            # Late observations drifted away from the question — contradicting earlier retrieval
            modes.append("CONFLICTING_EVIDENCE")
            evidences["conflicting_evidence"] = (
                f"thought_var={thought_var:.4f} AND obs_contradiction={obs_contradiction:.3f} "
                f"— later observations less relevant, earlier evidence contradicted")

        # ANSWER_UNSUPPORTED: final answer not grounded in reasoning
        # Improved (v0.54): combine word-overlap gap WITH embedding similarity
        # Both signals must agree to reduce FP rate (v0.53 word-only F1=0.046)
        ans_emb_low = (answer_sim > 0.0 and answer_sim < _ANSWER_SIM_LOW)
        if answer_gap > _HIGH_ANSWER_GAP and fa_words and (ans_emb_low or answer_sim == 0.0):
            modes.append("ANSWER_UNSUPPORTED")
            evidences["answer_unsupported"] = (
                f"answer_gap={answer_gap:.3f} AND answer_sim={answer_sim:.3f} — "
                f"final answer semantically distant from reasoning chain "
                f"(agent may have introduced information not from retrieved evidence)")

        # LOW_RISK: no failure signals
        if not modes:
            modes.append("LOW_RISK")

        # ── Priority ranking for primary_mode ─────────────────────────────────
        priority = ["RETRIEVAL_FAILURE", "CONFLICTING_EVIDENCE", "EXCESSIVE_SEARCH",
                    "INSUFFICIENT_EVIDENCE", "ANSWER_UNSUPPORTED", "PREMATURE_STOP", "LOW_RISK"]
        primary = next((m for m in priority if m in modes), modes[0])

        # ── Production-safe mode (non-experimental) ───────────────────────────
        # Validated modes (production-safe):
        #   RETRIEVAL_FAILURE: F1=0.701 (exp_failure_taxonomy_llm, 2026-03-20)
        #   EXCESSIVE_SEARCH:  +20pp SelfHealer improvement (exp_selfhealer_ab, 2026-03-20)
        #   LOW_RISK: pass-through — "no failure detected" is always production-safe.
        # All other failure modes collapse to "OTHER" unless experimental=True.
        if not self._experimental:
            _validated = {"RETRIEVAL_FAILURE", "EXCESSIVE_SEARCH", "LOW_RISK"}
            primary = primary if primary in _validated else "OTHER"
            modes   = [m if m in _validated else "OTHER" for m in modes]
            modes   = list(dict.fromkeys(modes))  # deduplicate, preserve order
            if not modes:
                modes = ["OTHER"]

        # ── Confidence ────────────────────────────────────────────────────────
        # More modes = higher confidence something is wrong; use n_modes as proxy
        n_failure_modes = len([m for m in modes if m != "LOW_RISK"])
        confidence = min(1.0, 0.5 + 0.15 * n_failure_modes)
        if primary == "LOW_RISK":
            confidence = 0.85   # high confidence nothing is wrong

        # ── Explanation + recommendation ──────────────────────────────────────
        explanation = _build_explanation(primary, evidences, n_steps, mean_sim, finished)
        recommendation = _build_recommendation(modes, n_steps, mean_sim, finished, thought_var)

        return FailureMode(
            primary_mode    = primary,
            all_modes       = modes,
            confidence      = confidence,
            explanation     = explanation,
            recommendation  = recommendation,
            signals = {
                "n_search_steps":       n_steps,
                "finished":             finished,
                "thought_variance":     round(thought_var, 5),
                "uncertainty_density":  round(uncertainty, 3),
                "answer_gap":           round(answer_gap, 3),
                "answer_sim":           round(answer_sim, 4),    # v0.54: embedding similarity
                "mean_sim":             round(mean_sim, 4),
                "max_sim":              round(max_sim, 4),
                "min_sim":              round(min_sim, 4),
                "query_diversity":      round(query_diversity, 3),
                "obs_contradiction":    round(obs_contradiction, 4),  # v0.54
                "is_multi_hop":         is_multi_hop,                 # v0.54
            },
        )

    def classify_batch(self, chains: list[dict]) -> list[FailureMode]:
        """Classify a batch of chains."""
        return [
            self.classify(
                c["question"],
                c.get("steps", []),
                c.get("final_answer", ""),
                c.get("finished", True),
            )
            for c in chains
        ]

    def discover_failure_modes(
        self,
        chains: list[dict],
        risk_threshold: float = 0.7,
    ) -> list[dict]:
        """Auto-discover dominant failure modes in chains above risk_threshold.

        Heuristic analysis (zero LLM cost). Returns failure mode dicts sorted by count.

        Args:
            chains: list of chain dicts with keys 'steps', 'final_answer', and optionally 'risk'.
            risk_threshold: only analyze chains with risk >= this value.
        Returns:
            list of dicts: [{"mode": str, "count": int, "chain_indices": list[int]}, ...]
        """
        high_risk = [
            (i, c) for i, c in enumerate(chains)
            if c.get("risk", 0.0) >= risk_threshold
        ]

        buckets: dict[str, list[int]] = {
            "EMPTY_ANSWER": [],
            "EXCESSIVE_SEARCH": [],
            "RETRIEVAL_FAILURE": [],
            "SHORT_REASONING": [],
            "OTHER": [],
        }

        for idx, chain in high_risk:
            steps = chain.get("steps", [])
            answer = chain.get("final_answer", "")
            n_search = sum(
                1 for s in steps
                if "search" in str(s.get("action", "")).lower()
            )

            if not answer or len(answer.strip()) < 5:
                buckets["EMPTY_ANSWER"].append(idx)
            elif n_search > 6:
                buckets["EXCESSIVE_SEARCH"].append(idx)
            elif len(steps) < 2:
                buckets["SHORT_REASONING"].append(idx)
            elif n_search > 0 and len(answer) < 20:
                buckets["RETRIEVAL_FAILURE"].append(idx)
            else:
                buckets["OTHER"].append(idx)

        result = [
            {"mode": mode, "count": len(indices), "chain_indices": indices}
            for mode, indices in buckets.items()
            if indices
        ]
        return sorted(result, key=lambda x: x["count"], reverse=True)


# ── Helper functions ───────────────────────────────────────────────────────────

def _build_explanation(primary: str, evidences: dict, n_steps: int,
                       mean_sim: float, finished: bool) -> str:
    if primary == "LOW_RISK":
        return "No significant failure signals detected. Answer appears well-supported."
    explanations = {
        "RETRIEVAL_FAILURE": (
            f"The search engine returned results with low relevance to the question "
            f"(mean similarity {mean_sim:.2f}). The agent may have used poor search queries "
            f"or the topic is outside the retrieval system's coverage."
        ),
        "CONFLICTING_EVIDENCE": (
            f"The agent's reasoning pivoted significantly across {n_steps} searches, "
            f"suggesting the retrieved evidence was contradictory or the agent was confused."
        ),
        "EXCESSIVE_SEARCH": (
            f"The agent performed {n_steps} searches, more than typical for correct chains "
            f"(avg 2.3 for correct vs 4.75 for incorrect). Excessive searching correlates "
            f"strongly with incorrect final answers. Validated: SelfHealer +20pp (exp_selfhealer_ab)."
        ),
        "INSUFFICIENT_EVIDENCE": (
            f"Evidence was found but was only weakly relevant (mean sim {mean_sim:.2f}). "
            f"The agent may need more or better searches to answer confidently."
        ),
        "ANSWER_UNSUPPORTED": (
            f"The final answer contains words not present in the reasoning chain, "
            f"suggesting the agent may have introduced information not in the retrieved evidence."
        ),
        "PREMATURE_STOP": (
            f"The agent stopped with {n_steps} search(es) and "
            f"{'a clean finish' if finished else 'no clean finish'}. "
            f"This suggests the agent guessed without sufficient evidence."
        ),
    }
    base = explanations.get(primary, "Unknown failure mode.")
    if len(evidences) > 1:
        other = [v for k, v in evidences.items() if k != primary.lower()]
        if other:
            base += f" Additional signals: {'; '.join(other[:2])}."
    return base


def _build_recommendation(modes: list[str], n_steps: int, mean_sim: float,
                           finished: bool, thought_var: float) -> str:
    recs = []
    if "RETRIEVAL_FAILURE" in modes:
        recs.append("rephrase search queries to be more specific to the question domain")
    if "EXCESSIVE_SEARCH" in modes:
        recs.append("inject consolidation prompt: 'Based on what you've found so far, what is the best answer?'")
    if "INSUFFICIENT_EVIDENCE" in modes:
        recs.append("try alternative search queries or increase max_steps")
    if "CONFLICTING_EVIDENCE" in modes:
        recs.append("inject reasoning prompt: 'Reconcile the conflicting evidence you found'")
    if "ANSWER_UNSUPPORTED" in modes:
        recs.append("request answer verification: 'Cite which observation supports your answer'")
    if "PREMATURE_STOP" in modes:
        recs.append("force at least one search before allowing Finish[]")
    if not recs or modes == ["LOW_RISK"]:
        return "No intervention needed."
    return " | ".join(recs)
