"""
StepTransformerVerifier — 2-layer Transformer encoder over MiniLM step embeddings.

Replaces LSTMRiskAccumulator with a format-agnostic sequence model.  Each step is
encoded as free text (thought + action + observation[:150]) → MiniLM-L6-v2
(frozen, 384-dim) → linear projection → TransformerEncoder → [CLS] head → risk.

Key design decisions
--------------------
- Format-agnostic step encoding: works with ReAct (thought/action/obs) AND
  function-calling (tool_call/result) and any dict-based step format.
- Frozen backbone: only the projection + Transformer + head are trained.
  Backbone weights are shared with MiniJudge (already a required dep).
- Tiny model (~150K trainable params): CPU-feasible, 30-epoch training on 200 chains
  takes < 3 minutes on a laptop.
- [CLS] token at position-0 aggregates the full chain.
- Positional embeddings are learned (not sinusoidal) so the model can adapt to
  variable-length chains without relying on step index as a feature.
- Fallback path (no torch): degrades to a sklearn LogReg on mean-pooled MiniLM
  embeddings + 3 hand-crafted step-count features.

Validated performance (exp_step_transformer)
--------------------------------------------
    Within-domain (HP 200 chains, 5-fold CV): reported in results.json
    Cross-domain (HP→TV):                     reported in results.json

Requirements
------------
    sentence-transformers (already required by llm-guard-kit)
    torch >= 1.12  (optional; fallback activates if absent)

Quick start
-----------
    from llm_guard import StepTransformerVerifier

    v = StepTransformerVerifier()
    v.fit(labeled_runs)   # [{"question", "steps", "final_answer", "correct": bool}]
    risk = v.score(question, steps, final_answer)  # float [0, 1]

    # Save / load
    v.save("step_transformer.pkl")
    v = StepTransformerVerifier.load("step_transformer.pkl")
Purpose:    Encodes each step as free text, projects through MiniLM-L6-v2 (frozen), and aggregates via a Transformer [CLS] head. Beats LSTM within-domain (0.737 vs 0.723) but under-performs cross-domain (0.658 vs 0.744).
Exports:    StepTransformerVerifier
Added:      v0.25.0
Status:     Experimental
"""

from __future__ import annotations

import pickle
import warnings
from typing import Dict, List, Optional

import numpy as np

# ── Universal step encoder ────────────────────────────────────────────────────

def encode_step(step: dict) -> str:
    """
    Convert any agent step dict to a short text snippet for MiniLM encoding.

    Handles ReAct (thought / action_arg / observation),
    OpenAI tool-call (content / tool_calls / tool result),
    LangChain (log / tool_input / observation),
    and custom formats with 'reasoning', 'input', 'output' keys.

    Lengths are capped so the full snippet fits inside MiniLM's 256-token window.
    """
    thought = (
        step.get("thought", "")
        or step.get("reasoning", "")
        or step.get("thinking", "")
        or step.get("log", "")
        or ""
    )
    action = (
        step.get("action_arg", "")
        or step.get("tool_input", "")
        or step.get("input", "")
        or step.get("query", "")
        or ""
    )
    obs = (
        step.get("observation", "")
        or step.get("output", "")
        or step.get("result", "")
        or ""
    )
    parts = [
        str(thought)[:100].strip(),
        str(action)[:100].strip(),
        str(obs)[:150].strip(),
    ]
    return " ".join(p for p in parts if p).strip()


# ── Lazy sentence-transformer singleton ──────────────────────────────────────

class _MiniLM:
    """Singleton wrapper; loads once and reuses across instances."""

    _instance: Optional["_MiniLM"] = None

    def __init__(self) -> None:
        self._model = None
        self._available: bool = False
        try:
            from sentence_transformers import SentenceTransformer  # type: ignore
            self._model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
            self._available = True
        except Exception:
            self._available = False

    @classmethod
    def get(cls) -> "_MiniLM":
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    def encode(self, texts: List[str]) -> np.ndarray:
        """Return (N, 384) float32 array; zeros if model unavailable."""
        if self._available and self._model is not None:
            vecs = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
            return vecs.astype(np.float32)
        return np.zeros((len(texts), 384), dtype=np.float32)

    @property
    def available(self) -> bool:
        return self._available


# ── Torch model definition ────────────────────────────────────────────────────

