"""
LLM Reliability Guard
=====================
Predicts, diagnoses, and repairs LLM failures automatically.

Core method: KNN anomaly scoring on sentence-transformer embeddings.
  - Calibrate on known-correct examples (or auto-calibrate without labels)
  - Score every query before the LLM responds — higher = more likely to fail
  - Latency overhead: <15ms per query after warm-up

  Note on AUROC: this class uses KNN on query embeddings, which generalises
  well within-distribution but degrades cross-domain. Earlier docs cited
  0.966–0.993 measured in-sample on the same calibration data — that figure
  is not representative of held-out performance.
  For validated held-out AUROC (~0.81 within / ~0.74 cross), use AgentGuard
  from llm_guard.agent_guard, which uses SC_OLD behavioral signals + optional
  Sonnet judge (exp88/89/91).

Three calibration paths:
  guard.fit(correct_questions)               # you have labeled correct examples
  guard.fit_from_consistency(questions)      # no labels: uses self-consistency
  guard.fit_from_execution(questions, fn)    # no labels: uses a verifier function

Usage:
    from qppg.guard import QPPGLLMGuard

    guard = QPPGLLMGuard(api_key="sk-ant-...")

    # Path 1 — you have a golden set
    guard.fit(correct_questions=[...])

    # Path 2 — no labels, let self-consistency decide
    guard.fit_from_consistency(questions=[...], n_samples=5)

    # Path 3 — automated verifier (math, code, SQL, schema)
    guard.fit_from_execution(questions=[...],
                             verifier_fn=lambda q, r: check(q, r))

    # Query with reliability scoring + automatic repair
    result = guard.query("What is 15% of 240?")
    print(result.answer)       # "36"
    print(result.risk_score)   # 0.12   (low = familiar territory)
    print(result.confidence)   # "high"

    # Diagnose accumulated failures
    clusters = guard.diagnose(failed_qs, model_answers, correct_answers)

    # Synthesise targeted repair tools from failure patterns
    guard.learn_from_errors(failed_qs, model_answers, correct_answers)
    # Subsequent queries in a known-failure cluster get the tool auto-applied
"""

import os
import pickle
import numpy as np
from collections import Counter
from typing import Callable, Dict, List, Optional
from dataclasses import dataclass, field


# ── QARA: cross-domain adapter (train-time torch, inference-time numpy) ────────

@dataclass
class _QARAWeights:
    """Portable numpy weights for the QARA 2-layer MLP adapter."""
    w1:   np.ndarray   # (384, 256)
    b1:   np.ndarray   # (256,)
    ln_w: np.ndarray   # (256,)
    ln_b: np.ndarray   # (256,)
    w2:   np.ndarray   # (256, 64)
    b2:   np.ndarray   # (64,)


def _qara_supcon_loss(z, labels_correct, temp=0.1):
    """Supervised contrastive loss: correct-correct pairs are positives."""
    import torch
    pos_mask = labels_correct.unsqueeze(1) * labels_correct.unsqueeze(0)
    pos_mask.fill_diagonal_(0)
    sim = (z @ z.T).clamp(-20, 20) / temp
    sim_masked = sim.clone()
    sim_masked.fill_diagonal_(float("-inf"))
    log_probs = sim_masked - torch.logsumexp(sim_masked, dim=1, keepdim=True)
    log_probs_pos = log_probs.clone()
    log_probs_pos[pos_mask == 0] = 0.0   # avoid 0 × (−inf) = NaN
    n_pos = pos_mask.sum(dim=1).clamp(min=1)
    return (-log_probs_pos.sum(dim=1) / n_pos).mean()


