"""Clause-level NLI shadow check — #000049 §7 #5 algorithm, Phase 2.

For a (claim, source) pair the lexical verifier resolved only *weakly*
(paraphrase path, or an entity path with one verified unit), run a
sentence-pair NLI model over each source clause and report whether a
*demotion-only contradiction veto* would fire:

    would_demote  iff  max_contradiction >= theta_contra
                  and  max_entailment   <  theta_entail

The second conjunct (the entailment guard) is load-bearing — a source
can carry both a clause that loosely contradicts the claim and a clause
that directly entails it (the competing-superlative case from the
#000049 §7 #18 bench); the guard suppresses the veto when some clause
entails. Defaults theta_contra=0.5, theta_entail=0.9 are the
bench-validated point (arborist-nli-bench a1cb28d).

This is SHADOW infrastructure: the result is telemetry, never an
``audit_mode`` input. The model runtime is the optional ``[nli]``
extra; everything degrades to ``available=False`` when it is missing.
"""
from __future__ import annotations

import json
import os
import re
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Optional

_MANIFEST_PATH = Path(__file__).resolve().parent / "manifest.json"
_SENT_SPLIT = re.compile(r"(?<=[.!?;])\s+")
_WORD_RE = re.compile(r"[A-Za-z][A-Za-z'-]+")
# A tiny stopword set — just the function words long enough to clear the
# >3-char length filter and frequent enough to dominate overlap noise.
_STOP = frozenset((
    "the", "and", "that", "this", "with", "from", "have", "has", "had",
    "was", "were", "are", "been", "being", "for", "into", "than", "then",
    "they", "them", "their", "there", "which", "what", "when", "where",
    "who", "whom", "whose", "while", "also", "such", "some", "any", "all",
    "more", "most", "other", "over", "under", "between", "about", "because",
    "since", "though", "although", "however", "would", "could", "should",
    "will", "shall", "can", "may", "might", "must", "does", "did", "not",
))


def load_manifest(path: Optional[Path] = None) -> dict:
    return json.loads((path or _MANIFEST_PATH).read_text())


def clauses(text: str) -> list[str]:
    """Split source text into clauses (the §7 #5 step-1 segmenter).

    Standalone here (a sibling-repo / bench-callable surface) rather than
    importing ``parse_claims``' pointer-aware segmenter — the shadow
    check runs on raw retrieved context, which has no pointer markup.
    """
    parts = [p.strip() for p in _SENT_SPLIT.split((text or "").strip()) if p.strip()]
    return parts or ([text.strip()] if text and text.strip() else [])


def _content_tokens(text: str) -> set[str]:
    return {w for w in (m.group(0).lower() for m in _WORD_RE.finditer(text or ""))
            if len(w) > 3 and w not in _STOP}


def candidate_clauses(claim: str, source: str, k: int) -> list[tuple[str, float]]:
    """The §7 #5 step-3 "candidate source clauses": clauses that share
    content tokens with the answer claim, ranked by overlap fraction,
    top ``k``, zero-overlap dropped.

    Restricting NLI to this small set is load-bearing — running it over
    *every* clause of a 30 KB retrieved context is a multiple-
    comparisons FP machine (ticket #000049 §7 #20). Standalone lexical
    proxy here; a runtime (Phase 3) hook would instead pass the clauses
    the verifier's quote/span/entity matching actually touched.
    """
    ctoks = _content_tokens(claim)
    if not ctoks:
        return []
    scored: list[tuple[str, float]] = []
    for cl in clauses(source):
        ov = ctoks & _content_tokens(cl)
        if ov:
            scored.append((cl, len(ov) / len(ctoks)))
    scored.sort(key=lambda t: t[1], reverse=True)
    return scored[:max(1, k)]


@dataclass
class ShadowResult:
    available: bool                 # False iff the [nli] extra / model could not load
    would_demote: bool              # the §7 #5 Demote() decision (shadow — not applied)
    max_contradiction: float
    max_entailment: float
    best_clause: Optional[str]      # the candidate clause with the highest contradiction prob
    n_clauses: int                  # total clauses in the source
    n_candidate_clauses: int        # clauses NLI was actually run on (§7 #5 step 3 restriction)
    best_clause_overlap: float      # max content-token overlap of any candidate clause with the claim
    recombination_risk: bool        # claim tokens scattered across >=2 candidate clauses, no single dominant one
    model_version: Optional[str]
    theta_contra: float
    theta_entail: float
    reason: str = ""                # human note ("ok", "deps_missing: …", "no_candidate_clauses", …)

    def as_dict(self) -> dict:
        return asdict(self)


