#!/usr/bin/env python3
"""RAG hallucination mitigation: faithfulness check via claim-level verification.

Splits an LLM-generated answer into claims and checks each against retrieved
contexts. Flags claims not supported by any context. This is a lightweight
simulation of the LLM-as-judge faithfulness check.
"""
from __future__ import annotations

import re
import sys


def split_claims(text: str) -> list[str]:
    """Naive claim splitter: break on sentence boundaries."""
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    return [s.strip() for s in sentences if len(s.strip()) > 5]


def claim_supported(claim: str, contexts: list[str], min_overlap: int = 3) -> bool:
    """
    Check if a claim has token overlap with any context chunk.
    This is a SIMPLE heuristic -- in production, use an NLI model or LLM judge.
    """
    claim_tokens = set(claim.lower().split())
    for ctx in contexts:
        ctx_tokens = set(ctx.lower().split())
        overlap = len(claim_tokens & ctx_tokens)
        if overlap >= min_overlap:
            return True
    return False


def faithfulness_score(answer: str, contexts: list[str]) -> tuple[float, list[dict]]:
    """
    Return (score, details) where score = supported_claims / total_claims.
    details lists each claim with support status.
    """
    claims = split_claims(answer)
    if not claims:
        return 1.0, []

    results = []
    supported = 0
    for claim in claims:
        ok = claim_supported(claim, contexts)
        if ok:
            supported += 1
        results.append({"claim": claim, "supported": ok})

    return supported / len(claims), results


def main() -> None:
    # Simulated RAG output
    test_cases = [
        {
            "question": "What is RAG?",
            "answer": "RAG is retrieval-augmented generation. "
                      "It combines search with LLMs. "
                      "RAG was invented in 2020 by Meta AI. "
                      "It always produces perfect answers.",
            "contexts": [
                "RAG stands for Retrieval-Augmented Generation.",
                "RAG combines information retrieval with large language models.",
                "This technique grounds LLM outputs in retrieved documents.",
            ],
        },
        {
            "question": "Who wrote Hamlet?",
            "answer": "Hamlet was written by William Shakespeare.",
            "contexts": [
                "William Shakespeare authored Hamlet around 1600.",
                "Hamlet is one of Shakespeare's most famous tragedies.",
            ],
        },
    ]

    for i, tc in enumerate(test_cases, 1):
        score, details = faithfulness_score(tc["answer"], tc["contexts"])
        print(f"\n=== Test {i}: {tc['question']} ===")
        print(f"Faithfulness: {score:.2f} ({int(score * 100)}%)")
        for item in details:
            mark = "+" if item["supported"] else "-"
            print(f"  [{mark}] {item['claim'][:80]}")
        if score < 0.8:
            print("  WARNING: Low faithfulness -- claims not backed by context")

    # Threshold for CI
    for tc in test_cases:
        score, _ = faithfulness_score(tc["answer"], tc["contexts"])
        if score < 0.7:
            print(f"\nFAIL: faithfulness {score:.2f} below threshold for: {tc['question'][:50]}")
            sys.exit(1)

    print("\nAll faithfulness checks passed.")


if __name__ == "__main__":
    main()