"""
nli_grounding.py — NLI-based answer-observation entailment verifier

Problem
-------
Semantic embedding (QARA, StepTransformer) learns domain-specific surface features
and overfits to the training distribution.  The RIGHT task is:
  "Does the final answer follow from the retrieved evidence?"

This is a textual entailment / NLI problem.  A pre-trained NLI cross-encoder
solves it without any domain-specific training.

Solution
--------
Use a pre-trained NLI cross-encoder (cross-encoder/nli-deberta-v3-small) to
measure premise-hypothesis entailment where:
  premise   = question + " " + key_observation  (what the agent retrieved)
  hypothesis = final_answer                      (what the agent concluded)

Three NLI labels map to risk:
  ENTAILMENT    → risk ≈ 0.20  (evidence supports the answer)
  NEUTRAL       → risk ≈ 0.50  (evidence doesn't confirm or deny)
  CONTRADICTION → risk ≈ 0.80  (evidence contradicts the answer)

Cross-domain hypothesis: Since this relies on semantic entailment rather than
surface format features, it should generalise across domains without retraining.

Model selection (in priority order, all CPU-feasible):
  1. cross-encoder/nli-deberta-v3-small  (~180MB, best NLI quality)
  2. cross-encoder/nli-MiniLM2-L6-H768  (~90MB, faster)
  3. cross-encoder/qnli-electra-base     (~135MB, QNLI variant)

Usage
-----
    from llm_guard.nli_grounding import NLIGroundingVerifier

    verifier = NLIGroundingVerifier()
    risk = verifier.score(question, steps, final_answer)
    # risk ∈ [0, 1]; higher = answer less supported by evidence

Validated (exp_nli_grounding):
    HP within-domain AUROC:       TBD
    TV cross-domain AUROC:        TBD
    $0 inference cost (local model)
    No training needed — uses pre-trained model weights
"""

from __future__ import annotations

import warnings
from typing import List, Optional


# NLI label ordering returned by cross-encoder models
# cross-encoder/nli-deberta-v3-small outputs [contradiction, entailment, neutral]
# We detect the actual order at load time.
_LABEL_ORDER_DEFAULT = ["contradiction", "entailment", "neutral"]

# Risk map: lower risk = evidence supports answer
_LABEL_RISK_MAP = {
    "entailment":    0.20,
    "neutral":       0.50,
    "contradiction": 0.80,
}

# Fallback models if the primary is unavailable
_MODEL_PRIORITY = [
    "cross-encoder/nli-deberta-v3-small",
    "cross-encoder/nli-MiniLM2-L6-H768",
    "cross-encoder/qnli-electra-base",
]