def _build_torch_model(
    d_model: int = 64,
    nhead: int = 4,
    num_layers: int = 2,
    dim_ff: int = 128,
    dropout: float = 0.1,
    max_seq: int = 12,
    embed_dim: int = 384,
):
    """
    Build and return the Transformer chain scorer (PyTorch).

    Architecture
    ------------
    Linear(384→64) projection
    Learned positional embedding (max_seq+1 positions; +1 for [CLS])
    2-layer TransformerEncoder: d_model=64, nhead=4, dim_ff=128, dropout=0.1
    Linear(64→1) head + sigmoid
    """
    try:
        import torch
        import torch.nn as nn
    except ImportError as e:
        raise ImportError(
            "PyTorch is required for StepTransformerVerifier's Transformer mode. "
            "Install with: pip install torch\n"
            "Alternatively, the sklearn fallback is used automatically when torch is absent."
        ) from e

    class _TransformerChainScorer(nn.Module):
        def __init__(
            inner,
            d_model: int = d_model,
            nhead: int = nhead,
            num_layers: int = num_layers,
            dim_ff: int = dim_ff,
            dropout: float = dropout,
            max_seq: int = max_seq,
            embed_dim: int = embed_dim,
        ):
            super().__init__()
            inner.proj = nn.Linear(embed_dim, d_model)
            inner.pos_emb = nn.Embedding(max_seq + 1, d_model)  # +1 for [CLS] at pos 0
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_ff,
                dropout=dropout,
                batch_first=True,
            )
            inner.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
            inner.head = nn.Linear(d_model, 1)
            inner.max_seq = max_seq
            inner.d_model = d_model

        def forward(inner, step_embeds: "torch.Tensor", lengths: List[int]) -> "torch.Tensor":
            """
            step_embeds : (B, T, 384) — pre-computed MiniLM embeddings (frozen backbone)
            lengths     : list of int — actual number of steps per chain (before padding)
            Returns     : (B,) sigmoid risk scores
            """
            import torch

            B, T, _ = step_embeds.shape

            # Project step embeddings
            x = inner.proj(step_embeds)  # (B, T, d_model)

            # Prepend learnable [CLS] token (zero-init projected embedding at pos 0)
            cls_tok = torch.zeros(B, 1, inner.d_model, device=x.device)
            x = torch.cat([cls_tok, x], dim=1)  # (B, T+1, d_model)

            # Add positional embeddings
            seq_len = x.shape[1]
            pos_ids = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(B, -1)
            x = x + inner.pos_emb(pos_ids)

            # Build padding mask (True = masked = ignore): positions beyond actual length
            # +1 because we prepended [CLS]
            pad_mask = torch.ones(B, T + 1, dtype=torch.bool, device=x.device)
            pad_mask[:, 0] = False  # [CLS] is always valid
            for i, l in enumerate(lengths):
                valid = min(int(l), T)
                pad_mask[i, 1:valid + 1] = False  # steps 1..l are valid

            x = inner.transformer(x, src_key_padding_mask=pad_mask)  # (B, T+1, d_model)

            # Use [CLS] token output (position 0) as chain representation
            cls_out = x[:, 0, :]  # (B, d_model)
            return torch.sigmoid(inner.head(cls_out)).squeeze(-1)  # (B,)

    return _TransformerChainScorer()


# ── Sklearn fallback classifier ───────────────────────────────────────────────

def _build_sklearn_fallback():
    """Return a sklearn Pipeline for the no-torch fallback."""
    from sklearn.linear_model import LogisticRegression
    from sklearn.pipeline import Pipeline
    from sklearn.preprocessing import StandardScaler

    return Pipeline([
        ("sc", StandardScaler()),
        ("clf", LogisticRegression(
            C=1.0, max_iter=1000, random_state=42, class_weight="balanced"
        )),
    ])


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

