"""
AgentGuard — Real-time reliability monitoring for multi-step LLM agents.
=========================================================================

Validated performance (HotpotQA within-domain / NQ cross-domain):
  Chain-level AUROC (behavioral SC_OLD only):  ~0.81 within / ~0.67 cross
  Chain-level AUROC (SC_OLD + Sonnet judge):   ~0.78 within / ~0.74 cross
  Chain-level AUROC (local verifier, n=200):   ~0.80 within (exp111, $0 inference)
  Conformal alerting (J5 at FPR ≤ 10%):        Precision=0.908, Recall=0.595
  Mid-chain judge at step 2:                    AUROC=0.683 (Δ+0.156 vs SC_OLD)

  Source: exp88 (behavioral CV), exp89 (Sonnet judge), exp92 (conformal),
          exp107 (mid-chain), exp105 (A2A trust object), exp111 (local verifier)

Note on earlier AUROC 1.000 claim
----------------------------------
  Earlier versions of this module reported step-level AUROC 1.000 measured
  on the same data used for KNN calibration (in-sample evaluation).
  That figure does NOT generalise to unseen agent runs. Realistic estimates
  on held-out data are ~0.81 within-domain and ~0.74 cross-domain.
  See exp88 5-fold CV and exp89 for reliable figures.

Three operating modes
---------------------
  1. Behavioral only (use_judge=False, default)
     Zero API cost. AUROC ~0.81 within-domain. Good for high-throughput.

  2. Behavioral + Sonnet judge (use_judge=True)
     ~$0.007/chain. AUROC ~0.74 cross-domain. Best for cross-domain pipelines
     or when precision of the alert tier is critical.

  3. Behavioral + local verifier (use_local_verifier=True, requires 200 labels)
     Zero inference cost after training. AUROC ~0.80 within-domain (exp111).
     Beats Sonnet judge within-domain; no cross-domain validation yet.
     Requires: guard.fit_verifier(labeled_runs) before scoring.

Scoring signals used
--------------------
  SC1   Completion flag (0=Finish, 1=fallback)
  SC2   Step count (single strongest signal, AUROC ~0.88)
  SC3   Thought variance (embedding-based)
  SC5   Uncertainty word density
  SC6   Answer-thought token gap
  SC8   Backtracking rate (repeated queries)
  SC9   Observation utilisation
  SC10  Thought-observation coherence
  SC11  Answer-question mismatch
  SC12  Risk-monotone slope (coherence decay across steps)
  Judge Sonnet CoT assessment: GOOD / BORDERLINE / POOR (optional)

J5 ensemble (exp89): SC_OLD × 1 + Sonnet judge × 3, normalised to [0, 1].
  Judge risk mapping: GOOD→0.25, BORDERLINE→0.55, POOR→0.85

Confidence tiers (exp92 conformal calibration):
  HIGH   risk_score < 0.50 → proceed normally
  MEDIUM 0.50 ≤ risk < 0.70 → proceed with monitoring
  LOW    risk_score ≥ 0.70 → alert / escalate / rewrite query

A2A trust object
----------------
  generate_trust_object() emits an A2ATrustObject with the full trust envelope.
  Downstream agents import and condition their search strategy on it.
  When confidence_tier == "LOW", pass to QueryRewriter for query diversification.

Quick start
-----------
    from llm_guard import AgentGuard

    guard = AgentGuard()                              # behavioral only, $0
    # guard = AgentGuard(api_key="sk-ant-...", use_judge=True)  # + Sonnet judge

    # 1. Score a completed chain
    result = guard.score_chain(
        question="Who was older, Einstein or Bohr?",
        steps=[
            {"thought": "Search for Einstein birth year",
             "action_type": "Search", "action_arg": "Einstein birth year",
             "observation": "Albert Einstein was born on March 14, 1879..."},
            {"thought": "Einstein (1879) is older than Bohr (1885)",
             "action_type": "Finish", "action_arg": "Einstein", "observation": ""},
        ],
        final_answer="Einstein",
    )
    print(result.confidence_tier)   # HIGH / MEDIUM / LOW
    print(result.risk_score)        # 0.0-1.0, higher = more likely wrong
    print(result.needs_alert)       # True when risk >= 0.70

    # 2. Get A2A trust object (for handoff to Agent B)
    trust = guard.generate_trust_object(question, steps, final_answer)
    payload = trust.to_dict()       # JSON-serialisable for wire transport

    # 3. Mid-chain monitoring (call at each step during agent loop)
    step_result = guard.monitor_step(
        question="Who was older, Einstein or Bohr?",
        steps_so_far=[step1],
        current_action="Search[Niels Bohr birth year]",
    )
    if step_result.risk == "high":
        pass  # intervene: retry, escalate, or abort early

    # 4. Calibrate on your own correct chains (improves GMM + obs-pool signals)
    guard.fit_from_agent_runs(my_correct_runs)

    # 5. Pre-screen before running the agent
    pre = guard.score_chain_start("Who was president when X happened?")
    print(pre["predicted_outcome"])  # "likely_success" | "uncertain" | "likely_failure"

    # 6. Local verifier (replaces Sonnet judge after 200 labeled runs, $0 inference)
    guard = AgentGuard(use_local_verifier=True)
    guard.fit_verifier(labeled_runs)  # [{"question":..,"steps":..,"final_answer":..,"correct":bool}]
    result = guard.score_chain(question, steps, final_answer)
    # Uses J_LOCAL = SC_OLD×1 + LocalVerifier×3 (AUROC ~0.80 within-domain, exp111)

    # 7. Pseudo-label conformal calibration (no human labels needed)
    guard = AgentGuard(use_judge=True)
    guard.calibrate_from_agreement(unlabeled_runs_with_agent_b_answers)
    # Uses 3-agent agreement as pseudo-labels for conformal alerting (exp110)
"""

