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

Validated on FARL-generated adversarial chains vs HP-correct chains:
  CV AUROC       : 0.9960 ± 0.0025   CI [0.9906, 0.9996]
  Precision@FPR10: 0.9551

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)
"""
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.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 (exp_adversarial_detection):
      CV AUROC        : 0.9960 ± 0.0025   CI [0.9906, 0.9996]
      Precision@FPR10 : 0.9551
    """

    def __init__(self, threshold: float = 0.50, C: float = 1.0) -> None:
        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) -> "AdversarialChainDetector":
        """
        Return a pre-trained detector fitted on the FARL validation set.
        Ships with hardcoded coefficients from exp_adversarial_detection.
        """
        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)

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

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