"""Citation validation: strip any [bracketed] citation not in the retrieved set.

Guards against hallucinated case names. A citation is kept only if its text
matches (substring, case-insensitive) one of the retrieved source case names.
"""
from __future__ import annotations

import re

_CITATION_RE = re.compile(r"\[([^\[\]]+)\]")


def _matches_any(cited: str, allowed: list[str]) -> bool:
    c = cited.strip().lower()
    if not c:
        return False
    for name in allowed:
        n = name.lower()
        # Accept exact, or either being a substring of the other (handles the
        # model truncating/expanding the case name slightly).
        if c == n or c in n or n in c:
            return True
    return False


def validate_answer(answer: str, sources: list[dict]) -> dict:
    """Drop hallucinated citations. Return {answer, removed:[...]}.

    Kept citations stay bracketed. Unknown ones are unwrapped (text retained,
    brackets removed) and reported in `removed`.
    """
    allowed = [s["case_name"] for s in sources]
    removed: list[str] = []

    def _sub(m: re.Match) -> str:
        cited = m.group(1)
        if _matches_any(cited, allowed):
            return m.group(0)  # keep bracketed citation
        removed.append(cited)
        return cited  # strip brackets, keep raw text

    cleaned = _CITATION_RE.sub(_sub, answer)
    return {"answer": cleaned, "removed": removed}


if __name__ == "__main__":
    srcs = [{"case_name": "Jimmy Jahangir Madan vs Bolly Cariyappa Hindley By L.Rs on 10 August, 2001"}]
    raw = (
        "The application under Section 302 was held not maintainable "
        "[Jimmy Jahangir Madan vs Bolly Cariyappa Hindley By L.Rs on 10 August, 2001]. "
        "A different rule applied [Fake Hallucinated Case v Nobody 2099]."
    )
    out = validate_answer(raw, srcs)
    print("CLEANED:\n", out["answer"])
    print("\nREMOVED:", out["removed"])