from __future__ import annotations

import json
import os
import re
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Dict, List, Optional

import numpy as np

# Ensure qppg_service is importable whether running from repo root or installed
_repo_root = Path(__file__).resolve().parent.parent
if str(_repo_root) not in sys.path:
    sys.path.insert(0, str(_repo_root))

from qppg_service.label_free_scorer import LabelFreeScorer
from llm_guard.trust_object import A2ATrustObject, TemporalValidity
from llm_guard.local_verifier import LocalVerifier


# ── Judge prompt (exp89 4-step CoT, Sonnet) ───────────────────────────────────

_JUDGE_SYSTEM = """\
You are a strict evaluator of AI reasoning chains. Assess whether the agent \
produced a correct, well-supported answer.

Evaluate in JSON only:
{
  "diagnosis": "one sentence: what the agent did and why it may be wrong",
  "label": "GOOD or BORDERLINE or POOR",
  "confidence": 1-5
}

GOOD      = searches are productive, reasoning is sound, answer is well-supported
BORDERLINE = some gaps or uncertainty but answer may be correct
POOR      = clear reasoning failures, answer is unsupported or likely wrong

Output ONLY the JSON object."""

_JUDGE_USER_TMPL = """\
QUESTION: {question}

AGENT CHAIN:
{chain_text}

FINAL ANSWER: {final_answer}"""

# Judge label → risk score mapping (calibrated on exp89 distribution)
_JUDGE_RISK: Dict[str, float] = {"GOOD": 0.25, "BORDERLINE": 0.55, "POOR": 0.85}

# J5 ensemble weights: SC_OLD × 1, Sonnet judge × 3
_J5_SC_WEIGHT    = 1.0
_J5_JUDGE_WEIGHT = 3.0


# ── Result dataclasses ────────────────────────────────────────────────────────