def _qara_train(embs: np.ndarray, labels: np.ndarray,
                epochs: int = 200, lr: float = 3e-4,
                temp: float = 0.1, verbose: bool = True) -> "_QARAWeights":
    """
    Train 384→256→64 MLP adapter with SupCon loss.
    Returns portable _QARAWeights (numpy only, no torch at inference time).
    """
    import torch
    import torch.nn as nn

    class _Adapter(nn.Module):
        def __init__(self):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(384, 256), nn.ReLU(), nn.LayerNorm(256),
                nn.Linear(256, 64),
            )
        def forward(self, x):
            z = self.net(x)
            return z / (z.norm(dim=1, keepdim=True) + 1e-8)

    X = torch.tensor(embs, dtype=torch.float32)
    y = torch.tensor(labels, dtype=torch.float32)
    adapter = _Adapter()
    opt = torch.optim.Adam(adapter.parameters(), lr=lr)

    for ep in range(epochs):
        adapter.train()
        z = adapter(X)
        loss = _qara_supcon_loss(z, y, temp=temp)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
        opt.step()
        if verbose and (ep + 1) % 50 == 0:
            print(f"  QARA epoch {ep+1}/{epochs}  loss={loss.item():.4f}")

    adapter.eval()
    sd = adapter.state_dict()
    return _QARAWeights(
        w1   = sd["net.0.weight"].detach().numpy().T,  # (384, 256)
        b1   = sd["net.0.bias"].detach().numpy(),
        ln_w = sd["net.2.weight"].detach().numpy(),
        ln_b = sd["net.2.bias"].detach().numpy(),
        w2   = sd["net.3.weight"].detach().numpy().T,  # (256, 64)
        b2   = sd["net.3.bias"].detach().numpy(),
    )


# ── Result dataclass ───────────────────────────────────────────────────────────

@dataclass
class GuardResult:
    """Result of a guarded LLM query."""
    answer: str
    risk_score: float        # higher = more likely to fail
    confidence: str          # "high" | "medium" | "low"
    tool_used: Optional[str] = None
    cluster_id: Optional[int] = None
    was_retried: bool = False   # True if a resource-failure retry fired
    raw_response: str = ""

    # Backward-compat alias used by older code
    @property
    def blindness_score(self) -> float:
        return self.risk_score


# ── Guard class ────────────────────────────────────────────────────────────────

