"""
AdversarialChainDetector — Detect Adversarially-Crafted Failure Chains
=======================================================================

Detects "confident wrong" chains — chains where the agent gives a wrong answer
with high apparent confidence, defeating standard risk scoring.

WARNING — THIS DETECTOR DOES NOT GENERALIZE TO NATURAL CHAINS:
  AUROC on natural HP wrong-answer chains : 0.397  (negatively correlated, exp_adversarial_holdout)
  AUROC on FARL-engineered chains (CV)    : 0.9960 CI [0.9906, 0.9996]  ← FARL domain only

  The 0.9960 figure is achieved on chains purposely engineered to contain adversarial patterns.
  On real user queries the detector scores wrong answers HIGHER than correct ones (AUROC 0.397).
  Do NOT use for production reliability monitoring. Experimental=True required to instantiate.

Top discriminating features:
  1. frac_search       — adversarial chains do fewer searches relative to steps
  2. avg_obs_len       — adversarial chains retrieve longer but less diverse evidence
  3. n_unique_queries  — adversarial chains repeat fewer search queries
  4. q_wh_position     — question structure (WH-word position norm)

Design
------
  Features (16): 9 structural chain features + 7 question complexity features
  Classifier   : Logistic Regression (L2, balanced class weights)
  No external calls — $0 inference cost.

Quick start
-----------
    from llm_guard import AdversarialChainDetector

    # Train on your labeled chains
    detector = AdversarialChainDetector()
    detector.fit(adversarial_steps_and_qs, correct_steps_and_qs)

    # Score a new chain
    score = detector.score(steps, question)   # 0.0 = normal, 1.0 = adversarial
    flag  = detector.predict(steps, question) # True = adversarial

    # Export / import for federated or persisted use
    params = detector.export_params()
    detector2 = AdversarialChainDetector.from_params(params)
Purpose:    16-feature behavioral classifier pre-trained on FARL-generated adversarial chains. CV AUROC 0.9960 on FARL chains; AUROC 0.397 on natural HP chains (negatively correlated). Experimental=True gated. NOT for production reliability monitoring.
Exports:    AdversarialChainDetector
Added:      v0.31.0
Status:     Experimental
"""
from __future__ import annotations

import json
import warnings
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np

from llm_guard.adversarial.novelty_detector import _complexity_features


# ── Feature names (public for inspection) ─────────────────────────────────────

ADVERSARIAL_FEATURE_NAMES: List[str] = [
    "n_steps", "n_search", "frac_search",
    "avg_thought_len", "avg_obs_len", "std_obs_len",
    "min_obs_len", "final_thought_len", "n_unique_queries",
    "q_token_count", "q_clause_count", "q_multihop_cue",
    "q_wh_position", "q_ne_count", "q_comma_count", "q_char_len",
]


# ── Dataclasses ────────────────────────────────────────────────────────────────

@dataclass
class AdversarialResult:
    """
    Result from AdversarialChainDetector.score_chain().

    adv_score   : probability of chain being adversarially crafted (0.0–1.0).
    is_adversarial: True when adv_score ≥ threshold.
    routing     : "safe" or "adversarial_suspected".
    top_features: list of (feature_name, value) for the most influential features.
    """
    adv_score:      float
    is_adversarial: bool
    routing:        str
    top_features:   List[Tuple[str, float]] = field(default_factory=list)


@dataclass
class AdversarialParams:
    """Serialisable model parameters (no raw training data)."""
    coef:      List[float]
    intercept: float
    scale_mean: List[float]
    scale_std:  List[float]
    threshold:  float = 0.50
    auroc:      Optional[float] = None

    def to_dict(self) -> Dict:
        return {
            "coef":        self.coef,
            "intercept":   self.intercept,
            "scale_mean":  self.scale_mean,
            "scale_std":   self.scale_std,
            "threshold":   self.threshold,
            "auroc":       self.auroc,
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict())

    @classmethod
    def from_dict(cls, d: Dict) -> "AdversarialParams":
        return cls(
            coef=[float(x) for x in d["coef"]],
            intercept=float(d["intercept"]),
            scale_mean=[float(x) for x in d["scale_mean"]],
            scale_std=[float(x) for x in d["scale_std"]],
            threshold=float(d.get("threshold", 0.50)),
            auroc=d.get("auroc"),
        )

    @classmethod
    def from_json(cls, s: str) -> "AdversarialParams":
        return cls.from_dict(json.loads(s))