def _softmax(xs: list[float]) -> list[float]:
    import math
    m = max(xs)
    es = [math.exp(x - m) for x in xs]
    s = sum(es)
    return [e / s for e in es]


def _resolve_label_indices(id2label: dict) -> tuple[int, int, int]:
    inv: dict[str, int] = {}
    for i, lab in id2label.items():
        low = str(lab).lower()
        if "entail" in low:
            inv["e"] = int(i)
        elif "neutral" in low:
            inv["n"] = int(i)
        elif "contra" in low:
            inv["c"] = int(i)
    missing = {"e", "n", "c"} - set(inv)
    if missing:
        raise RuntimeError(f"cannot map NLI labels from id2label={id2label!r} (missing {missing})")
    return inv["e"], inv["n"], inv["c"]


class ShadowNLI:
    """Lazily-loaded pinned NLI model + the §7 #5 clause-level check.

    Construction never raises: if the ``[nli]`` extra is missing or the
    pinned checkpoint will not load, :attr:`available` stays False and
    :meth:`check` returns an unavailable :class:`ShadowResult`.
    """

    def __init__(self, manifest: Optional[dict] = None, device: Optional[str] = None):
        self.manifest = manifest or load_manifest()
        self.model_version: Optional[str] = self.manifest.get("nli_model_version")
        th = self.manifest.get("thresholds", {})
        self.theta_contra: float = float(th.get("contradiction_veto", 0.5))
        self.theta_entail: float = float(th.get("entailment_block_veto", 0.9))
        self.max_length: int = int(self.manifest.get("max_length", 256))
        # §7 #5 step-3 / §7 #20 — cap how many source clauses NLI runs on.
        self.max_candidate_clauses: int = int(self.manifest.get("max_candidate_clauses", 6))
        # device: explicit arg > ARBORIST_NLI_DEVICE env > auto (cuda if available, else cpu).
        self.device_pref: str = device or os.environ.get("ARBORIST_NLI_DEVICE") or "auto"
        self.device: Optional[str] = None  # resolved at load
        self.batch_size: int = int(os.environ.get("ARBORIST_NLI_BATCH", "64"))
        # opt-in engine override (§7 #28). Default "" = the proven auto path
        # (onnx-int8 export if present, else torch). "tinygrad" is opt-in
        # only: its ONNX-frontend op-coverage for the large MNLI
        # cross-encoders is unverified, so it never auto-preempts the
        # proven path — honest-verdicts discipline + the cage (engine is
        # never a trust dependency; see docs/onnx-vendor-capture-immunity.md).
        self.backend_pref: str = (os.environ.get("ARBORIST_NLI_BACKEND") or "").strip().lower()
        self.backend: Optional[str] = None  # "onnx" | "onnx-int8" | "torch" | "tinygrad", set at load
        self.available = False
        self._reason = "uninitialised"
        self._tok = None
        self._model = None
        self._ei = self._ni = self._ci = None

    def _onnx_dir(self) -> Optional[Path]:
        """Where an ONNX export of the pinned checkpoint would live, if any.
        `ARBORIST_NLI_ONNX_DIR` overrides; default is
        `~/.arborist/models/nli/<model_version>/onnx/` (populated by
        `bench/scripts/export_nli_onnx.py`)."""
        env = os.environ.get("ARBORIST_NLI_ONNX_DIR")
        if env:
            return Path(env)
        mv = self.model_version or "nli"
        return Path.home() / ".arborist" / "models" / "nli" / mv / "onnx"

    def _ensure_loaded(self) -> None:
        if self.available or self._reason.startswith(("deps_missing", "load_failed")):
            return
        try:
            import torch  # noqa: F401
            from transformers import AutoTokenizer, AutoModelForSequenceClassification
        except ImportError as e:
            self._reason = f"deps_missing: {e} (install: pip install 'arborist[nli]')"
            return
        repo = self.manifest["hf_repo"]
        rev = self.manifest.get("pinned_revision")
        # resolve device
        if self.device_pref == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = self.device_pref
        # §7 #28 — opt-in tinygrad ONNX-frontend backend. Tried ONLY when
        # ARBORIST_NLI_BACKEND=tinygrad is set; never auto-preferred (its
        # op-coverage for these checkpoints is unverified — measured by
        # `make bench-nli-backends`, not assumed). SHADOW only; the engine
        # swap is invisible to the proof path by construction. Any failure
        # falls through to the proven onnx/torch auto path below.
        if self.backend_pref == "tinygrad":
            try:
                self._tok, self._model = self._try_tinygrad()
                if self._model is not None:
                    self.backend = "tinygrad"
            except Exception:  # noqa: BLE001 — best-effort; fall through
                self._tok = self._model = None
                self.backend = None

        # prefer an ONNX export if one exists and `optimum` is importable
        # (§3.1 of the speedup plan — ~2-4x on CPU, drops the torch fwd path;
        # on cuda uses CUDAExecutionProvider). Falls back to torch silently.
        onnx_dir = self._onnx_dir()
        try:
            if self._model is None and onnx_dir and onnx_dir.exists():
                from optimum.onnxruntime import ORTModelForSequenceClassification
                provider = ("CUDAExecutionProvider" if self.device == "cuda"
                            else "CPUExecutionProvider")
                self._tok = AutoTokenizer.from_pretrained(str(onnx_dir))
                # prefer the int8-quantized graph if the export produced one
                # (the §3 CPU speedup); else the fp32 model.onnx.
                kw = {}
                if (onnx_dir / "model_quantized.onnx").exists():
                    kw["file_name"] = "model_quantized.onnx"
                self._model = ORTModelForSequenceClassification.from_pretrained(
                    str(onnx_dir), provider=provider, **kw)
                self.backend = "onnx-int8" if kw else "onnx"
        except Exception:  # noqa: BLE001 — ONNX is best-effort; fall through to torch
            self._tok = self._model = None
            self.backend = None
        if self._model is None:
            try:
                self._tok = AutoTokenizer.from_pretrained(repo, revision=rev)
                self._model = AutoModelForSequenceClassification.from_pretrained(repo, revision=rev)
                self._model.eval()
                if self.device == "cuda":
                    self._model = self._model.to("cuda")
                self.backend = "torch"
            except Exception as e:  # noqa: BLE001 — any load failure is "unavailable"
                self._reason = f"load_failed: {type(e).__name__}: {e}"
                self._tok = self._model = None
                return
        try:
            id2label = {int(k): v for k, v in self._model.config.id2label.items()}
            self._ei, self._ni, self._ci = _resolve_label_indices(id2label)
        except Exception as e:  # noqa: BLE001
            self._reason = f"load_failed: {type(e).__name__}: {e}"
            self._tok = self._model = None
            return
        self.available = True
        self._reason = "ok"

    def _try_tinygrad(self):
        """§7 #28 — load the pinned checkpoint's ONNX export through
        tinygrad's ONNX frontend, wrapped so it satisfies the exact
        ``model(**enc).logits`` (torch tensor) contract ``_nli_batch``
        already uses — so the proven batched forward path is byte-
        unchanged (zero regression risk to torch/onnx).

        Returns ``(tokenizer, adapter)`` or ``(None, None)``. tinygrad
        is *not* an arborist dependency; this fails closed (op-coverage
        for the large MNLI cross-encoders is unverified by design —
        `make bench-nli-backends` measures it, this does not assume it).
        Prefers the fp32 ``model.onnx`` (int8 QLinear* ops are unlikely
        to be covered by the frontend).
        """
        onnx_dir = self._onnx_dir()
        if not (onnx_dir and onnx_dir.exists()):
            return None, None
        onnx_path = onnx_dir / "model.onnx"
        if not onnx_path.exists():
            return None, None

        import numpy as np
        import onnx as _onnx
        import torch
        from transformers import AutoTokenizer, AutoConfig

        proto = _onnx.load(str(onnx_path))
        graph_inputs = [i.name for i in proto.graph.input]
        out_names = [o.name for o in proto.graph.output]
        logit_out = next((n for n in out_names if "logit" in n.lower()),
                         out_names[0] if out_names else None)

        # Defensive multi-version frontend import. tinygrad has moved
        # this module across releases; try the known paths in order.
        runner = None
        try:
            from tinygrad.frontend.onnx import OnnxRunner  # newer tinygrad
            runner = OnnxRunner(proto)
            _kind = "OnnxRunner"
        except Exception:  # noqa: BLE001
            try:
                from extra.onnx import get_run_onnx  # older tinygrad layout
                runner = get_run_onnx(proto)
                _kind = "get_run_onnx"
            except Exception:  # noqa: BLE001
                return None, None

        tok = AutoTokenizer.from_pretrained(str(onnx_dir))
        cfg = AutoConfig.from_pretrained(str(onnx_dir))

        class _TinygradAdapter:
            """Mimics a HF seq-classification model's call contract."""

            def __init__(self):
                self.config = cfg

            def __call__(self, **enc):
                feed = {}
                for k, v in enc.items():
                    if k not in graph_inputs:
                        continue
                    arr = v.detach().cpu().numpy() if hasattr(v, "detach") else np.asarray(v)
                    feed[k] = arr
                out = runner(feed)
                if isinstance(out, dict):
                    val = out.get(logit_out, next(iter(out.values())))
                else:
                    val = out
                if hasattr(val, "numpy"):      # tinygrad Tensor
                    val = val.numpy()
                logits = torch.as_tensor(np.asarray(val, dtype=np.float32))
                import types
                return types.SimpleNamespace(logits=logits)

        return tok, _TinygradAdapter()

    def _nli_batch(self, pairs: list[tuple[str, str]]) -> list[tuple[float, float, float]]:
        """One forward pass over a list of (premise, hypothesis) pairs.
        Replaces N batch-1 forwards with ⌈N/batch_size⌉ batched ones —
        the §3.1 speedup; on CPU ~3-5x, on cuda far more."""
        import torch
        out: list[tuple[float, float, float]] = []
        for i in range(0, len(pairs), self.batch_size):
            chunk = pairs[i:i + self.batch_size]
            prem = [p for p, _ in chunk]
            hyp = [h for _, h in chunk]
            enc = self._tok(prem, hyp, return_tensors="pt", truncation=True,
                            padding=True, max_length=self.max_length)
            if self.device == "cuda" and self.backend == "torch":
                enc = {k: v.to("cuda") for k, v in enc.items()}
            with torch.no_grad():
                logits = self._model(**enc).logits
            logits = logits.detach().cpu().tolist()
            for row in logits:
                p = _softmax(row)
                out.append((p[self._ei], p[self._ni], p[self._ci]))
        return out

    def _nli(self, premise: str, hypothesis: str) -> tuple[float, float, float]:
        return self._nli_batch([(premise, hypothesis)])[0]

    def check(self, claim: str, source: str) -> ShadowResult:
        n_clauses = len(clauses(source))
        cand = candidate_clauses(claim, source, self.max_candidate_clauses)
        best_overlap = cand[0][1] if cand else 0.0
        # recombination risk: claim tokens scattered across >=2 candidate
        # clauses, none of which alone dominates the claim. (Not a gate
        # here — recorded so the bench can split the would-demote rate by
        # it and confirm whether the §7 #5 risk-gate would help.)
        recomb_risk = (len(cand) >= 2 and cand[1][1] >= 0.2 and best_overlap < 0.85)

        def _result(**kw) -> ShadowResult:
            base = dict(n_clauses=n_clauses, n_candidate_clauses=len(cand),
                        best_clause_overlap=round(best_overlap, 4), recombination_risk=recomb_risk,
                        model_version=self.model_version, theta_contra=self.theta_contra,
                        theta_entail=self.theta_entail)
            base.update(kw)
            return ShadowResult(**base)  # type: ignore[arg-type]

        self._ensure_loaded()
        if not self.available:
            return _result(available=False, would_demote=False, max_contradiction=0.0,
                           max_entailment=0.0, best_clause=None, reason=self._reason)
        if not (claim or "").strip() or n_clauses == 0:
            return _result(available=True, would_demote=False, max_contradiction=0.0,
                           max_entailment=0.0, best_clause=None, reason="empty_input")
        if not cand:
            # nothing in the source shares content tokens with the claim →
            # the lexical verifier would not have matched it either; no
            # contradiction veto. (This is *not* "grounded" — it's "NLI has
            # nothing to say"; UNGROUNDED-by-no-overlap is the verifier's job.)
            return _result(available=True, would_demote=False, max_contradiction=0.0,
                           max_entailment=0.0, best_clause=None, reason="no_candidate_clauses")
        max_e = max_n = max_c = 0.0
        best_clause = None
        for (cl, _ov), (pe, pn, pc) in zip(cand, self._nli_batch([(cl, claim) for cl, _ in cand])):
            max_e = max(max_e, pe)
            max_n = max(max_n, pn)
            if pc > max_c:
                max_c = pc
                best_clause = cl
        would_demote = (max_c >= self.theta_contra) and (max_e < self.theta_entail)
        return _result(available=True, would_demote=would_demote,
                       max_contradiction=round(max_c, 4), max_entailment=round(max_e, 4),
                       best_clause=best_clause, reason="ok")


_DEFAULT: Optional[ShadowNLI] = None


def shadow_check(claim: str, source: str) -> ShadowResult:
    """Module-level convenience: run :meth:`ShadowNLI.check` on a shared,
    lazily-built model. Safe to call regardless of whether ``[nli]`` is
    installed (returns ``available=False`` in that case)."""
    global _DEFAULT
    if _DEFAULT is None:
        _DEFAULT = ShadowNLI()
    return _DEFAULT.check(claim, source)
