"""Verify the citations in one logged RAG trace and classify each fault (PB-RAG-06-04).

Usage:
    python verify_citations.py trace.json [--current current_chunks.json]

trace.json uses the answer schema from PB-RAG-06-01:
    {"trace_id": "...",
     "retrieved": [{"chunk_id": "...", "content": "...", "source_url": "...", "page_number": 3}],
     "answer": {"answerable": true,
                "claims": [{"text": "...", "citations": [{"chunk_id": "...", "quote": "..."}]}]},
     "rendered": [{"chunk_id": "...", "source_url": "...", "page_number": 3}]}   # optional

"rendered" is the citation list from the renderer payload (build_renderer_contract
in answer_with_citations.py). The old field name "page" is still accepted.

current_chunks.json (optional) maps chunk_id to its content in the index today,
or to null if the key no longer exists. Build it with the Documents - Get call in step 6.

Exit code: 0 if no faults, 1 if any fault, 2 if the input is unreadable.
"""
import argparse
import json
import re
import sys

# Fault class -> where the fix belongs (section 5 of the playbook).
GENERATION = "generation"   # the model cited or quoted something it was not given
RENDERER = "renderer"       # the answer is right, the link or page shown to the user is wrong
STALE = "stale_index"       # the chunk changed or disappeared after the answer was generated


def normalise(text):
    """Compare quotes on words only: ignore case, curly quotes and runs of white space."""
    text = text.replace("’", "'").replace("‘", "'").replace("“", '"').replace("”", '"')
    return re.sub(r"\s+", " ", text).strip().lower()


def page_of(item):
    """Page number of a chunk or rendered citation; accepts page_number or the older page."""
    return item.get("page_number", item.get("page"))


def verify(trace, current=None):
    findings = []
    retrieved = {c["chunk_id"]: c for c in trace.get("retrieved", [])}
    claims = trace.get("answer", {}).get("claims", [])

    for i, claim in enumerate(claims):
        if not claim.get("citations"):
            findings.append((GENERATION, f"claim {i}: no citation"))
        for cit in claim.get("citations", []):
            chunk = retrieved.get(cit["chunk_id"])
            if chunk is None:
                findings.append((GENERATION, f"claim {i}: {cit['chunk_id']} was never retrieved (invented id)"))
            elif cit.get("quote") and normalise(cit["quote"]) not in normalise(chunk["content"]):
                findings.append((GENERATION, f"claim {i}: quote not found in {cit['chunk_id']}"))

    for shown in trace.get("rendered", []):
        chunk = retrieved.get(shown["chunk_id"])
        if chunk is None:
            continue
        if shown.get("source_url") != chunk.get("source_url"):
            findings.append((RENDERER, f"{shown['chunk_id']}: rendered source_url {shown.get('source_url')!r} "
                                       f"but chunk has {chunk.get('source_url')!r}"))
        if page_of(shown) != page_of(chunk):
            findings.append((RENDERER, f"{shown['chunk_id']}: rendered page {page_of(shown)!r} "
                                       f"but chunk has {page_of(chunk)!r}"))

    for chunk_id, chunk in retrieved.items():
        if current is None or chunk_id not in current:
            continue
        if current[chunk_id] is None:
            findings.append((STALE, f"{chunk_id}: no longer in the index (deleted or re-chunked)"))
        elif normalise(current[chunk_id]) != normalise(chunk["content"]):
            findings.append((STALE, f"{chunk_id}: content in the index has changed since the answer"))

    return findings


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("trace")
    parser.add_argument("--current", help="JSON map of chunk_id to current indexed content")
    args = parser.parse_args()
    try:
        with open(args.trace, encoding="utf-8") as f:
            trace = json.load(f)
        current = None
        if args.current:
            with open(args.current, encoding="utf-8") as f:
                current = json.load(f)
    except (OSError, json.JSONDecodeError) as exc:
        print(f"cannot read input: {exc}")
        return 2

    findings = verify(trace, current)
    print(f"trace {trace.get('trace_id', '?')}: {len(findings)} fault(s)")
    for fault, detail in findings:
        print(f"  [{fault}] {detail}")
    if not findings:
        print("  all citations verified")
    return 1 if findings else 0


if __name__ == "__main__":
    sys.exit(main())
