"""
A2ATrustObject — Structured trust handoff between agents.

Defined and validated in exp105 (A2A Handoff Confidence Protocol).
This is the standard envelope an agent emits at the end of a chain so that
a downstream agent can condition its strategy on upstream confidence.

Key findings (exp105):
  - Inter-agent error correlation ρ = +0.108 (errors are weakly correlated)
  - P(B fails | A fails)  = 0.679 vs P(B fails | A correct) = 0.574, lift +0.105
  - Best downstream strategy: Agent B scores itself (AUROC 0.808) rather than
    trusting the chain product formula (AUROC 0.661)
  - When confidence_tier = LOW, Agent B should rewrite the query before
    proceeding (see QueryRewriter in query_rewriter.py)

Confidence tier thresholds (from exp92 conformal calibration):
  HIGH   risk_score < 0.50  → proceed normally
  MEDIUM 0.50 ≤ risk < 0.70 → proceed with monitoring
  LOW    risk_score ≥ 0.70  → trigger query rewriter or human review

Usage
-----
    # Agent A produces a trust object after its chain completes
    from llm_guard.agent_guard import AgentGuard

    guard = AgentGuard(api_key="sk-ant-...", use_judge=True)
    trust = guard.generate_trust_object(question, steps, final_answer)

    # Serialise for wire transport (JSON, message queue, etc.)
    payload = trust.to_dict()

    # Agent B deserialises and conditions its strategy
    trust = A2ATrustObject.from_dict(payload)
    if trust.should_rewrite:
        from llm_guard.query_rewriter import QueryRewriter
        rewriter = QueryRewriter(api_key="sk-ant-...")
        alt_queries = rewriter.rewrite(question, trust)
        # use alt_queries for Agent B's search strategy
"""

from __future__ import annotations

from dataclasses import dataclass, field, asdict
from typing import Optional, Dict, Any


@dataclass
class TemporalValidity:
    """
    Time-sensitivity metadata for a chain (exp106).

    Activated only when the question is classified as time-sensitive.
    Agent B should re-verify if tv_risk > 0.5.
    """
    is_time_sensitive: bool = False
    stale_score: float = 0.0       # 0=fresh, 1=stale observations
    recency_score: float = 0.0     # fraction of steps with recency signals
    tv_risk: float = 0.0           # combined temporal validity risk [0, 1]

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "TemporalValidity":
        return cls(
            is_time_sensitive=d.get("is_time_sensitive", False),
            stale_score=d.get("stale_score", 0.0),
            recency_score=d.get("recency_score", 0.0),
            tv_risk=d.get("tv_risk", 0.0),
        )


@dataclass
class A2ATrustObject:
    """
    Structured confidence envelope emitted by an agent at chain completion.

    Carry this alongside the answer when passing output to a downstream
    agent, orchestrator, or human review queue. The downstream consumer
    uses confidence_tier and failure_mode to decide whether to proceed,
    re-verify, or rewrite the query.

    Fields
    ------
    answer : str
        The agent's final answer.
    risk_score : float
        Composite risk score [0, 1]. Higher = more likely wrong.
        From J5 ensemble (SC_OLD behavioral + Sonnet judge) when
        use_judge=True, else SC_OLD behavioral only.
    confidence_tier : str
        "HIGH"   risk < 0.50  — proceed normally
        "MEDIUM" 0.50 ≤ risk < 0.70  — proceed with monitoring
        "LOW"    risk ≥ 0.70  — rewrite query or escalate
    failure_mode : str or None
        Detected failure pattern from LabelFreeScorer:
          "retrieval_fail"        observation contains no-results signal
          "repeated_query"        same search issued twice
          "long_chain"            4+ steps (budget exhaustion)
          "empty_answer"          final answer is absent or trivially short
          "low_retrieval_quality" mean cosine(q, obs) < 0.30
          "no_evidence"           no token overlap between queries and obs
          None                    no failure mode detected
    step_count : int
        Number of Search steps in the chain.
    judge_label : str or None
        Sonnet judge verdict: "GOOD" | "BORDERLINE" | "POOR" | None
        None when use_judge=False (saves ~$0.007/chain).
    downstream_hint : str
        Plain-English recommendation for the downstream agent.
          "proceed"                    HIGH confidence, no action needed
          "proceed_with_caution"       MEDIUM, validate key claims
          "rewrite_query"              LOW, try alternative formulation
          "rewrite_and_verify"         LOW + temporal validity issue
          "escalate_to_human"          POOR judge label, high risk
    should_rewrite : bool
        Shortcut: True when confidence_tier == "LOW".
        Set to True when QueryRewriter should be invoked before Agent B runs.
    behavioral_components : dict
        Raw SC_OLD component scores (sc1..sc12) for debugging / logging.
    temporal_validity : TemporalValidity or None
        Time-sensitivity metadata. None when TV scoring is disabled.
    """

    answer: str
    risk_score: float
    confidence_tier: str                # HIGH / MEDIUM / LOW
    failure_mode: Optional[str]
    step_count: int
    judge_label: Optional[str]          # GOOD / BORDERLINE / POOR / None
    downstream_hint: str
    should_rewrite: bool
    behavioral_components: Dict[str, float] = field(default_factory=dict)
    temporal_validity: Optional[TemporalValidity] = None

    # ── Derived helpers ───────────────────────────────────────────────────────

    @property
    def is_high_risk(self) -> bool:
        """True when the chain is likely wrong (risk >= 0.70)."""
        return self.risk_score >= 0.70

    @property
    def precision_at_alert(self) -> float:
        """
        Expected precision when this object triggers an alert.
        Based on exp92 conformal calibration:
          α=0.10 (FPR ≤ 10%) → Precision = 0.908
          α=0.05 (FPR ≤  5%) → Precision = 0.903
        Returns 0.0 when confidence_tier != "LOW" (no alert triggered).
        """
        return 0.908 if self.confidence_tier == "LOW" else 0.0

    # ── Serialisation ─────────────────────────────────────────────────────────

    def to_dict(self) -> Dict[str, Any]:
        """Serialise to a plain dict suitable for JSON transport."""
        d = {
            "answer":                 self.answer,
            "risk_score":             round(self.risk_score, 4),
            "confidence_tier":        self.confidence_tier,
            "failure_mode":           self.failure_mode,
            "step_count":             self.step_count,
            "judge_label":            self.judge_label,
            "downstream_hint":        self.downstream_hint,
            "should_rewrite":         self.should_rewrite,
            "behavioral_components":  {k: round(v, 4) for k, v in self.behavioral_components.items()},
            "temporal_validity":      self.temporal_validity.to_dict() if self.temporal_validity else None,
        }
        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "A2ATrustObject":
        """Deserialise from a plain dict (wire format)."""
        tv = d.get("temporal_validity")
        return cls(
            answer=d.get("answer", ""),
            risk_score=float(d.get("risk_score", 0.5)),
            confidence_tier=d.get("confidence_tier", "MEDIUM"),
            failure_mode=d.get("failure_mode"),
            step_count=int(d.get("step_count", 0)),
            judge_label=d.get("judge_label"),
            downstream_hint=d.get("downstream_hint", "proceed_with_caution"),
            should_rewrite=bool(d.get("should_rewrite", False)),
            behavioral_components=d.get("behavioral_components", {}),
            temporal_validity=TemporalValidity.from_dict(tv) if tv else None,
        )

    def __repr__(self) -> str:
        return (
            f"A2ATrustObject("
            f"tier={self.confidence_tier}, "
            f"risk={self.risk_score:.3f}, "
            f"judge={self.judge_label}, "
            f"hint={self.downstream_hint!r}"
            f")"
        )
