"""Checks a generated answer against the evidence it was built from.

The answer layer may only rephrase retrieved text. This module is what makes
that a checkable property rather than a promise in a prompt: it recomputes,
from the evidence alone, whether every claim in a generated answer traces
back to the specific source block it cites. A generation that fails is
discarded, never shown.

Binding is per-citation, not global. The answer is split at each citation
marker group (one or more consecutive `[n]`); the text immediately before a
group is that group's claim, and only the evidence block(s) named in that
group may support it. A number that is true of evidence block 2 does not
make a claim citing `[1]` grounded — the old implementation pooled every
number from every evidence block into one set, which let a number attributed
to the wrong source pass silently. `evidence_texts` is positional: `[n]`
refers to `evidence_texts[n - 1]`.

A claim with no valid citation group is rejected outright — a citation
nobody can follow is not a citation, and an uncited clinical statement is not
verifiable, numeric or not. This catches a missing-citation defect that the
old check never looked for at all (it only ever checked numbers already
carrying a marker).

Numbers are compared **character for character**, deliberately. "7,5" and
"7.5" are not treated as equal, and no attempt is made to parse either into a
quantity. Parsing invites the one error that matters most here: `1.500` is
1500 under one reading and 1.5 under another, and a normaliser that strips
separators maps "7,5" and "75" to the same key — a tenfold dose error scored
as a match. The model is told to copy figures verbatim, so an exact match is
achievable, and every deviation from it is refused rather than interpreted.

What this module still cannot do: confirm that a citation-bearing nonnumeric
claim is actually *entailed* by the block it cites (e.g. "chữa ung thư [1]"
where evidence 1 is only about "điều trị đái tháo đường" — same drug name,
unrelated indication). Regex-level number/citation checking has no notion of
semantic content. That gap is closed separately by an LLM entailment pass
(`rag/answer.py`'s post-generation verifier call), not by this module.
"""
from __future__ import annotations

import re
from dataclasses import dataclass

# A digit run with internal separators kept: "500", "7,5", "1.000".
# Ranges ("4 - 6 giờ") yield two tokens, and each is checked on its own.
_NUMBER = re.compile(r"\d+(?:[.,]\d+)*")

# Citation markers are stripped before number extraction so that "[2]" is
# never mistaken for the quantity 2.
_CITATION = re.compile(r"\[(\d+)\]")

# One or more consecutive markers ("[1]", "[1][2]") count as a single group:
# the prompt allows citing more than one source for one claim, and each is
# checked against the union of just those sources, not all evidence.
_CITATION_GROUP = re.compile(r"(?:\[\d+\])+")

# Any word character — letter (Vietnamese diacritics included) or digit —
# used to tell "claim with actual content" apart from bare punctuation or
# whitespace trailing a citation, which needs no citation of its own.
_LETTER = re.compile(r"\w", re.UNICODE)


@dataclass(frozen=True)
class GroundingReport:
    grounded: bool
    unsupported_numbers: tuple[str, ...]
    invalid_citations: tuple[int, ...]
    uncited_claim: bool
    cited_indices: tuple[int, ...]

    @property
    def reason(self) -> str:
        if self.unsupported_numbers:
            return "ungrounded_number"
        if self.invalid_citations:
            return "invalid_citation"
        if self.uncited_claim:
            return "uncited_claim"
        return "grounded"


def numbers_in(text: str) -> tuple[str, ...]:
    """Numeric tokens, with citation markers removed first."""
    return tuple(_NUMBER.findall(_CITATION.sub(" ", text)))


def citations_in(text: str) -> tuple[int, ...]:
    return tuple(int(marker) for marker in _CITATION.findall(text))


def has_content(text: str) -> bool:
    """True once `text` carries any letter or digit — i.e. more than
    punctuation or whitespace left over between/after citation markers."""
    return _LETTER.search(text) is not None


@dataclass(frozen=True)
class Claim:
    """One citation-bounded segment of an answer: the text before a citation
    group, and the (in-range) evidence indices that group names.

    `indices` is empty for the trailing segment after the last citation
    group, or for a claim whose only marker(s) were out of range — in both
    cases there is no evidence block left to check the claim against.
    """

    text: str
    indices: tuple[int, ...]


def split_claims(answer: str, evidence_count: int) -> tuple[Claim, ...]:
    """The claim segmentation `verify` checks numbers against, exposed so a
    semantic entailment pass can run the same per-claim binding — each claim
    checked only against the evidence block(s) it actually cites, never the
    whole evidence set.
    """
    claims: list[Claim] = []
    cursor = 0
    for group in _CITATION_GROUP.finditer(answer):
        text = answer[cursor:group.start()]
        cursor = group.end()
        indices = tuple(
            i for i in citations_in(group.group(0)) if 1 <= i <= evidence_count
        )
        claims.append(Claim(text, indices))
    claims.append(Claim(answer[cursor:], ()))
    return tuple(claims)


def verify(answer: str, evidence_texts: tuple[str, ...]) -> GroundingReport:
    """Whether `answer` states only figures and sources traceable to the
    specific evidence block(s) cited immediately after each claim.

    See module docstring for the binding rule and its known limit (no
    semantic entailment check).
    """
    unsupported: list[str] = []
    invalid: list[int] = []
    uncited = False
    cited_all: set[int] = set()

    cursor = 0
    for group in _CITATION_GROUP.finditer(answer):
        claim = answer[cursor:group.start()]
        cursor = group.end()

        indices = citations_in(group.group(0))
        bad = [i for i in indices if not 1 <= i <= len(evidence_texts)]
        good = [i for i in indices if i not in bad]
        invalid.extend(bad)
        cited_all.update(good)

        claim_numbers = numbers_in(claim)
        if good:
            source_numbers: set[str] = set()
            for index in good:
                source_numbers.update(numbers_in(evidence_texts[index - 1]))
            unsupported.extend(n for n in claim_numbers if n not in source_numbers)
        else:
            # Every marker in this group was out of range: nothing to bind
            # the claim to, numeric or not.
            unsupported.extend(claim_numbers)
            if has_content(claim):
                uncited = True

    tail = answer[cursor:]
    unsupported.extend(numbers_in(tail))
    if has_content(tail):
        uncited = True

    return GroundingReport(
        grounded=not unsupported and not invalid and not uncited,
        unsupported_numbers=tuple(unsupported),
        invalid_citations=tuple(invalid),
        uncited_claim=uncited,
        cited_indices=tuple(sorted(cited_all)),
    )