class NLIGroundingVerifier:
    """
    Measures answer-observation entailment using a pre-trained NLI cross-encoder.

    Format-agnostic: only uses (question, key_observation, final_answer).
    No training needed — uses pre-trained model weights.

    Model: cross-encoder/nli-deberta-v3-small (~180MB, CPU-feasible)
    Downloads automatically from HuggingFace on first use.

    Parameters
    ----------
    model_name : str
        HuggingFace model ID.  Default: cross-encoder/nli-deberta-v3-small.
        Falls back to nli-MiniLM2 or qnli-electra if the primary fails.

    Notes
    -----
    - Model is lazily loaded on first call to score() or score_batch().
    - Requires: sentence-transformers (pip install sentence-transformers)
    - The premise is constructed as "Question: {q}  Evidence: {obs}".
      The hypothesis is the final_answer.  This framing focuses the NLI
      cross-encoder on factual consistency rather than stylistic similarity.
    - For chains with no usable observations, risk defaults to 0.5 (neutral).
    """

    def __init__(self, model_name: str = "cross-encoder/nli-deberta-v3-small"):
        self._model_name   = model_name
        self._model        = None   # lazy load
        self._label_order: Optional[List[str]] = None
        self._loaded_model_name: Optional[str] = None

    # ── Public API ─────────────────────────────────────────────────────────────

    def score(
        self,
        question: str,
        steps: List[dict],
        final_answer: str,
    ) -> float:
        """
        Return risk score in [0, 1].

        Low  = answer well-supported by retrieved evidence (entailment).
        High = answer contradicts or is unsupported by evidence (contradiction/neutral).

        Parameters
        ----------
        question     : str  — the original question
        steps        : list — ReAct steps with "observation" or "output" fields
        final_answer : str  — the agent's final answer

        Returns
        -------
        float : risk score in [0, 1].
                Returns 0.5 (neutral) if no usable observation found or model unavailable.
        """
        obs = self._extract_key_observation(steps)
        if not obs:
            return 0.5  # no evidence to judge

        model = self._get_model()
        if model is None:
            return 0.5  # transformers not installed

        premise    = f"Question: {question}  Evidence: {obs}"
        hypothesis = final_answer[:500] if final_answer else ""

        if not hypothesis.strip():
            return 0.5

        try:
            scores = model.predict([[premise, hypothesis]])
            risk   = self._scores_to_risk(scores[0])
            return float(risk)
        except Exception as exc:
            warnings.warn(
                f"NLIGroundingVerifier.score() failed: {exc}",
                RuntimeWarning,
                stacklevel=2,
            )
            return 0.5

    def score_batch(
        self,
        chains: List[dict],
    ) -> List[float]:
        """
        Score a batch of chains.  More efficient than calling score() repeatedly
        because it batches the cross-encoder inference.

        Parameters
        ----------
        chains : list of dict, each with "question", "steps", "final_answer".

        Returns
        -------
        list of float : risk scores in the same order as input.
                        Failed chains are assigned 0.5.
        """
        model = self._get_model()
        if model is None:
            return [0.5] * len(chains)

        pairs: List[list]  = []
        valid_idx: List[int] = []
        scores_out = [0.5] * len(chains)

        for i, chain in enumerate(chains):
            obs        = self._extract_key_observation(chain.get("steps", []))
            final_ans  = chain.get("final_answer", "")
            question   = chain.get("question", "")
            if obs and final_ans and question:
                premise    = f"Question: {question}  Evidence: {obs}"
                hypothesis = final_ans[:500]
                pairs.append([premise, hypothesis])
                valid_idx.append(i)

        if not pairs:
            return scores_out

        try:
            raw_scores = model.predict(pairs)
            for pos, chain_idx in enumerate(valid_idx):
                scores_out[chain_idx] = float(self._scores_to_risk(raw_scores[pos]))
        except Exception as exc:
            warnings.warn(
                f"NLIGroundingVerifier.score_batch() failed: {exc}",
                RuntimeWarning,
                stacklevel=2,
            )

        return scores_out

    def is_available(self) -> bool:
        """
        Check if sentence-transformers is installed (required for scoring).
        Does NOT attempt to download the model.
        """
        try:
            import sentence_transformers  # noqa: F401
            return True
        except ImportError:
            return False

    @property
    def model_name(self) -> str:
        """The model name that was successfully loaded (or the configured name)."""
        return self._loaded_model_name or self._model_name

    @property
    def is_loaded(self) -> bool:
        """True if the NLI model has been loaded into memory."""
        return self._model is not None

    # ── Key observation extraction ─────────────────────────────────────────────

    def _extract_key_observation(self, steps: List[dict]) -> str:
        """
        Extract the most informative observation from the agent chain.

        Strategy: use the longest observation (most content).  Longer observations
        typically contain more factual content from search results.
        Fallback: last non-empty observation.

        Parameters
        ----------
        steps : list of step dicts

        Returns
        -------
        str : the selected observation, capped at 500 chars.  Empty string if none.
        """
        obs_list = []
        for s in steps:
            obs = s.get("observation", "") or s.get("output", "")
            if obs and len(obs) > 20:
                obs_list.append(obs)

        if not obs_list:
            return ""

        # Use the longest observation (most informational content)
        key_obs = max(obs_list, key=len)
        return key_obs[:500]

    # ── Model management ───────────────────────────────────────────────────────

    def _get_model(self):
        """
        Lazy-load the NLI cross-encoder.  Tries models in priority order
        if the primary model fails to load.

        Returns
        -------
        CrossEncoder or None if sentence-transformers is not installed.
        """
        if self._model is not None:
            return self._model

        try:
            from sentence_transformers.cross_encoder import CrossEncoder
        except ImportError:
            warnings.warn(
                "NLIGroundingVerifier requires sentence-transformers: "
                "pip install sentence-transformers",
                ImportWarning,
                stacklevel=3,
            )
            return None

        # Try configured model first, then fallback list
        candidates = [self._model_name] + [
            m for m in _MODEL_PRIORITY if m != self._model_name
        ]

        for model_name in candidates:
            try:
                model = CrossEncoder(model_name)
                # Detect label order from model config
                self._label_order = self._detect_label_order(model)
                self._model       = model
                self._loaded_model_name = model_name
                if model_name != self._model_name:
                    warnings.warn(
                        f"NLIGroundingVerifier: primary model '{self._model_name}' "
                        f"unavailable; using fallback '{model_name}'.",
                        RuntimeWarning,
                        stacklevel=3,
                    )
                return model
            except Exception:
                continue

        warnings.warn(
            "NLIGroundingVerifier: all candidate models failed to load. "
            "Scoring will return 0.5 for all chains.",
            RuntimeWarning,
            stacklevel=3,
        )
        return None

    def _detect_label_order(self, model) -> List[str]:
        """
        Detect the label ordering of the loaded NLI model.

        cross-encoder models may output logits in different orders.
        We read from model.config.id2label if available; otherwise use default.

        Returns
        -------
        list of str : label names in the order of model output logits.
        """
        try:
            config = model.model.config
            if hasattr(config, "id2label"):
                labels = [config.id2label[i].lower() for i in range(len(config.id2label))]
                # Normalise: map any label containing "entail" / "contradict" / "neutral"
                normalised = []
                for lbl in labels:
                    if "entail" in lbl:
                        normalised.append("entailment")
                    elif "contradict" in lbl:
                        normalised.append("contradiction")
                    else:
                        normalised.append("neutral")
                return normalised
        except Exception:
            pass
        return _LABEL_ORDER_DEFAULT

    def _scores_to_risk(self, logits) -> float:
        """
        Convert raw logits (or probabilities) from the NLI model to a risk score.

        Steps:
          1. Apply softmax if the values look like logits (any value > 1 or < 0).
          2. Map each label's probability to a risk using _LABEL_RISK_MAP.
          3. Return the probability-weighted average risk.

        Parameters
        ----------
        logits : array-like of shape (n_labels,)

        Returns
        -------
        float : risk score in [0, 1].
        """
        import numpy as np

        arr = np.array(logits, dtype=float)

        # Apply softmax if values don't look like probabilities
        if arr.min() < 0 or arr.max() > 1 or abs(arr.sum() - 1.0) > 0.01:
            arr = arr - arr.max()  # numerical stability
            arr = np.exp(arr)
            arr = arr / arr.sum()

        label_order = self._label_order or _LABEL_ORDER_DEFAULT

        # Ensure we have the right number of labels
        n = min(len(arr), len(label_order))
        risk = 0.0
        for i in range(n):
            lbl  = label_order[i]
            prob = float(arr[i])
            risk += prob * _LABEL_RISK_MAP.get(lbl, 0.5)

        # Normalise if label count mismatched
        if n < len(arr):
            remainder = float(arr[n:].sum())
            risk += remainder * 0.5  # neutral for unknown labels

        return float(np.clip(risk, 0.0, 1.0))

    def __repr__(self) -> str:
        return (
            f"NLIGroundingVerifier("
            f"model='{self.model_name}', "
            f"loaded={self.is_loaded}"
            f")"
        )