@dataclass
class AgentStepResult:
    """Result of AgentGuard.monitor_step() — per-step mid-chain scoring."""
    risk_score: float        # 0-1; higher = more likely in a failing chain
    risk: str                # "low" | "medium" | "high"
    confidence: str          # "high" | "medium" | "low"  (inverse of risk)
    predicted_outcome: str   # "likely_success" | "uncertain" | "likely_failure"
    step_index: int          # 0-based index of the current step
    failure_mode: Optional[str] = None   # detected failure mode or None


@dataclass
class ChainTrustResult:
    """
    Result of AgentGuard.score_chain() — full-chain assessment.

    When needs_alert=True, treat the chain as a likely failure.
    Expected precision at that threshold: ~0.908 (FPR ≤ 10%, exp92).
    """
    risk_score: float              # J5 composite score [0, 1]
    confidence_tier: str           # HIGH / MEDIUM / LOW
    needs_alert: bool              # True if risk >= alert_threshold (default 0.70)
    failure_mode: Optional[str]    # detected failure pattern or None
    judge_label: Optional[str]     # GOOD / BORDERLINE / POOR / None
    step_count: int                # number of Search steps
    behavioral_score: float        # SC_OLD score before judge blending
    behavioral_components: Dict[str, float] = field(default_factory=dict)
    latency_ms: float = 0.0


# ── AgentGuard ────────────────────────────────────────────────────────────────

