app/extraction/citation_validator.py

"""Deterministic citation & source-span validation.

This module enforces the core guarantee: **every claim must map to a literal
source span.** It is pure Python with no network dependency, so it can be unit
tested offline and run as a hard gate on any LLM output.

Validation policy (audit-safe — bias toward dropping, never inventing):

* A cited ``chunk_id`` must exist in the retrieved evidence pool. Fabricated
  citations invalidate the span.
* An evidence span's text must actually appear in its cited chunk (verbatim
  after normalization, or via high token-containment to tolerate OCR drift).
* Every NUMBER, MONETARY AMOUNT, and flagged PROPER NOUN in the finding text
  must be present in the cited evidence. This blocks altered ratings, fabricated
  totals, and entity substitution ("London plane tree" -> "cedar tree").
* Escalation/severity words (catastrophic, deadly, unsafe, ...) may only appear
  if present verbatim in the source.

Anything failing a *critical* check is forced to ``NOT_FOUND`` and dropped.
"""

from __future__ import annotations

import re
import unicodedata
from collections.abc import Iterable, Mapping

from app.extraction.schemas import SupportLevel, SurveyFinding

# Severity / escalation vocabulary that must never be introduced by the model.
# These may ONLY survive validation if present verbatim in the cited evidence.
BANNED_SEVERITY_TERMS: frozenset[str] = frozenset({
    "catastrophic", "deadly", "lethal", "unsafe", "dangerous", "hazardous",
    "severe health", "health consequences", "environmental disaster", "disaster",
    "life-threatening", "life threatening", "toxic", "fatal", "collapse imminent",
})

# Containment thresholds for whole-finding grounding (token overlap with evidence).
_SUPPORTED_CONTAINMENT = 0.85
_PARTIAL_CONTAINMENT = 0.55
# Span-level containment for tolerating OCR/quote drift when a verbatim
# substring match fails.
_SPAN_CONTAINMENT = 0.90

_STOPWORDS: frozenset[str] = frozenset({
    "the", "a", "an", "and", "or", "of", "to", "in", "on", "at", "is", "are",
    "was", "were", "be", "been", "being", "this", "that", "these", "those",
    "with", "for", "as", "by", "from", "it", "its", "has", "have", "had",
    "which", "but", "not", "no", "there", "their", "they", "we", "you", "i",
    "will", "would", "should", "could", "may", "can", "also", "some", "any",
})

_WORD_RE = re.compile(r"[A-Za-z0-9£%.\-/]+")
# Numbers, percentages, money, and ratings like "CR2"/"1"/"£12,500"/"3.5m".
_NUMBER_RE = re.compile(r"£?\d[\d,]*(?:\.\d+)?%?")
# TitleCase proper-noun runs of >=2 words, or a single word with >=2 capitals.
_PROPER_NOUN_RE = re.compile(r"\b(?:[A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\b")


def normalize_for_match(text: str) -> str:
    """Lowercase, fold accents, normalize quotes/dashes, collapse whitespace."""
    if not text:
        return ""
    text = unicodedata.normalize("NFKD", text)
    text = "".join(c for c in text if not unicodedata.combining(c))
    text = (
        text.replace("\u2019", "'").replace("\u2018", "'")
        .replace("\u201c", '"').replace("\u201d", '"')
        .replace("\u2013", "-").replace("\u2014", "-")
        .replace("\u00a0", " ")
    )
    text = text.lower()
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def content_tokens(text: str) -> list[str]:
    """Tokens used for containment scoring (stopwords removed)."""
    norm = normalize_for_match(text)
    return [t for t in _WORD_RE.findall(norm) if t not in _STOPWORDS and len(t) > 1]


def _containment(needle_tokens: Iterable[str], haystack_tokens: set[str]) -> float:
    needle = list(needle_tokens)
    if not needle:
        return 1.0
    hits = sum(1 for t in needle if t in haystack_tokens)
    return hits / len(needle)


def _coerce_pool(pool: Mapping[str, str] | Iterable[object]) -> dict[str, str]:
    """Accept a ``{chunk_id: text}`` map or an iterable of SearchResult-like rows."""
    if isinstance(pool, Mapping):
        return {str(k): str(v) for k, v in pool.items()}
    out: dict[str, str] = {}
    for row in pool:
        cid = getattr(row, "chunk_id", None)
        txt = getattr(row, "text", None)
        if cid is not None and txt is not None:
            out[str(cid)] = str(txt)
    return out


def span_in_chunk(span_text: str, chunk_text: str) -> bool:
    """True when ``span_text`` is grounded in ``chunk_text``.

    Verbatim (normalized) substring match first; falls back to high
    token-containment to tolerate OCR/whitespace/quote drift while still
    rejecting genuinely absent content.
    """
    if not span_text or not chunk_text:
        return False
    nspan = normalize_for_match(span_text)
    nchunk = normalize_for_match(chunk_text)
    if not nspan:
        return False
    if nspan in nchunk:
        return True
    chunk_tok = set(_WORD_RE.findall(nchunk))
    span_tok = _WORD_RE.findall(nspan)
    return _containment(span_tok, chunk_tok) >= _SPAN_CONTAINMENT


