"""Cross-encoder aboutness / relevance shadow check — #000052 §3.2.

For a (query, document) pair returns a single relevance logit (higher
= more relevant). The shadow check has two surface methods that wrap
that single primitive:

  - :meth:`check_question_answer(question, answer)` — does the answer
    address the question? (catches deflection / topic-drift the
    lexical ``diagnose_deflection`` sidecar misses on token-collision
    questions.)
  - :meth:`check_claim_source(claim, source_title_or_clause)` — is
    each claim *about* its cited source? (catches the 'Zionist
    entity'-style mis-cite where the cited title shares a polysemous
    token with the claim but is structurally about a different topic.)

Both are SHADOW: the result is telemetry, never an ``audit_mode``
input. **No demote threshold is hardcoded** — the bench-maxing
discipline from #000049 §7 #18→#27 (proved six times that clean-eval
thresholds do not transfer to real bench-qa precision *or* recall)
requires the threshold to be set by a shadow sweep against pooled
real STRICT, *not* by a literature number. The model runtime is the
optional ``[nli]`` extra (same transformers + torch dependency the
NLI shadow uses); everything degrades to ``available=False`` when it
is missing.
"""
from __future__ import annotations

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

_MANIFEST_PATH = Path(__file__).resolve().parent / "manifest.json"

# Claim-lattice overlay markup that depresses cross-encoder scores when
# fed verbatim (the reranker tries to "score" hash strings + bracket
# punctuation). Stripping it before scoring was measured to drop the
# real-STRICT FP rate by 22% relative on the §3.2.2 step 2 pool
# without losing any true-positive deflection signal (the broad-
# question / narrow-answer cases still scored low). #000052 §3.2.2.
_CLAIM_LATTICE_BLOCK_RE = re.compile(r"\[E\d+\s*\|[^\]]*\]", re.DOTALL)
_TRAILING_BRACKET_TAIL_RE = re.compile(r"\.{2,}[\"'’”]?\s*[\]>]?\s*$")


def clean_for_relevance(text: str) -> str:
    """Strip claim-lattice overlay markup ([E\\d+ | title | hash: …])
    and truncated trailing-bracket tails (..."]) before relevance
    scoring. These artifacts are verifier overlay, not part of the
    natural-language answer; scoring them hurts precision. Pure
    function; safe on already-clean inputs (no-op)."""
    if not text:
        return text
    out = _CLAIM_LATTICE_BLOCK_RE.sub("", text)
    out = _TRAILING_BRACKET_TAIL_RE.sub("", out)
    return " ".join(out.split())


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


@dataclass
class ShadowRelevanceResult:
    available: bool                 # False iff the [nli] extra / model could not load
    relevance_score: float          # the cross-encoder logit (higher = more relevant); NOT a probability
    would_demote: bool              # only meaningful when a demote_below_score is configured (else always False)
    demote_below_score: Optional[float]
    model_version: Optional[str]
    pair_kind: str                  # "question_answer" | "claim_source"
    reason: str = ""

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


def _resolve_relevance_index(id2label: dict) -> int:
    """Cross-encoder reranker checkpoints have a single output, usually
    `LABEL_0` or `relevant`. Return that index; if there are multiple
    labels, return the one whose name suggests positive relevance.
    Substring traps: 'irrelevant' contains 'relevant' — match on tokens,
    not raw substrings."""
    if not id2label:
        return 0
    if len(id2label) == 1:
        return int(next(iter(id2label.keys())))
    # tokenize each label, look for a positive-relevance token without a
    # negative-relevance modifier
    def _toks(s: str) -> set[str]:
        return set(re.findall(r"[a-z0-9]+", str(s).lower()))
    pos_tokens = {"relevant", "positive", "entail", "entailment", "label_1"}
    neg_tokens = {"irrelevant", "negative", "not", "contradiction", "label_0"}
    # first pass: a label whose token set contains a pos token but no neg token
    for k, v in id2label.items():
        ts = _toks(v)
        if ts & pos_tokens and not (ts & neg_tokens):
            return int(k)
    # second pass: "label_1" convention
    for k, v in id2label.items():
        if str(v).lower().strip() in ("label_1", "1"):
            return int(k)
    return 0