class QPPGLLMGuard:
    """
    Drop-in reliability layer for any LLM API call.

    Predicts query failure risk using KNN anomaly scoring:
      1. Embed the query with sentence-transformers (<10ms)
      2. Compute mean distance to bank of known-correct query embeddings
      3. High distance = unfamiliar territory = likely failure

    Routing based on risk score (thresholds auto-calibrated from training data):
      Low risk   → direct LLM call                        (confidence: "high")
      Medium risk → apply cluster repair tool if available (confidence: "medium")
      High risk  → call with explicit uncertainty flag     (confidence: "low")

    Failure-type detection (applied at medium/high risk):
      Resource failure (hit max_tokens) → retry with 2x tokens, no tool
      Reasoning failure (wrong but complete) → apply synthesised cluster tool

    Parameters
    ----------
    api_key : str
        Anthropic API key (falls back to ANTHROPIC_API_KEY env var).
    model : str
        Claude model to use for LLM calls.
    embedding_model : str
        sentence-transformers model name (default: all-MiniLM-L6-v2).
    n_neighbors : int
        k for KNN anomaly scoring (default 5).
    """

    def __init__(
        self,
        api_key: str = None,
        model: str = "claude-haiku-4-5-20251001",
        embedding_model: str = "all-MiniLM-L6-v2",
        n_neighbors: int = 5,
    ):
        self.model = model
        self.n_neighbors = n_neighbors

        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self._client = None
        self._embedder = None
        self._embedding_model_name = embedding_model

        # KNN state — set during any fit_* call
        self._knn = None
        self._risk_low_threshold: Optional[float] = None   # auto-calibrated
        self._risk_high_threshold: Optional[float] = None  # auto-calibrated
        self._fitted = False
        self._cal_embs: Optional[np.ndarray] = None        # raw calibration embeddings (for QARA re-fit)

        # QARA adapter — set by fit_qara()
        self._qara: Optional[_QARAWeights] = None

        # Cluster repair tool library — set during learn_from_errors
        self._tools: Dict[str, dict] = {}
        self._cluster_centers: Optional[np.ndarray] = None
        self._tool_usage: Dict[str, int] = {}

        # Usage tracking
        self.total_calls = 0
        self.total_input_tokens = 0
        self.total_output_tokens = 0

    # ── Private: infrastructure ────────────────────────────────────────────────

    def _get_client(self):
        if self._client is None:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self._api_key)
        return self._client

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

    def _embed(self, texts: List[str]) -> np.ndarray:
        """Embed a list of texts; returns (N, D) float32 array."""
        return self._get_embedder().encode(
            texts, normalize_embeddings=True, show_progress_bar=False
        )

    def _adapt_embs(self, embs: np.ndarray) -> np.ndarray:
        """Apply the QARA adapter (numpy forward pass). No-op if not fitted."""
        if self._qara is None or len(embs) == 0:
            return embs
        q = self._qara
        x = embs @ q.w1 + q.b1                                          # (N, 256)
        x = np.maximum(x, 0)                                             # ReLU
        mu  = x.mean(axis=1, keepdims=True)
        var = x.var(axis=1, keepdims=True)
        x = (x - mu) / np.sqrt(var + 1e-5) * q.ln_w + q.ln_b           # LayerNorm
        x = x @ q.w2 + q.b2                                             # (N, 64)
        return x / (np.linalg.norm(x, axis=1, keepdims=True) + 1e-8)   # L2-norm

    def _call_llm(
        self,
        system_prompt: str,
        user_prompt: str,
        max_tokens: int = 500,
        temperature: float = 0.0,
    ):
        """Call the LLM. Returns (text, stop_reason)."""
        client = self._get_client()
        resp = client.messages.create(
            model=self.model,
            max_tokens=max_tokens,
            temperature=temperature,
            system=system_prompt,
            messages=[{"role": "user", "content": user_prompt}],
        )
        self.total_calls += 1
        self.total_input_tokens += resp.usage.input_tokens
        self.total_output_tokens += resp.usage.output_tokens
        return resp.content[0].text, resp.stop_reason

    # ── Private: KNN core ──────────────────────────────────────────────────────

    def _fit_knn(self, embeddings: np.ndarray) -> None:
        """
        Fit KNN on correct-example embeddings and auto-calibrate risk thresholds.

        Stores raw embeddings so fit_qara() can re-fit after training the adapter.
        If QARA adapter is already set, fits in adapted (64-d) space instead.

        Thresholds are derived from the training distribution so they work
        regardless of domain, dataset size, or embedding space scale:
          risk_low_threshold  = 75th percentile of within-correct-set distances
          risk_high_threshold = 95th percentile of within-correct-set distances
        """
        from sklearn.neighbors import NearestNeighbors

        self._cal_embs = embeddings          # save raw for re-fit after fit_qara()
        adapted = self._adapt_embs(embeddings)

        k = min(self.n_neighbors, len(adapted) - 1)
        self._knn = NearestNeighbors(n_neighbors=k, metric="euclidean")
        self._knn.fit(adapted)

        # Self-distances of training examples give us the baseline distribution.
        dists, _ = self._knn.kneighbors(adapted)
        scores = dists.mean(axis=1)
        self._risk_low_threshold  = float(np.percentile(scores, 75))
        self._risk_high_threshold = float(np.percentile(scores, 95))
        self._fitted = True

    def _compute_risk_score(self, question: str) -> float:
        """Mean KNN distance to correct-example bank. Higher = more likely to fail."""
        if not self._fitted:
            return 0.5
        emb = self._embed([question])
        emb = self._adapt_embs(emb)
        dists, _ = self._knn.kneighbors(emb)
        return float(dists.mean())

    def _match_tool(self, emb: np.ndarray) -> Optional[dict]:
        """Find the best cluster repair tool for a query embedding."""
        if not self._tools or self._cluster_centers is None:
            return None
        dists = np.linalg.norm(self._cluster_centers - emb[0], axis=1)
        best_idx = int(np.argmin(dists))
        for tool in self._tools.values():
            if tool.get("cluster_idx") == best_idx:
                return tool
        return None

    # ── Calibration paths ──────────────────────────────────────────────────────

    def fit(
        self,
        correct_questions: List[str],
        correct_answers: Optional[List[str]] = None,
        correct_responses: Optional[List[str]] = None,
    ) -> "QPPGLLMGuard":
        """
        Calibrate on known-correct examples.

        Use this when you already have a golden/labeled dataset.

        Parameters
        ----------
        correct_questions : list of str
            Questions the LLM is known to answer correctly.
        correct_answers : list of str, optional
            Ground truth answers (not used for fitting; reserved for future use).
        correct_responses : list of str, optional
            Full LLM response texts (not used currently).
        """
        emb = self._embed(correct_questions)
        self._fit_knn(emb)
        return self

    def fit_from_consistency(
        self,
        questions: List[str],
        n_samples: int = 5,
        system_prompt: str = None,
        max_tokens: int = 500,
        agreement_threshold: float = 0.8,
    ) -> "QPPGLLMGuard":
        """
        Calibrate without labeled data using self-consistency.

        Runs each question n_samples times at temperature=0.7.
        Questions where >= agreement_threshold of responses agree are treated as
        "probably correct" and used to build the KNN reference bank.

        Best suited to domains where model accuracy > ~70%.

        Parameters
        ----------
        questions : list of str
            Pool of questions to evaluate (100–500 recommended).
        n_samples : int
            Samples per question (default 5; higher = more reliable).
        system_prompt : str, optional
            System prompt for calibration calls.
        max_tokens : int
            Max tokens per sample call.
        agreement_threshold : float
            Fraction of samples that must agree (default 0.8 = 80%).
        """
        if system_prompt is None:
            system_prompt = "You are a helpful assistant. Answer concisely."

        def _short(text: str) -> str:
            """Normalise to a short canonical form for agreement comparison."""
            lines = [l.strip() for l in text.strip().splitlines() if l.strip()]
            return (lines[-1] if lines else text.strip()).lower()[:100]

        print(
            f"fit_from_consistency: sampling {len(questions)} questions "
            f"x{n_samples} at temperature=0.7 ..."
        )

        probably_correct: List[str] = []
        for i, q in enumerate(questions):
            responses = []
            for _ in range(n_samples):
                text, _ = self._call_llm(
                    system_prompt, q, max_tokens=max_tokens, temperature=0.7
                )
                responses.append(_short(text))

            top_count = Counter(responses).most_common(1)[0][1]
            if top_count / n_samples >= agreement_threshold:
                probably_correct.append(q)

            if (i + 1) % 20 == 0:
                print(f"  [{i+1}/{len(questions)}] agreed: {len(probably_correct)}")

        n_agreed = len(probably_correct)
        print(
            f"  Calibration set: {n_agreed}/{len(questions)} questions "
            f"({n_agreed/len(questions):.1%} agreement rate)"
        )

        if n_agreed < self.n_neighbors + 1:
            raise ValueError(
                f"Only {n_agreed} consistently-answered questions found. "
                f"Need at least {self.n_neighbors + 1}. "
                "Try lowering agreement_threshold or supplying more questions."
            )

        emb = self._embed(probably_correct)
        self._fit_knn(emb)
        return self

    def fit_from_execution(
        self,
        questions: List[str],
        verifier_fn: Callable[[str, str], bool],
        system_prompt: str = None,
        max_tokens: int = 500,
    ) -> "QPPGLLMGuard":
        """
        Calibrate without labeled data using an automated verifier.

        Runs each question through the LLM, calls verifier_fn(question, response)
        to determine correctness, then fits KNN on the correctly-answered questions.

        Ideal for domains with automated correctness checks:
          - Math:   verifier checks eval(response) == expected_value
          - Code:   verifier runs unit tests against the response
          - SQL:    verifier executes and validates the result set
          - Schema: verifier validates JSON / XML structure

        Parameters
        ----------
        questions : list of str
            Questions to run through the LLM.
        verifier_fn : callable(question: str, response: str) -> bool
            Returns True if the response is correct.
        system_prompt : str, optional
            System prompt for calibration calls.
        max_tokens : int
            Max tokens per call.

        Example
        -------
        def math_verifier(question, response):
            import re
            nums = re.findall(r'-?\\d+\\.?\\d*', response)
            return nums and nums[-1] == expected_answers[question]

        guard.fit_from_execution(questions, math_verifier)
        """
        if system_prompt is None:
            system_prompt = "You are a helpful assistant. Answer concisely."

        print(f"fit_from_execution: running {len(questions)} questions ...")

        correct_questions: List[str] = []
        for i, q in enumerate(questions):
            text, _ = self._call_llm(system_prompt, q, max_tokens=max_tokens)
            try:
                is_correct = bool(verifier_fn(q, text))
            except Exception:
                is_correct = False
            if is_correct:
                correct_questions.append(q)

            if (i + 1) % 20 == 0:
                acc = len(correct_questions) / (i + 1)
                print(f"  [{i+1}/{len(questions)}] correct: {len(correct_questions)} ({acc:.1%})")

        n_correct = len(correct_questions)
        print(
            f"  Calibration set: {n_correct}/{len(questions)} correct "
            f"({n_correct/len(questions):.1%} accuracy)"
        )

        if n_correct < self.n_neighbors + 1:
            raise ValueError(
                f"Only {n_correct} correctly-answered questions found. "
                f"Need at least {self.n_neighbors + 1}. "
                "Check your verifier_fn or try a different system prompt."
            )

        emb = self._embed(correct_questions)
        self._fit_knn(emb)
        return self

    def fit_qara(
        self,
        domains: List[Dict],
        epochs: int = 200,
        lr: float = 3e-4,
        verbose: bool = True,
    ) -> Dict:
        """
        Train the QARA cross-domain quality adapter.

        QARA (Quality-Aware Reasoning Adapter) learns to separate correct from
        incorrect LLM outputs in a 64-d space where domain identity is suppressed.
        After training, all KNN scoring happens in the adapted space, improving
        cross-domain AUROC from near-chance (~0.48) to ~0.63 on HP+NQ.

        Requires a prior call to fit() / fit_from_consistency() / fit_from_execution()
        so that calibration embeddings exist to re-fit the KNN in adapted space.

        Parameters
        ----------
        domains : list of dict
            Each dict must have:
              "name"      : str  — domain identifier (for logging)
              "questions" : list of str
              "labels"    : list of int/float  — 1=correct, 0=incorrect
        epochs : int
            Training epochs (default 200).
        lr : float
            Adam learning rate (default 3e-4).
        verbose : bool
            Print per-epoch loss every 50 steps.

        Returns
        -------
        dict with keys: n_examples, n_correct, n_domains, epochs

        Example
        -------
        guard.fit(hotpot_correct_questions)        # initial KNN calibration

        guard.fit_qara([
            {"name": "hotpot", "questions": hp_qs,  "labels": hp_labels},
            {"name": "nq",     "questions": nq_qs,  "labels": nq_labels},
        ])

        # Now cross-domain scoring benefits from QARA
        result = guard.query("What year was the Eiffel Tower built?")
        """
        all_embs, all_labels = [], []
        for d in domains:
            qs, ls = d["questions"], d["labels"]
            if len(qs) != len(ls):
                raise ValueError(
                    f"Domain '{d.get('name', '?')}': "
                    f"questions ({len(qs)}) and labels ({len(ls)}) must have equal length."
                )
            all_embs.append(self._embed(qs))
            all_labels.extend(ls)

        all_embs_arr = np.vstack(all_embs)
        all_labels_arr = np.array(all_labels, dtype=np.float32)
        n_correct = int(all_labels_arr.sum())
        n_total   = len(all_labels_arr)

        if verbose:
            print(
                f"fit_qara: {n_total} examples "
                f"({n_correct} correct, {n_total - n_correct} incorrect) "
                f"across {len(domains)} domain(s)"
            )

        self._qara = _qara_train(
            all_embs_arr, all_labels_arr,
            epochs=epochs, lr=lr, verbose=verbose,
        )

        # Re-fit KNN in the new adapted space using stored calibration embeddings.
        if self._cal_embs is not None:
            if verbose:
                print("fit_qara: re-fitting KNN in adapted space ...")
            self._fit_knn(self._cal_embs)
        elif verbose:
            print(
                "fit_qara: WARNING — no calibration embeddings stored. "
                "Call fit() first so the KNN is re-fitted in adapted space."
            )

        return {
            "n_examples": n_total,
            "n_correct":  n_correct,
            "n_domains":  len(domains),
            "epochs":     epochs,
        }

    def save_qara(self, path: str) -> None:
        """Persist the QARA adapter weights to a file (pickle)."""
        if self._qara is None:
            raise RuntimeError("No QARA adapter fitted. Call fit_qara() first.")
        with open(path, "wb") as f:
            pickle.dump(self._qara, f)

    def load_qara(self, path: str) -> "QPPGLLMGuard":
        """
        Load a previously saved QARA adapter and re-fit the KNN in adapted space.

        Parameters
        ----------
        path : str
            Path written by save_qara().
        """
        with open(path, "rb") as f:
            self._qara = pickle.load(f)
        if self._cal_embs is not None:
            self._fit_knn(self._cal_embs)
        return self

    # ── Query ──────────────────────────────────────────────────────────────────

    def query(
        self,
        question: str,
        system_prompt: str = None,
        max_tokens: int = 500,
    ) -> GuardResult:
        """
        Query the LLM with reliability scoring and automatic repair.

        Routing:
          Low risk   → direct LLM call                        (confidence "high")
          Medium risk → apply cluster repair tool if known     (confidence "medium")
          High risk  → call with explicit uncertainty flag     (confidence "low")

        Failure-type detection at medium/high risk:
          stop_reason == "max_tokens" → resource failure → retry with 2x tokens
          Otherwise                  → reasoning failure → apply synthesised tool

        Parameters
        ----------
        question : str
        system_prompt : str, optional
        max_tokens : int

        Returns
        -------
        GuardResult
        """
        if system_prompt is None:
            system_prompt = (
                "You are a helpful assistant. "
                "Think step by step and give your final answer on the last line."
            )

        risk = self._compute_risk_score(question)
        low_t  = self._risk_low_threshold  or 0.3
        high_t = self._risk_high_threshold or 0.7

        # ── Low risk ────────────────────────────────────────────────────────────
        if risk <= low_t:
            text, stop_reason = self._call_llm(system_prompt, question, max_tokens)
            retried = False
            if stop_reason == "max_tokens":
                # Resource failure — retry with more tokens
                text, _ = self._call_llm(system_prompt, question, max_tokens * 2)
                retried = True
            return GuardResult(
                answer=text, risk_score=risk, confidence="high",
                was_retried=retried, raw_response=text,
            )

        # ── Medium risk ─────────────────────────────────────────────────────────
        elif risk <= high_t:
            emb  = self._embed([question])
            tool = self._match_tool(emb)

            if tool:
                aug = (
                    f"{system_prompt}\n\n"
                    f"IMPORTANT: {tool.get('system_addition', '')}\n"
                    "Double-check your work before giving the final answer."
                )
                text, stop_reason = self._call_llm(aug, question, max_tokens)
                tool_id = tool.get("tool_name", "unknown")
                self._tool_usage[tool_id] = self._tool_usage.get(tool_id, 0) + 1
            else:
                text, stop_reason = self._call_llm(
                    system_prompt + "\n\nThis may be a challenging problem. Be extra careful.",
                    question, max_tokens,
                )
                tool_id = None

            retried = False
            if stop_reason == "max_tokens":
                # Resource failure — retry with more tokens, drop the tool
                text, _ = self._call_llm(system_prompt, question, max_tokens * 2)
                retried = True
                tool_id = None

            return GuardResult(
                answer=text, risk_score=risk, confidence="medium",
                tool_used=tool_id,
                cluster_id=tool.get("cluster_idx") if tool and not retried else None,
                was_retried=retried, raw_response=text,
            )

        # ── High risk ───────────────────────────────────────────────────────────
        else:
            text, stop_reason = self._call_llm(
                system_prompt
                + "\n\nIf you are not certain about any part of this, "
                "say so explicitly rather than guessing.",
                question, max_tokens,
            )
            retried = False
            if stop_reason == "max_tokens":
                text, _ = self._call_llm(system_prompt, question, max_tokens * 2)
                retried = True
            return GuardResult(
                answer=text, risk_score=risk, confidence="low",
                was_retried=retried, raw_response=text,
            )

    # ── Diagnose + Heal ────────────────────────────────────────────────────────

    def diagnose(
        self,
        failed_questions: List[str],
        model_answers: List[str],
        correct_answers: Optional[List[str]] = None,
        n_clusters: int = None,
    ) -> List[Dict]:
        """
        Error Autopsy: cluster failures into a labeled taxonomy.

        Read-only — does NOT modify guard state.
        Call learn_from_errors() to also synthesise repair tools.

        Returns
        -------
        List of cluster dicts, each containing:
          cluster_id    int
          size          number of failures in cluster
          label         LLM-generated description of the error pattern
          examples      up to 3 representative (question, model_answer) pairs
          suggested_fix one-sentence mitigation (only if correct_answers provided)
        """
        if len(failed_questions) < 5:
            return []

        emb = self._embed(failed_questions)
        labels, _, n_clusters = self._cluster_failures(emb, n_clusters)

        clusters = []
        for cid in range(n_clusters):
            mask      = labels == cid
            cq        = [failed_questions[i] for i, m in enumerate(mask) if m]
            ca        = [model_answers[i]    for i, m in enumerate(mask) if m]
            gt        = ([correct_answers[i] for i, m in enumerate(mask) if m]
                         if correct_answers else None)

            examples_str = ""
            for j in range(min(3, len(cq))):
                examples_str += f"\nQ: {cq[j][:300]}\nModel answer: {ca[j]}"
                if gt:
                    examples_str += f"\nCorrect answer: {gt[j]}"
                examples_str += "\n"

            label_text, _ = self._call_llm(
                "You are an error analysis expert.",
                f"These {mask.sum()} failures share a common error pattern:\n"
                f"{examples_str}\n"
                "In 1–2 sentences, describe the specific error pattern. "
                "Be concrete, not vague.",
                max_tokens=150,
            )

            cluster: Dict = {
                "cluster_id": cid,
                "size":       int(mask.sum()),
                "label":      label_text.strip(),
                "examples":   [
                    {"question": cq[j], "model_answer": ca[j]}
                    for j in range(min(3, len(cq)))
                ],
            }

            if gt:
                fix_text, _ = self._call_llm(
                    "You are an error analysis expert.",
                    f"Given this error pattern: {label_text}\n"
                    "Write a 1-sentence instruction to prevent this error.",
                    max_tokens=80,
                )
                cluster["suggested_fix"] = fix_text.strip()

            clusters.append(cluster)

        return clusters

    def learn_from_errors(
        self,
        failed_questions: List[str],
        model_answers: List[str],
        correct_answers: List[str],
        n_clusters: int = None,
    ) -> None:
        """
        Prompt Healer: cluster failures and synthesise targeted repair tools.

        After calling this, future queries that fall into a known failure cluster's
        embedding territory will automatically have the repair tool injected into
        their system prompt.

        Parameters
        ----------
        failed_questions : list of str
        model_answers    : list of str  — the LLM's wrong answers
        correct_answers  : list of str  — the ground truth answers
        n_clusters       : int, optional — auto-selected if None
        """
        if len(failed_questions) < 5:
            return

        emb = self._embed(failed_questions)
        labels, centers, n_clusters = self._cluster_failures(emb, n_clusters)
        self._cluster_centers = centers

        for cid in range(n_clusters):
            mask = labels == cid
            cq   = [failed_questions[i] for i, m in enumerate(mask) if m]
            cma  = [model_answers[i]    for i, m in enumerate(mask) if m]
            cgt  = [correct_answers[i]  for i, m in enumerate(mask) if m]

            if len(cq) < 2:
                continue

            examples = ""
            for j in range(min(3, len(cq))):
                examples += (
                    f"\nQ: {cq[j][:300]}\n"
                    f"Wrong: {cma[j]}\nCorrect: {cgt[j]}\n"
                )

            tool_text, _ = self._call_llm(
                "You are an error analysis expert.",
                f"These {mask.sum()} failures share a common error pattern:\n"
                f"{examples}\n"
                "Write a 2-sentence instruction that prevents this specific error. "
                "Be concrete and actionable.",
                max_tokens=200,
            )

            self._tools[str(cid)] = {
                "tool_name":       f"error_fix_{cid}",
                "system_addition": tool_text,
                "cluster_idx":     cid,
                "cluster_size":    int(mask.sum()),
            }

    # ── Internal: clustering ───────────────────────────────────────────────────

    def _cluster_failures(
        self,
        embeddings: np.ndarray,
        n_clusters: int = None,
    ):
        """
        Cluster failure embeddings. Auto-selects k via silhouette score if
        n_clusters is None. Returns (labels, centers, k).
        """
        from sklearn.cluster import KMeans
        from sklearn.metrics import silhouette_score

        if n_clusters is None:
            best_k, best_sil = 2, -1.0
            max_k = min(10, len(embeddings) // 3)
            for k in range(2, max(3, max_k)):
                km  = KMeans(n_clusters=k, n_init=10, random_state=42)
                lbl = km.fit_predict(embeddings)
                sil = silhouette_score(embeddings, lbl)
                if sil > best_sil:
                    best_k, best_sil = k, sil
            n_clusters = best_k

        km     = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
        labels = km.fit_predict(embeddings)
        return labels, km.cluster_centers_, n_clusters

    # ── Stats ──────────────────────────────────────────────────────────────────

    def get_stats(self) -> Dict:
        """Return usage and state statistics."""
        return {
            "total_calls":         self.total_calls,
            "total_input_tokens":  self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "n_tools":             len(self._tools),
            "tool_usage":          dict(self._tool_usage),
            "fitted":              self._fitted,
            "qara_fitted":         self._qara is not None,
            "risk_thresholds": {
                "low":  self._risk_low_threshold,
                "high": self._risk_high_threshold,
            },
            "cost_usd": round(
                self.total_input_tokens  * 0.80 / 1e6
                + self.total_output_tokens * 4.00 / 1e6,
                4,
            ),
        }