class AgentGuard:
    """
    Real-time reliability monitor for multi-step LLM agents.

    Uses validated behavioral signals (SC_OLD, exp88) as the zero-cost baseline,
    optionally enhanced with a Sonnet judge (exp89) for cross-domain accuracy.

    Parameters
    ----------
    api_key : str, optional
        Anthropic API key. Required only when use_judge=True.
        Falls back to ANTHROPIC_API_KEY env var.
    use_judge : bool
        If True, call Sonnet after each completed chain (~$0.007/chain).
        Improves cross-domain AUROC from ~0.67 to ~0.74.
        Default False (behavioral only, $0 cost).
    use_local_verifier : bool
        If True, use a trained LocalVerifier instead of the Sonnet judge.
        $0 inference cost. AUROC ~0.80 within-domain (exp111).
        Requires calling fit_verifier() before scoring.
        Cannot be combined with use_judge=True (local verifier takes precedence).
    judge_model : str
        Model for the judge. Default: claude-sonnet-4-6 (exp89 validated).
    alert_threshold : float
        Risk score above which needs_alert=True. Default 0.70.
        At this threshold: Precision=0.908, Recall=0.595 (exp92 conformal).
    review_threshold : float
        Threshold passed to LabelFreeScorer for needs_review flag. Default 0.65.
    on_alert : callable, optional
        Called with ChainTrustResult whenever needs_alert=True.
        Use for webhook dispatch, Slack notifications, or routing to human review.
        Example: on_alert=lambda r: requests.post(webhook_url, json=r.__dict__)
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        use_judge: bool = False,
        use_local_verifier: bool = False,
        judge_model: str = "claude-sonnet-4-6",
        alert_threshold: float = 0.70,
        review_threshold: float = 0.65,
        on_alert: Optional[Callable[["ChainTrustResult"], None]] = None,
    ):
        self._api_key             = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self._use_judge           = use_judge
        self._use_local_verifier  = use_local_verifier
        self._judge_model         = judge_model
        self._alert_threshold     = alert_threshold
        self._on_alert            = on_alert
        self._scorer              = LabelFreeScorer(review_threshold=review_threshold)
        self._local_verifier: Optional[LocalVerifier] = None
        self._client              = None
        self._fitted              = False

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

    # ── Calibration ────────────────────────────────────────────────────────────

    def fit_from_agent_runs(
        self,
        runs: List[Dict],
        correct_only: bool = True,
    ) -> "AgentGuard":
        """
        Calibrate the internal LabelFreeScorer on completed agent runs.

        Feeds chains into the GMM + obs-pool calibration (exp60).
        Requires ≥ 20 chains for density signals to activate.
        Below that threshold, behavioral SC_OLD already provides AUROC ~0.81.

        Parameters
        ----------
        runs : list of dicts
            Each dict must have "question", "steps", "final_answer" keys.
            Optional: "chain_correct": bool for filtering.
        correct_only : bool
            If True, use only runs where chain_correct=True (or key is absent).
        """
        chains = []
        for run in runs:
            if correct_only and not run.get("chain_correct", True):
                continue
            chains.append({
                "question":     run["question"],
                "steps":        run.get("steps", []),
                "final_answer": run.get("final_answer", ""),
            })
        if len(chains) >= 5:
            self._scorer.calibrate(chains)
            self._fitted = True
        return self

    def fit_verifier(
        self,
        runs: List[Dict],
    ) -> "AgentGuard":
        """
        Train the local verifier on labeled agent runs (exp111).

        Trains a LogisticRegression on 12 Jaccard behavioral features.
        After training, use_local_verifier=True routes score_chain() through
        J_LOCAL = SC_OLD×1 + LocalVerifier×3 (AUROC ~0.80 within-domain).

        Requires ≥ 50 labeled chains (recommended: ≥ 200 for stable CV scores).
        Automatically activates use_local_verifier after fitting.

        Parameters
        ----------
        runs : list of dicts
            Each dict must have "question", "steps", "final_answer", "correct" keys.
            "correct" (bool): True if the chain produced the right answer.
        """
        if len(runs) < 50:
            import warnings
            warnings.warn(
                f"fit_verifier called with only {len(runs)} runs. "
                "Recommend ≥ 200 for stable AUROC estimates (exp111).",
                UserWarning,
                stacklevel=2,
            )
        verifier = LocalVerifier()
        verifier.fit(runs)
        self._local_verifier = verifier
        self._use_local_verifier = True
        return self

    def calibrate_from_agreement(
        self,
        runs: List[Dict],
        agreement_threshold: float = 0.5,
        alpha: float = 0.10,
    ) -> "AgentGuard":
        """
        Pseudo-label conformal calibration using 3-agent agreement (exp110).

        Generates pseudo-labels from multi-agent agreement (no human annotation),
        then computes a conformal alerting threshold at the target FPR (alpha).

        Validated precision gap vs true labels: ~6.1% (marginal).
        Recall is substantially lower than true-label conformal (0.207 vs 0.552).
        Use with caveat — best for cold-start deployment before labels accumulate.

        Parameters
        ----------
        runs : list of dicts
            Each dict must have "question", "steps", "final_answer" keys.
            Also needs "agent_b_answer" and "agent_c_answer" (str) for agreement.
        agreement_threshold : float
            Mean pairwise token-F1 >= this → pseudo-label as correct. Default 0.5.
        alpha : float
            Target FPR for conformal coverage (1 - alpha). Default 0.10 (10% FPR).
        """
        import re as _re

        def _toks_local(text: str):
            w = _re.findall(r"[a-zA-Z]+", text.lower())
            return {x for x in w if len(x) > 1}

        def _f1_local(a: str, b: str) -> float:
            p, r = _toks_local(a), _toks_local(b)
            if not p or not r:
                return 0.0
            c = p & r
            pr = len(c) / len(p)
            rc = len(c) / len(r)
            return 2 * pr * rc / (pr + rc) if pr + rc else 0.0

        pseudo_correct_risks = []
        for run in runs:
            fa  = run.get("final_answer", "")
            ans_b = run.get("agent_b_answer", "")
            ans_c = run.get("agent_c_answer", "")
            if not (fa and ans_b and ans_c):
                continue
            f1_ab = _f1_local(fa, ans_b)
            f1_ac = _f1_local(fa, ans_c)
            f1_bc = _f1_local(ans_b, ans_c)
            mean_f1 = (f1_ab + f1_ac + f1_bc) / 3.0
            if mean_f1 >= agreement_threshold:
                lf_result = self._scorer.score(
                    run["question"], run.get("steps", []), fa
                )
                pseudo_correct_risks.append(float(lf_result.risk_score))

        if len(pseudo_correct_risks) < 10:
            return self  # not enough data; keep existing threshold

        risks_arr = np.array(pseudo_correct_risks)
        n = len(risks_arr)
        adjusted_q = min(1.0, (1 + 1.0 / n) * (1 - alpha))
        new_threshold = float(np.quantile(risks_arr, adjusted_q))
        self._alert_threshold = new_threshold
        return self

    # ── Chain scoring ──────────────────────────────────────────────────────────

    def score_chain(
        self,
        question: str,
        steps: List[Dict],
        final_answer: str,
        finished: bool = True,
    ) -> ChainTrustResult:
        """
        Score a completed reasoning chain.

        Runs SC_OLD behavioral signals (always) and optionally the Sonnet judge
        (when use_judge=True). Combines as J5 = SC_OLD×1 + judge×3.

        Parameters
        ----------
        question : str
        steps : list of dicts
            Each step: {"thought": str, "action_type": str,
                        "action_arg": str, "observation": str}
        final_answer : str
        finished : bool
            True if the chain reached a Finish action.

        Returns
        -------
        ChainTrustResult
        """
        t0 = time.time()

        # 1. Behavioral SC_OLD (always runs, $0)
        lf_result    = self._scorer.score(question, steps, final_answer, finished)
        sc_old_risk  = float(lf_result.risk_score)
        failure_mode = lf_result.failure_mode
        components   = dict(lf_result.components) if hasattr(lf_result, "components") else {}
        step_count   = sum(1 for s in steps if s.get("action_type") == "Search")

        # 2. Optional Sonnet judge (~$0.007/chain) or local verifier ($0)
        judge_label = None
        judge_risk  = None

        if self._use_local_verifier and self._local_verifier is not None and self._local_verifier.is_fitted:
            # J_LOCAL: SC_OLD×1 + LocalVerifier×3 (exp111, AUROC ~0.80 within-domain)
            judge_risk  = self._local_verifier.predict_risk(question, steps, final_answer)
            judge_label = "LOCAL"
        elif self._use_judge and self._api_key:
            # J5: SC_OLD×1 + Sonnet judge×3 (exp89, AUROC ~0.78 within / ~0.74 cross)
            judge_label, judge_risk = self._run_judge(question, steps, final_answer)

        # 3. Ensemble: SC_OLD×1 + secondary signal×3 (normalised)
        if judge_risk is not None:
            risk_score = (
                _J5_SC_WEIGHT * sc_old_risk + _J5_JUDGE_WEIGHT * judge_risk
            ) / (_J5_SC_WEIGHT + _J5_JUDGE_WEIGHT)
        else:
            risk_score = sc_old_risk

        confidence_tier = _risk_to_tier(risk_score)
        needs_alert     = risk_score >= self._alert_threshold
        latency         = (time.time() - t0) * 1000

        result = ChainTrustResult(
            risk_score=round(float(risk_score), 4),
            confidence_tier=confidence_tier,
            needs_alert=needs_alert,
            failure_mode=failure_mode,
            judge_label=judge_label,
            step_count=step_count,
            behavioral_score=round(sc_old_risk, 4),
            behavioral_components=components,
            latency_ms=round(latency, 1),
        )

        if needs_alert and self._on_alert is not None:
            try:
                self._on_alert(result)
            except Exception:
                pass  # never let the alert callback break the scoring path

        return result

    # ── A2A trust object ───────────────────────────────────────────────────────

    def generate_trust_object(
        self,
        question: str,
        steps: List[Dict],
        final_answer: str,
        finished: bool = True,
    ) -> A2ATrustObject:
        """
        Produce an A2ATrustObject for agent-to-agent handoff (exp105).

        The trust object is the standardised confidence envelope that downstream
        agents consume to decide whether to proceed, verify, or rewrite the query.

        Serialise with trust.to_dict() for JSON/queue transport.
        Deserialise with A2ATrustObject.from_dict(payload).

        Returns
        -------
        A2ATrustObject
        """
        result = self.score_chain(question, steps, final_answer, finished)

        downstream_hint = _build_downstream_hint(
            result.confidence_tier, result.judge_label, result.failure_mode
        )

        return A2ATrustObject(
            answer=final_answer,
            risk_score=result.risk_score,
            confidence_tier=result.confidence_tier,
            failure_mode=result.failure_mode,
            step_count=result.step_count,
            judge_label=result.judge_label,
            downstream_hint=downstream_hint,
            should_rewrite=result.confidence_tier == "LOW",
            behavioral_components=result.behavioral_components,
            temporal_validity=None,
        )

    # ── Mid-chain monitoring ───────────────────────────────────────────────────

    def monitor_step(
        self,
        question: str,
        steps_so_far: List[Dict],
        current_action: str,
    ) -> AgentStepResult:
        """
        Score a single agent step in real time (mid-chain intervention, exp107).

        Call BEFORE executing each action in your agent loop. High risk at step 2
        is a strong early-failure signal (AUROC=0.683, Δ+0.156 vs SC_OLD at step 2).

        Parameters
        ----------
        question : str
        steps_so_far : list of dicts
            Steps already completed (empty list = first step).
        current_action : str
            The action the agent is about to take.

        Returns
        -------
        AgentStepResult
        """
        t = len(steps_so_far)

        if t == 0:
            return AgentStepResult(
                risk_score=0.5, risk="medium", confidence="medium",
                predicted_outcome="uncertain", step_index=0,
            )

        lf_result  = self._scorer.score_prefix(question, steps_so_far, t)
        risk_score = float(lf_result.risk_score)
        risk, confidence, predicted = _risk_to_labels(risk_score)

        return AgentStepResult(
            risk_score=round(risk_score, 4),
            risk=risk,
            confidence=confidence,
            predicted_outcome=predicted,
            step_index=t,
            failure_mode=lf_result.failure_mode,
        )

    # ── Pre-screening ──────────────────────────────────────────────────────────

    def score_chain_start(self, question: str) -> Dict:
        """
        Pre-screen a question BEFORE running the agent at all.

        Most signal comes from the chain itself — this is primarily useful for
        routing obviously difficult questions to a stronger model upfront.

        Returns dict with: risk_score, tier, confidence, predicted_outcome,
        recommended_action.
        """
        lf_result  = self._scorer.score(question, [], "", False)
        risk_score = float(lf_result.risk_score)
        tier       = _risk_to_tier(risk_score)
        _, confidence, predicted = _risk_to_labels(risk_score)

        recommendation = {
            "LOW":    "use_stronger_model_or_add_verification",
            "MEDIUM": "proceed_with_monitoring",
            "HIGH":   "proceed",
        }[tier]

        return {
            "risk_score":         round(risk_score, 4),
            "tier":               tier,
            "confidence":         confidence,
            "predicted_outcome":  predicted,
            "recommended_action": recommendation,
        }

    # ── Diagnostics ────────────────────────────────────────────────────────────

    def diagnostics(self) -> Dict:
        """Return a summary of the guard's current configuration and validated performance."""
        mode = "behavioral"
        if self._use_local_verifier and self._local_verifier is not None:
            mode = f"j_local (verifier trained on {self._local_verifier.n_train} chains)"
        elif self._use_judge:
            mode = f"j5_sonnet ({self._judge_model})"

        return {
            "mode":                mode,
            "use_judge":           self._use_judge,
            "use_local_verifier":  self._use_local_verifier,
            "judge_model":         self._judge_model if self._use_judge else None,
            "alert_threshold":     self._alert_threshold,
            "calibrated":          self._fitted,
            "verifier_fitted":     self._local_verifier.is_fitted if self._local_verifier else False,
            "expected_auroc": {
                "within_domain_behavioral":      "~0.81 (exp88 5-fold CV)",
                "within_domain_j5_sonnet":       "~0.78 (exp89 single run)",
                "within_domain_j_local_200":     "~0.80 (exp111 5-fold CV)",
                "cross_domain_behavioral":       "~0.67 (exp91 NQ)",
                "cross_domain_j5_sonnet":        "~0.74 (exp91 NQ)",
            },
            "conformal_alerting_at_threshold_0.70": {
                "fpr_guarantee": "≤ 10%",
                "precision":     0.908,
                "recall":        0.595,
                "source":        "exp92",
            },
        }

    # ── Backward-compatible helper (kept for existing integrations) ────────────

    @staticmethod
    def format_step_context(
        question: str,
        steps_so_far: List[Dict],
        current_action: str,
    ) -> str:
        """Serialise a step into a string (legacy format, kept for compatibility)."""
        parts = [f"TASK: {question}"]
        for i, s in enumerate(steps_so_far, 1):
            thought = s.get("thought", "")[:200]
            if "action_type" in s and "action_arg" in s:
                action = f"{s['action_type']}[{s['action_arg']}]"[:200]
            else:
                action = str(s.get("action", ""))[:200]
            parts.append(f"STEP {i}: Thought: {thought} | Action: {action}")
        parts.append(f"CURRENT: {current_action[:300]}")
        return "\n".join(parts)

    # ── Private helpers ────────────────────────────────────────────────────────

    def _run_judge(self, question: str, steps: List[Dict], final_answer: str):
        """Call Sonnet judge (exp89). Returns (label, risk_float) or (None, None)."""
        chain_text = _format_chain_for_judge(steps)
        user_msg   = _JUDGE_USER_TMPL.format(
            question=question,
            chain_text=chain_text,
            final_answer=final_answer,
        )
        try:
            resp = self._get_client().messages.create(
                model=self._judge_model,
                max_tokens=150,
                temperature=0.0,
                system=_JUDGE_SYSTEM,
                messages=[{"role": "user", "content": user_msg}],
            )
            raw = resp.content[0].text.strip()
            m   = re.search(r'\{.*\}', raw, re.DOTALL)
            if not m:
                return None, None
            data  = json.loads(m.group())
            label = str(data.get("label", "BORDERLINE")).upper().strip()
            if label not in _JUDGE_RISK:
                label = "BORDERLINE"
            return label, _JUDGE_RISK[label]
        except Exception:
            return None, None


