"""
evaluate.py — the evaluation harness (Knowledge Module 08 §2-6).

Metrics computed over a golden set of (question, contexts, answer, ...):
  - retrieval : recall@k, context precision
  - generation: groundedness/faithfulness (LLM-judge stand-in), answer relevance
  - safety    : refusal-correctness (did it refuse what it should?)
The harness + thresholds are exactly what you'd gate a release on; only the
JUDGE swaps for a real LLM-judge / RAGAS / Foundry evaluator in production.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field

_WORD = re.compile(r"[A-Za-z0-9.%]+")
_CITATION = re.compile(r"\[[^\]]*\]")     # strip [source:section] before scoring
_STOP = {"the", "a", "an", "is", "are", "of", "to", "for", "and", "or", "in",
         "on", "your", "you", "i", "it", "with", "can", "be", "this", "that",
         "have", "has", "from", "by", "at", "may", "within", "do", "does"}


def _content(text: str) -> set[str]:
    text = _CITATION.sub(" ", text)
    return {w.lower() for w in _WORD.findall(text) if w.lower() not in _STOP}


@dataclass
class GoldenCase:
    question: str
    contexts: list[str]          # retrieved chunks shown to the model
    answer: str                  # the model's answer (to be judged)
    expected_source_idx: int | None = None   # which context SHOULD be retrieved (recall)
    retrieved_idx: list[int] = field(default_factory=list)  # which were retrieved
    should_refuse: bool = False
    refused: bool = False
    language: str = "en"         # for stratified evaluation (en / ar)


# --- retrieval metrics ------------------------------------------------------
def recall_at_k(case: GoldenCase, k: int = 5) -> float:
    if case.expected_source_idx is None:
        return 1.0
    return 1.0 if case.expected_source_idx in case.retrieved_idx[:k] else 0.0


# --- generation metrics -----------------------------------------------------
def groundedness(case: GoldenCase) -> float:
    """LLM-JUDGE STAND-IN. A real judge reads CONTEXT + ANSWER and rates whether
    every claim is supported. We approximate 'supported' as: the answer's content
    words are (almost entirely) present in the union of the contexts. An answer
    that introduces words not in any context is (likely) hallucinated -> low
    score. This deterministically discriminates grounded vs invented answers,
    which is the whole point of the metric (K08 §3)."""
    if case.refused:
        return 1.0                      # abstaining is grounded behavior
    ctx = set().union(*[_content(c) for c in case.contexts]) if case.contexts else set()
    ans = _content(case.answer)
    if not ans:
        return 1.0
    supported = len(ans & ctx) / len(ans)
    return round(supported, 3)


def answer_relevance(case: GoldenCase) -> float:
    """Does the answer address the question? Approximated by content overlap of
    question terms with the answer (a real judge is more nuanced)."""
    if case.refused:
        return 1.0 if case.should_refuse else 0.0
    q, a = _content(case.question), _content(case.answer)
    if not q:
        return 1.0
    return round(len(q & a) / len(q), 3)


# --- safety metric ----------------------------------------------------------
def refusal_correct(case: GoldenCase) -> float:
    """Did the system refuse exactly when it should? (refused == should_refuse)"""
    return 1.0 if case.refused == case.should_refuse else 0.0


# --- aggregate report -------------------------------------------------------
def evaluate(cases: list[GoldenCase], k: int = 5) -> dict:
    n = len(cases)
    return {
        "n": n,
        "recall@k": round(sum(recall_at_k(c, k) for c in cases) / n, 3),
        "groundedness": round(sum(groundedness(c) for c in cases) / n, 3),
        "answer_relevance": round(sum(answer_relevance(c) for c in cases) / n, 3),
        "refusal_correct": round(sum(refusal_correct(c) for c in cases) / n, 3),
    }


def evaluate_stratified(cases: list[GoldenCase], by: str, k: int = 5) -> dict:
    """Stratify metrics by a grouping (e.g. language) so an Arabic regression
    can't hide in the average (K08 §5). `by` reads an attribute set on each case."""
    groups: dict[str, list[GoldenCase]] = {}
    for c in cases:
        groups.setdefault(getattr(c, by, "?"), []).append(c)
    return {g: evaluate(cs, k) for g, cs in groups.items()}