class StepTransformerVerifier:
    """
    2-layer Transformer encoder over frozen MiniLM step embeddings.

    Architecture
    ------------
    MiniLM-L6-v2 (frozen, 384-dim) → Linear(384→64) → positional embedding
    → 2-layer TransformerEncoder (nhead=4, ff=128, dropout=0.1)
    → [CLS] at position 0 → Linear(64→1) → sigmoid → risk ∈ [0, 1]

    Format-agnostic: works with ReAct (thought/action/obs) and function-calling
    (tool_call/result) chains.  Any step dict with observation/output/result
    fields is handled via the universal ``encode_step`` function.

    Cross-domain advantage over LSTM
    ---------------------------------
    The LSTM's per-step hand-crafted Jaccard features are format-specific and
    don't transfer across domains (cross-domain AUROC 0.545, exp143).
    MiniLM embeddings capture semantic content irrespective of step-count or
    chain format, so domain gap should be smaller.

    Fallback
    --------
    When torch is not installed, a sklearn LogReg on mean-pooled MiniLM step
    embeddings is used instead.  Performance may be lower but the API is
    identical.

    Parameters
    ----------
    max_steps : int
        Maximum steps to encode per chain (excess steps are truncated).
        Default: 10.
    d_model : int
        Transformer hidden dimension. Default: 64.
    epochs : int
        Training epochs. Default: 30.
    lr : float
        Adam learning rate. Default: 1e-3.
    random_state : int
        Seed for reproducibility. Default: 42.
    """

    D_MODEL     = 64
    NHEAD       = 4
    NUM_LAYERS  = 2
    DIM_FF      = 128
    DROPOUT     = 0.1
    MAX_STEPS   = 10
    BATCH_SIZE  = 32
    EPOCHS      = 30
    LR          = 1e-3

    def __init__(
        self,
        max_steps: int = MAX_STEPS,
        d_model: int = D_MODEL,
        epochs: int = EPOCHS,
        lr: float = LR,
        random_state: int = 42,
    ) -> None:
        self.max_steps    = max_steps
        self.d_model      = d_model
        self.epochs       = epochs
        self.lr           = lr
        self.random_state = random_state

        self._model       = None          # torch model
        self._fallback    = None          # sklearn pipeline (no-torch mode)
        self._use_torch   = False
        self._fitted      = False

    # ── Feature prep ─────────────────────────────────────────────────────────

    def _embed_chain(self, steps: List[dict]) -> np.ndarray:
        """
        Return (max_steps, 384) padded MiniLM embedding matrix for a chain.
        Also returns the actual number of valid steps (before padding).
        """
        snippets = [encode_step(s) for s in steps[:self.max_steps]]
        if not snippets:
            snippets = [""]
        miniLM = _MiniLM.get()
        vecs = miniLM.encode(snippets)  # (k, 384), k ≤ max_steps
        n_real = vecs.shape[0]
        # Pad to max_steps
        if n_real < self.max_steps:
            pad = np.zeros((self.max_steps - n_real, 384), dtype=np.float32)
            vecs = np.concatenate([vecs, pad], axis=0)
        return vecs, n_real  # (max_steps, 384), int

    def _embed_chain_mean(self, steps: List[dict]) -> np.ndarray:
        """Mean-pooled embedding for the sklearn fallback + 3 structural features."""
        snippets = [encode_step(s) for s in steps[:self.max_steps]] or [""]
        miniLM = _MiniLM.get()
        vecs = miniLM.encode(snippets)
        mean_emb = vecs.mean(0)  # (384,)
        # 3 structural features
        n_norm  = min(len(steps) / 10.0, 1.0)
        has_obs = np.mean([
            1.0 if (s.get("observation","") or s.get("output","") or "").strip() else 0.0
            for s in steps
        ]) if steps else 0.0
        empty_obs = 1.0 - has_obs
        return np.concatenate([mean_emb, [n_norm, has_obs, empty_obs]]).astype(np.float32)

    # ── Training ─────────────────────────────────────────────────────────────

    def fit(self, labeled_runs: List[Dict]) -> "StepTransformerVerifier":
        """
        Train on a list of labeled chain dicts.

        Parameters
        ----------
        labeled_runs : list of dict
            Each dict must have:
                "steps"        : list of step dicts
                "final_answer" : str  (or "answer")
                "correct"      : bool  (True = answer is correct)

        Returns
        -------
        self
        """
        np.random.seed(self.random_state)

        labels = np.array(
            [0 if r.get("correct", True) else 1 for r in labeled_runs],
            dtype=np.float32,
        )

        # Try torch first
        try:
            import torch
            import torch.nn as nn
            import torch.optim as optim
            self._use_torch = True
        except ImportError:
            self._use_torch = False
            warnings.warn(
                "torch not installed — StepTransformerVerifier falling back to "
                "sklearn LogReg on mean-pooled MiniLM embeddings. "
                "Install torch for the full Transformer model: pip install torch",
                UserWarning,
                stacklevel=2,
            )

        if self._use_torch:
            self._fit_torch(labeled_runs, labels)
        else:
            self._fit_sklearn(labeled_runs, labels)

        self._fitted = True
        return self

    def _fit_torch(self, runs: List[Dict], labels: np.ndarray) -> None:
        import torch
        import torch.nn as nn
        import torch.optim as optim

        torch.manual_seed(self.random_state)

        # Build embeddings (CPU; no gradient)
        all_vecs, all_n = [], []
        for r in runs:
            vecs, n_real = self._embed_chain(r.get("steps", []))
            all_vecs.append(vecs)
            all_n.append(n_real)

        X_t = torch.FloatTensor(np.stack(all_vecs, axis=0))  # (N, T, 384)
        y_t = torch.FloatTensor(labels)
        n   = len(labels)

        # Class-imbalance weight
        n_pos  = labels.sum()
        n_neg  = n - n_pos
        pos_w  = float(n_neg / max(n_pos, 1))
        neg_w  = 1.0
        sample_w = torch.FloatTensor(
            [pos_w if y == 1 else neg_w for y in labels]
        )

        self._model = _build_torch_model(
            d_model=self.d_model,
            nhead=self.NHEAD,
            num_layers=self.NUM_LAYERS,
            dim_ff=self.DIM_FF,
            dropout=self.DROPOUT,
            max_seq=self.max_steps,
            embed_dim=384,
        )
        opt     = optim.Adam(self._model.parameters(), lr=self.lr)
        loss_fn = nn.BCELoss(reduction="none")

        for epoch in range(self.epochs):
            idx = np.random.permutation(n)
            self._model.train()
            for start in range(0, n, self.BATCH_SIZE):
                b_idx = idx[start:start + self.BATCH_SIZE]
                bx    = X_t[b_idx]                          # (B, T, 384)
                by    = y_t[b_idx]
                bw    = sample_w[b_idx]
                bl    = [all_n[i] for i in b_idx]

                pred  = self._model(bx, bl)                 # (B,)
                raw   = loss_fn(pred, by)                   # (B,)
                loss  = (raw * bw).mean()
                opt.zero_grad()
                loss.backward()
                opt.step()

    def _fit_sklearn(self, runs: List[Dict], labels: np.ndarray) -> None:
        """Fallback: LogReg on mean-pooled MiniLM embeddings."""
        X = np.stack([
            self._embed_chain_mean(r.get("steps", []))
            for r in runs
        ])
        self._fallback = _build_sklearn_fallback()
        self._fallback.fit(X, labels.astype(int))

    # ── Inference ─────────────────────────────────────────────────────────────

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

        Parameters
        ----------
        question     : str — the user question
        steps        : list of step dicts (any supported format)
        final_answer : str — the agent's final answer

        Returns
        -------
        risk : float ∈ [0, 1] — failure probability (higher = more likely wrong)
        """
        if not self._fitted:
            raise RuntimeError("Call fit() before score().")

        if self._use_torch:
            return self._score_torch(steps)
        return self._score_sklearn(steps)

    def _score_torch(self, steps: List[dict]) -> float:
        import torch

        vecs, n_real = self._embed_chain(steps)
        x = torch.FloatTensor(vecs).unsqueeze(0)  # (1, T, 384)
        self._model.eval()
        with torch.no_grad():
            pred = self._model(x, [n_real])
        return float(pred.item())

    def _score_sklearn(self, steps: List[dict]) -> float:
        x = self._embed_chain_mean(steps).reshape(1, -1)
        return float(self._fallback.predict_proba(x)[0, 1])

    def score_run(self, run: dict) -> float:
        """Convenience wrapper: score a run dict directly."""
        return self.score(
            run.get("question", ""),
            run.get("steps", []),
            run.get("final_answer", run.get("answer", "")),
        )

    # ── Persistence ──────────────────────────────────────────────────────────

    def save(self, path: str) -> None:
        """Save the fitted model to a file."""
        state = {
            "max_steps":    self.max_steps,
            "d_model":      self.d_model,
            "epochs":       self.epochs,
            "lr":           self.lr,
            "random_state": self.random_state,
            "use_torch":    self._use_torch,
            "fitted":       self._fitted,
        }
        if self._use_torch and self._model is not None:
            import torch
            import io
            buf = io.BytesIO()
            torch.save(self._model.state_dict(), buf)
            state["model_state"] = buf.getvalue()
        elif self._fallback is not None:
            state["fallback"] = self._fallback

        with open(path, "wb") as f:
            pickle.dump(state, f)

    @classmethod
    def load(cls, path: str) -> "StepTransformerVerifier":
        """Load a previously saved StepTransformerVerifier."""
        with open(path, "rb") as f:
            state = pickle.load(f)

        obj = cls(
            max_steps=state["max_steps"],
            d_model=state["d_model"],
            epochs=state["epochs"],
            lr=state["lr"],
            random_state=state["random_state"],
        )
        obj._use_torch = state["use_torch"]
        obj._fitted    = state["fitted"]

        if state["use_torch"] and "model_state" in state:
            import torch
            import io
            obj._model = _build_torch_model(
                d_model=state["d_model"],
                nhead=StepTransformerVerifier.NHEAD,
                num_layers=StepTransformerVerifier.NUM_LAYERS,
                dim_ff=StepTransformerVerifier.DIM_FF,
                dropout=StepTransformerVerifier.DROPOUT,
                max_seq=state["max_steps"],
                embed_dim=384,
            )
            buf = io.BytesIO(state["model_state"])
            obj._model.load_state_dict(torch.load(buf, map_location="cpu"))
        elif "fallback" in state:
            obj._fallback = state["fallback"]

        return obj

    # ── Repr ─────────────────────────────────────────────────────────────────

    def __repr__(self) -> str:
        mode   = "Transformer (torch)" if self._use_torch else "LogReg fallback (sklearn)"
        status = "fitted" if self._fitted else "not fitted"
        return (
            f"StepTransformerVerifier("
            f"max_steps={self.max_steps}, d_model={self.d_model}, "
            f"mode={mode}, {status})"
        )