# ── Module-level helpers ──────────────────────────────────────────────────────

def _risk_to_tier(risk_score: float) -> str:
    if risk_score < 0.50:
        return "HIGH"
    elif risk_score < 0.70:
        return "MEDIUM"
    return "LOW"


def _risk_to_labels(risk_score: float):
    if risk_score < 0.50:
        return "low", "high", "likely_success"
    elif risk_score < 0.70:
        return "medium", "medium", "uncertain"
    return "high", "low", "likely_failure"


def _build_downstream_hint(
    tier: str,
    judge_label: Optional[str],
    failure_mode: Optional[str],
) -> str:
    if tier == "HIGH":
        return "proceed"
    if tier == "MEDIUM":
        return "proceed_with_caution"
    if judge_label == "POOR":
        return "escalate_to_human"
    if failure_mode in ("retrieval_fail", "no_evidence", "repeated_query"):
        return "rewrite_and_verify"
    return "rewrite_query"


def _format_chain_for_judge(steps: List[Dict]) -> str:
    parts = []
    for i, s in enumerate(steps, 1):
        thought = s.get("thought", "").strip()[:300]
        a_type  = s.get("action_type", "")
        a_arg   = s.get("action_arg", "").strip()[:200]
        obs     = s.get("observation", "").strip()[:300]
        parts.append(
            f"Step {i}: {thought}\n"
            f"  Action: {a_type}[{a_arg}]\n"
            f"  Result: {obs}"
        )
    return "\n".join(parts) if parts else "(no steps)"