class ShadowRelevance:
    """Lazily-loaded pinned cross-encoder reranker + the §3.2 aboutness
    check. Construction never raises; degrades to ``available=False``."""

    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("relevance_model_version")
        self.max_length: int = int(self.manifest.get("max_length", 256))
        self.demote_below_score: Optional[float] = self.manifest.get("demote_below_score")
        self.device_pref: str = device or os.environ.get("ARBORIST_RELEVANCE_DEVICE") or os.environ.get("ARBORIST_NLI_DEVICE") or "auto"
        self.device: Optional[str] = None
        self.batch_size: int = int(os.environ.get("ARBORIST_RELEVANCE_BATCH", "64"))
        self.backend: Optional[str] = None
        self.available = False
        self._reason = "uninitialised"
        self._tok = None
        self._model = None
        self._rel_idx: Optional[int] = None

    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
        if self.device_pref == "auto":
            import torch as _torch
            self.device = "cuda" if _torch.cuda.is_available() else "cpu"
        else:
            self.device = self.device_pref
        try:
            repo = self.manifest["hf_repo"]
            rev = self.manifest.get("pinned_revision")
            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")
            id2label = {int(k): v for k, v in self._model.config.id2label.items()}
            self._rel_idx = _resolve_relevance_index(id2label)
            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
        self.available = True
        self._reason = "ok"

    def _score_batch(self, pairs: list[tuple[str, str]]) -> list[float]:
        import torch
        scores: list[float] = []
        for i in range(0, len(pairs), self.batch_size):
            chunk = pairs[i:i + self.batch_size]
            q = [p for p, _ in chunk]
            d = [doc for _, doc in chunk]
            enc = self._tok(q, d, return_tensors="pt", truncation=True,
                            padding=True, max_length=self.max_length)
            if self.device == "cuda":
                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:
                scores.append(float(row[self._rel_idx] if len(row) > 1 else row[0]))
        return scores

    def _check(self, query: str, document: str, pair_kind: str,
               clean_input: bool = True) -> ShadowRelevanceResult:
        def _result(**kw) -> ShadowRelevanceResult:
            base = dict(model_version=self.model_version, pair_kind=pair_kind,
                        demote_below_score=self.demote_below_score)
            base.update(kw)
            return ShadowRelevanceResult(**base)  # type: ignore[arg-type]

        self._ensure_loaded()
        if not self.available:
            return _result(available=False, relevance_score=0.0, would_demote=False,
                           reason=self._reason)
        # Strip claim-lattice overlay markup by default — measured to
        # drop the real-STRICT FP rate by 22% relative without losing
        # true-positive deflection signal (#000052 §3.2.2 refinement).
        # Opt out with clean_input=False if you want the raw score.
        if clean_input:
            query = clean_for_relevance(query)
            document = clean_for_relevance(document)
        if not (query or "").strip() or not (document or "").strip():
            return _result(available=True, relevance_score=0.0, would_demote=False,
                           reason="empty_input")
        score = self._score_batch([(query, document)])[0]
        wd = (self.demote_below_score is not None and score < self.demote_below_score)
        return _result(available=True, relevance_score=round(score, 4),
                       would_demote=wd, reason="ok")

    def check_question_answer(self, question: str, answer: str) -> ShadowRelevanceResult:
        return self._check(question, answer, "question_answer")

    def check_claim_source(self, claim: str, source: str) -> ShadowRelevanceResult:
        return self._check(claim, source, "claim_source")


_DEFAULT: Optional[ShadowRelevance] = None


def _shared() -> ShadowRelevance:
    global _DEFAULT
    if _DEFAULT is None:
        _DEFAULT = ShadowRelevance()
    return _DEFAULT


def shadow_check_question_answer(question: str, answer: str) -> ShadowRelevanceResult:
    return _shared().check_question_answer(question, answer)


def shadow_check_claim_source(claim: str, source: str) -> ShadowRelevanceResult:
    return _shared().check_claim_source(claim, source)