def _numbers(text: str) -> set[str]:
    return {m.replace(",", "") for m in _NUMBER_RE.findall(text or "")}


def _proper_nouns(text: str) -> set[str]:
    return {normalize_for_match(m) for m in _PROPER_NOUN_RE.findall(text or "")}


def _banned_terms_present(text: str) -> set[str]:
    norm = normalize_for_match(text)
    return {term for term in BANNED_SEVERITY_TERMS if term in norm}


def validate_finding(
    finding: SurveyFinding,
    pool: Mapping[str, str] | Iterable[object],
) -> tuple[SupportLevel, list[str]]:
    """Validate one finding against the retrieved evidence pool.

    Returns the computed :class:`SupportLevel` and a list of human-readable
    violation reasons (empty when fully supported). Mutates nothing.
    """
    pool_map = _coerce_pool(pool)
    violations: list[str] = []

    # 1. Spans must cite real chunks and actually appear in them.
    valid_spans = []
    for span in finding.evidence:
        chunk_text = pool_map.get(span.chunk_id)
        if chunk_text is None:
            violations.append(f"citation to unknown chunk_id={span.chunk_id!r}")
            continue
        if not span_in_chunk(span.text, chunk_text):
            violations.append(f"span not found in chunk {span.chunk_id!r}: {span.text[:60]!r}")
            continue
        valid_spans.append((span, chunk_text))

    if not valid_spans:
        return SupportLevel.NOT_FOUND, violations or ["no valid supporting spans"]

    # Build the union of grounded evidence text from spans that validated.
    evidence_text = " \n ".join(s.text for s, _ in valid_spans)
    # Also allow grounding against the FULL cited chunk text (the span is a
    # window into it; numbers/entities elsewhere in the same chunk still count).
    evidence_text += " \n " + " \n ".join(ct for _, ct in valid_spans)
    evidence_norm = normalize_for_match(evidence_text)
    evidence_tok = set(_WORD_RE.findall(evidence_norm))
    evidence_numbers = _numbers(evidence_text)
    evidence_nouns = _proper_nouns(evidence_text)

    # 2. CRITICAL: every number in the finding must be in the evidence.
    for num in _numbers(finding.finding):
        if num not in evidence_numbers:
            violations.append(f"unsupported number/amount {num!r} not in evidence")

    # 3. CRITICAL: severity escalation must exist in the source.
    finding_banned = _banned_terms_present(finding.finding)
    evidence_banned = _banned_terms_present(evidence_text)
    for term in finding_banned - evidence_banned:
        violations.append(f"unsupported severity term {term!r} not in evidence")

    # 4. CRITICAL: proper nouns (materials, species, places) must not be substituted.
    for noun in _proper_nouns(finding.finding):
        if noun not in evidence_nouns and noun not in evidence_norm:
            violations.append(f"unsupported/altered entity {noun!r} not in evidence")

    has_critical = bool(violations)

    # 5. Whole-finding token containment (catches paraphrased fabrication).
    containment = _containment(content_tokens(finding.finding), evidence_tok)

    if has_critical:
        return SupportLevel.NOT_FOUND, violations
    if containment >= _SUPPORTED_CONTAINMENT:
        return SupportLevel.SUPPORTED, []
    if containment >= _PARTIAL_CONTAINMENT:
        return SupportLevel.PARTIAL, [f"partial grounding (containment={containment:.2f})"]
    return SupportLevel.NOT_FOUND, [f"insufficient grounding (containment={containment:.2f})"]


def validate_findings(
    findings: list[SurveyFinding],
    pool: Mapping[str, str] | Iterable[object],
    *,
    drop_partial: bool = False,
) -> tuple[list[SurveyFinding], list[str]]:
    """Validate every finding; return (kept_findings, dropped_claim_descriptions).

    ``NOT_FOUND`` findings are always dropped. ``PARTIAL`` findings are kept
    (with their downgraded support level) unless ``drop_partial`` is set, in
    which case only fully ``SUPPORTED`` findings survive — the strictest,
    audit-safe mode.
    """
    kept: list[SurveyFinding] = []
    dropped: list[str] = []
    for f in findings:
        support, reasons = validate_finding(f, pool)
        f.support = support
        if support == SupportLevel.NOT_FOUND:
            dropped.append(f"{f.element}: {'; '.join(reasons)}")
            continue
        if support == SupportLevel.PARTIAL and drop_partial:
            dropped.append(f"{f.element}: partial grounding dropped under strict mode")
            continue
        kept.append(f)
    return kept, dropped