# ── Feature extraction (module-level helper) ───────────────────────────────────

def extract_chain_features(steps: List[Dict], question: str) -> np.ndarray:
    """
    Extract 16 behavioral + structural features from a chain.

    Parameters
    ----------
    steps    : list of step dicts with keys {thought, action_type, action_arg, observation}
    question : the question string

    Returns
    -------
    np.ndarray of shape (16,)
    """
    if not steps:
        return np.zeros(16)

    n_steps      = len(steps)
    thoughts     = [s.get("thought", "") or "" for s in steps]
    obs          = [s.get("observation", "") or "" for s in steps]
    action_types = [str(s.get("action_type", "")).lower() for s in steps]
    queries      = [s.get("action_arg", "") or "" for s in steps
                    if str(s.get("action_type", "")).lower() == "search"]

    n_search     = sum(1 for a in action_types if a == "search")
    frac_search  = n_search / n_steps

    thought_lens = [len(t) for t in thoughts]
    obs_lens     = [len(o) for o in obs]

    avg_thought  = float(np.mean(thought_lens)) if thought_lens else 0.0
    avg_obs      = float(np.mean(obs_lens))     if obs_lens     else 0.0
    std_obs      = float(np.std(obs_lens))      if obs_lens     else 0.0
    min_obs      = float(min(obs_lens))         if obs_lens     else 0.0
    final_th_len = float(thought_lens[-1])      if thought_lens else 0.0
    n_unique_q   = float(len(set(q.strip().lower() for q in queries)))

    feat_struct = np.array([
        min(n_steps        / 5.0,  2.0),
        min(n_search       / 3.0,  2.0),
        frac_search,
        min(avg_thought    / 300., 2.0),
        min(avg_obs        / 300., 2.0),
        min(std_obs        / 200., 2.0),
        min(min_obs        / 200., 2.0),
        min(final_th_len   / 300., 2.0),
        min(n_unique_q     / 3.0,  2.0),
    ], dtype=float)

    feat_q = _complexity_features(question)

    return np.concatenate([feat_struct, feat_q])


# ── Main class ─────────────────────────────────────────────────────────────────

