"""
Solution: the pieces of a RAG answer you can check (Lab 04, 01_basic_rag.py).
"""
import re

IDK = "I don't have enough information"


def build_prompt(query: str, chunks: list[dict]) -> str:
    """The lab's build_prompt: each chunk as [source] then text, separated by ---."""
    blocks = [f"[{c['source']}]\n{c['text']}" for c in chunks]
    return "Context:\n\n" + "\n\n---\n\n".join(blocks) + f"\n\nQuestion: {query}"


def cited(answer: str) -> list[str]:
    """Sources the answer cites, as [name] tags, in order, without repeats."""
    return list(dict.fromkeys(re.findall(r"\[([a-z0-9_-]+)\]", answer)))


def citations_check(answer: str, sent: list[dict]) -> dict:
    """Which cited sources were actually sent, and whether the answer cites anything at all."""
    sources = {c["source"] for c in sent}
    tags = cited(answer)
    return {"cites": bool(tags), "not_sent": [t for t in tags if t not in sources]}


def grade(answer: str, rule: dict) -> str:
    """Label an answer against a written rule:
    idk      starts with the "not enough information" sentence
    wrong    starts with any of rule["wrong_if_starts"], or contains none of rule["needs_any"]
    partial  has one of needs_any but is missing every phrase in rule["complete_if_any"]
    right    otherwise
    """
    t = answer.strip()
    if t.startswith(IDK):
        return "idk"
    if any(t.startswith(p) for p in rule.get("wrong_if_starts", [])):
        return "wrong"
    needs = rule.get("needs_any", [])
    if needs and not any(n.lower() in t.lower() for n in needs):
        return "wrong"
    complete = rule.get("complete_if_any", [])
    if complete and not any(c.lower() in t.lower() for c in complete):
        return "partial"
    return "right"


def tally(labels: list[str]) -> dict:
    out = {"right": 0, "partial": 0, "wrong": 0, "idk": 0}
    for x in labels:
        out[x] += 1
    return out