"""
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
  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)

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.

Two 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.

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"
"""

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 Dict, List, Optional

# 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


# ── 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).
    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.
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        use_judge: bool = False,
        judge_model: str = "claude-sonnet-4-6",
        alert_threshold: float = 0.70,
        review_threshold: float = 0.65,
    ):
        self._api_key         = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self._use_judge       = use_judge
        self._judge_model     = judge_model
        self._alert_threshold = alert_threshold
        self._scorer          = LabelFreeScorer(review_threshold=review_threshold)
        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

    # ── 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)
        judge_label = None
        judge_risk  = None
        if self._use_judge and self._api_key:
            judge_label, judge_risk = self._run_judge(question, steps, final_answer)

        # 3. J5 ensemble: SC_OLD×1 + judge×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

        return 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),
        )

    # ── 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."""
        return {
            "use_judge":        self._use_judge,
            "judge_model":      self._judge_model if self._use_judge else None,
            "alert_threshold":  self._alert_threshold,
            "calibrated":       self._fitted,
            "expected_auroc": {
                "within_domain_behavioral": "~0.81 (exp88 5-fold CV)",
                "within_domain_j5":         "~0.78 (exp89 single run)",
                "cross_domain_behavioral":  "~0.67 (exp91 NQ)",
                "cross_domain_j5":          "~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)"