class AdversarialChainDetector:
    """
    Detect adversarially-crafted "confident wrong" chains via behavioral features.

    No external API calls — $0 inference cost.

    Parameters
    ----------
    threshold : float
        Probability threshold for is_adversarial.  Default 0.50.
    C : float
        Logistic Regression regularization strength.  Default 1.0.

    Validated performance:
      Natural HP chains (exp_adversarial_holdout): AUROC 0.397 — negatively correlated
      FARL-engineered chains (CV, exp_adversarial_detection): AUROC 0.9960 CI [0.9906, 0.9996]
      Use ONLY for detecting known FARL-style adversarial chains, not general wrong answers.
    """

    def __init__(self, threshold: float = 0.50, C: float = 1.0,
                 experimental: bool = False) -> None:
        if not experimental:
            raise RuntimeError(
                "AdversarialChainDetector requires experimental=True to instantiate.\n"
                "Natural-chain holdout AUROC=0.3974 (negatively correlated with natural errors).\n"
                "This detector is ONLY valid for FARL-style adversarial chains, not general "
                "wrong-answer detection. Set experimental=True to acknowledge this limitation.\n"
                "See docs/adversarial_methodology.md for the full evaluation."
            )
        self._threshold = threshold
        self._C         = C
        self._coef: Optional[np.ndarray] = None
        self._intercept: float = 0.0
        self._scale_mean: Optional[np.ndarray] = None
        self._scale_std:  Optional[np.ndarray] = None
        self._auroc: Optional[float] = None

    # ── Fit ──────────────────────────────────────────────────────────────────

    def fit(
        self,
        adversarial: List[Tuple[List[Dict], str]],
        correct:     List[Tuple[List[Dict], str]],
    ) -> "AdversarialChainDetector":
        """
        Train the detector.

        Parameters
        ----------
        adversarial : list of (steps, question) for adversarial/confident-wrong chains
        correct     : list of (steps, question) for correct chains

        Returns self.
        """
        from sklearn.linear_model import LogisticRegression
        from sklearn.preprocessing import StandardScaler

        X_list, y_list = [], []
        for steps, q in adversarial:
            X_list.append(extract_chain_features(steps, q))
            y_list.append(1)
        for steps, q in correct:
            X_list.append(extract_chain_features(steps, q))
            y_list.append(0)

        X = np.array(X_list, dtype=float)
        y = np.array(y_list, dtype=int)

        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)

        clf = LogisticRegression(
            C=self._C, class_weight="balanced",
            max_iter=1000, random_state=42,
        )
        clf.fit(X_scaled, y)

        self._coef       = clf.coef_[0].copy()
        self._intercept  = float(clf.intercept_[0])
        self._scale_mean = scaler.mean_.copy()
        self._scale_std  = scaler.scale_.copy()

        return self

    @classmethod
    def load_default(cls, experimental: bool = False) -> "AdversarialChainDetector":
        """
        Return a pre-trained detector fitted on the FARL validation set.
        Ships with hardcoded coefficients from exp_adversarial_detection.

        .. warning::
            **This detector is NOT a general wrong-answer detector.**

            Validation was performed on 174 self-curated FARL chains (n=86
            adversarial, n=88 correct) with features selected post-hoc after
            seeing the data. CV AUROC 0.9960 reflects FARL domain only.
            Natural-chain production estimate: AUROC 0.397.

            Natural-chain holdout (exp_adversarial_natural_holdout, 2026-03-20):
            AUROC **0.3974** on HP naturally-wrong chains — *negatively
            correlated*. Natural wrong chains score higher than correct chains
            because they have more search retries, not because they are
            adversarially crafted. Do NOT combine with standard risk scoring.

            Use only for detecting FARL-style "confident_wrong" adversarial
            chains. Do not use as a general reliability signal.
        """
        if not experimental:
            raise RuntimeError(
                "AdversarialChainDetector.load_default() requires experimental=True.\n"
                "Natural-chain AUROC=0.3974. Use experimental=True to acknowledge scope limits."
            )
        warnings.warn(
            "AdversarialChainDetector.load_default(): FARL-domain specific detector. "
            "Natural-chain holdout AUROC=0.3974 (negatively correlated with natural errors). "
            "Use only to detect FARL-style adversarial chains, not general wrong answers. "
            "See docs/adversarial_methodology.md for full evaluation.",
            UserWarning,
            stacklevel=2,
        )
        params = AdversarialParams(
            coef=[
                -0.602053, -0.487907, -1.782044,  0.039253,  1.647271,  1.647121,
                 0.007445,  0.655224, -1.361831,  0.033510,  0.365414, -0.044600,
                -1.373435, -0.869402, -1.163310, -0.835141,
            ],
            intercept=-1.757710,
            scale_mean=[
                0.442529, 0.423372, 0.508621, 0.607899, 0.711185, 1.017878,
                0.046466, 0.678180, 0.352490, 0.762356, 0.268199, 0.114943,
                0.105832, 0.189655, 0.082375, 0.738602,
            ],
            scale_std=[
                0.196013, 0.401390, 0.167876, 0.213069, 0.338608, 0.506664,
                0.275086, 0.245821, 0.146097, 0.454454, 0.408527, 0.318953,
                0.249204, 0.181649, 0.175755, 0.469388,
            ],
            threshold=0.50,
            auroc=0.9960,
        )
        return cls.from_params(params)

    # ── Scoring ───────────────────────────────────────────────────────────────

    def _check_fitted(self):
        if self._coef is None:
            raise RuntimeError(
                "AdversarialChainDetector is not fitted. Call fit() or use load_default()."
            )

    def _sigmoid(self, x: float) -> float:
        return float(1.0 / (1.0 + np.exp(-x)))

    def score(self, steps: List[Dict], question: str) -> float:
        """
        Score a single chain.

        Returns
        -------
        float in [0, 1] — probability of being adversarially crafted.
        Higher = more likely adversarial.
        """
        self._check_fitted()
        feat = extract_chain_features(steps, question)
        x_scaled = (feat - self._scale_mean) / np.where(
            self._scale_std > 1e-9, self._scale_std, 1.0
        )
        logit = float(np.dot(self._coef, x_scaled)) + self._intercept
        return self._sigmoid(logit)

    # Training distribution z-score bounds derived from 174 FARL training chains.
    # Features with |z| > _OOD_Z_THRESHOLD on > _OOD_MAX_VIOLATIONS features
    # indicate the chain is likely outside the HP/MuSiQue training distribution.
    _OOD_Z_THRESHOLD = 3.0
    _OOD_MAX_VIOLATIONS = 2  # warn if more than this many features are extreme

    def _check_domain(self, feat: "np.ndarray") -> Optional[str]:
        """Return OOD warning message if chain features are outside training distribution."""
        if self._scale_mean is None or self._scale_std is None:
            return None
        z_scores = np.abs(
            (feat - self._scale_mean) / np.where(self._scale_std > 1e-9, self._scale_std, 1.0)
        )
        violations = int((z_scores > self._OOD_Z_THRESHOLD).sum())
        if violations > self._OOD_MAX_VIOLATIONS:
            extreme_names = [
                f"{ADVERSARIAL_FEATURE_NAMES[i]}(z={z_scores[i]:.1f})"
                for i in np.argsort(-z_scores)[:3]
            ]
            return (
                f"Chain features are outside the training distribution on "
                f"{violations} dimensions (threshold |z|>{self._OOD_Z_THRESHOLD}): "
                f"{', '.join(extreme_names)}. "
                f"Training domain: HotpotQA/MuSiQue FARL chains. "
                f"adversarial_score may not be meaningful for this chain domain."
            )
        return None

    def score_chain(self, steps: List[Dict], question: str) -> AdversarialResult:
        """
        Score a chain and return a full AdversarialResult.

        Emits a UserWarning when chain features fall outside the HP/MuSiQue
        training distribution (|z| > 3.0 on more than 2 features).
        """
        adv_score = self.score(steps, question)
        is_adv    = adv_score >= self._threshold

        # Top 3 contributing features
        feat = extract_chain_features(steps, question)
        x_sc = (feat - self._scale_mean) / np.where(
            self._scale_std > 1e-9, self._scale_std, 1.0
        )
        contributions = self._coef * x_sc
        top_idx = np.argsort(-np.abs(contributions))[:3]
        top_feats = [(ADVERSARIAL_FEATURE_NAMES[i], float(feat[i])) for i in top_idx]

        # Domain guard: warn if out-of-distribution
        ood_msg = self._check_domain(feat)
        if ood_msg:
            warnings.warn(
                f"AdversarialChainDetector.score_chain(): {ood_msg}",
                UserWarning,
                stacklevel=2,
            )

        return AdversarialResult(
            adv_score=adv_score,
            is_adversarial=is_adv,
            routing="adversarial_suspected" if is_adv else "safe",
            top_features=top_feats,
        )

    def predict(self, steps: List[Dict], question: str) -> bool:
        """True if chain is classified as adversarial."""
        return self.score(steps, question) >= self._threshold

    def score_batch(
        self, chains: List[Tuple[List[Dict], str]]
    ) -> List[AdversarialResult]:
        """Score a batch of (steps, question) tuples."""
        return [self.score_chain(steps, q) for steps, q in chains]

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

    def export_params(self) -> AdversarialParams:
        """Export fitted parameters (no raw training data)."""
        self._check_fitted()
        return AdversarialParams(
            coef=[float(x) for x in self._coef],
            intercept=self._intercept,
            scale_mean=[float(x) for x in self._scale_mean],
            scale_std=[float(x) for x in self._scale_std],
            threshold=self._threshold,
            auroc=self._auroc,
        )

    @classmethod
    def from_params(cls, params: AdversarialParams) -> "AdversarialChainDetector":
        """Reconstruct a fitted detector from exported params."""
        det = cls(threshold=params.threshold, experimental=True)  # params implies acceptance
        det._coef       = np.array(params.coef, dtype=float)
        det._intercept  = params.intercept
        det._scale_mean = np.array(params.scale_mean, dtype=float)
        det._scale_std  = np.array(params.scale_std,  dtype=float)
        det._auroc      = params.auroc
        return det

    @property
    def is_fitted(self) -> bool:
        return self._coef is not None
