#!/usr/bin/env python3
"""RAG eval harness — retrieval and generation metrics.

Computes recall@k, precision@k, MRR for retrieval, and (mock) faithfulness for
generation. Runs a golden eval set and prints a scored table.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass
class EvalSample:
    query: str
    relevant_ids: list[int]
    expected_answer: str = ""
    category: str = "easy"  # easy | paraphrase | cross-doc | out-of-scope


# ── Metrics ─────────────────────────────────────────────────────────────

def recall_at_k(retrieved: list[int], relevant: list[int], k: int) -> float:
    if not relevant:
        return 1.0
    return len(set(retrieved[:k]) & set(relevant)) / len(relevant)


def precision_at_k(retrieved: list[int], relevant: list[int], k: int) -> float:
    return len(set(retrieved[:k]) & set(relevant)) / max(k, 1)


def mrr(retrieved: list[int], relevant: list[int]) -> float:
    for rank, idx in enumerate(retrieved, 1):
        if idx in relevant:
            return 1.0 / rank
    return 0.0


def nDCG_at_k(retrieved: list[int], relevant: list[int], k: int) -> float:
    """Normalized Discounted Cumulative Gain."""
    ideal = [1.0] * min(len(relevant), k)
    dcg = sum(
        (1.0 if idx in relevant else 0.0) / __import__("math").log2(i + 2)
        for i, idx in enumerate(retrieved[:k])
    )
    idcg = sum(rel / __import__("math").log2(i + 2) for i, rel in enumerate(ideal))
    return dcg / idcg if idcg > 0 else 0.0


# ── Mock retrievers ─────────────────────────────────────────────────────

def mock_retrieve(query: str, corpus: list[str], top_k: int = 5) -> list[int]:
    """Deterministic mock: rank docs by word overlap with the query."""
    query_terms = set(query.lower().split())
    scored = []
    for i, doc in enumerate(corpus):
        overlap = len(query_terms & set(doc.lower().split()))
        if overlap > 0:
            scored.append((i, overlap))
    scored.sort(key=lambda x: (-x[1], x[0]))
    return [idx for idx, _ in scored[:top_k]]


def mock_generate(query: str, retrieved_docs: list[str]) -> str:
    """Mock generation: naive answer from retrieved docs."""
    if not retrieved_docs:
        return "I don't know."
    return f"Based on the documents: {retrieved_docs[0]}"


# ── Harness ─────────────────────────────────────────────────────────────

def run_harness(
    corpus: list[str],
    eval_set: list[EvalSample],
    retrieve_fn,
    k_values: list[int],
) -> dict[str, Any]:
    results = {f"recall@{k}": [] for k in k_values}
    results.update({f"precision@{k}": [] for k in k_values})
    results["mrr"] = []
    results["ndcg@3"] = []
    results["answered_idk"] = []  # For out-of-scope samples

    for sample in eval_set:
        retrieved = retrieve_fn(sample.query, corpus, top_k=max(k_values))
        for k in k_values:
            results[f"recall@{k}"].append(recall_at_k(retrieved, sample.relevant_ids, k))
            results[f"precision@{k}"].append(precision_at_k(retrieved, sample.relevant_ids, k))
        results["mrr"].append(mrr(retrieved, sample.relevant_ids))
        results["ndcg@3"].append(nDCG_at_k(retrieved, sample.relevant_ids, 3))

        # Out-of-scope handling: should say "don't know" when no relevant docs
        if sample.category == "out-of-scope":
            answer = mock_generate(sample.query, [corpus[i] for i in retrieved])
            results["answered_idk"].append("don't know" in answer.lower())

    # Aggregate
    summary = {}
    for metric, values in results.items():
        if values:
            summary[metric] = round(sum(values) / len(values), 3)
    return summary


# ── Demo ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    corpus = [
        "Python is great for data science and machine learning.",
        "JavaScript powers interactive web applications.",
        "PostgreSQL is a robust relational database system.",
        "Docker containers package applications with dependencies.",
        "Machine learning models require evaluation and testing.",
        "FastAPI is a modern Python web framework for APIs.",
    ]

    eval_set = [
        EvalSample("Python web framework", [5], category="easy"),
        EvalSample("data science tools", [0, 4], category="paraphrase"),
        EvalSample("containerized database", [2, 3], category="cross-doc"),
        EvalSample("quantum computing", [], category="out-of-scope"),
    ]

    print("RAG Eval Harness Results\n" + "=" * 40)
    summary = run_harness(corpus, eval_set, mock_retrieve, k_values=[1, 3, 5])

    print("\nMetric summary:")
    for metric, value in summary.items():
        print(f"  {metric}: {value}")

    print("\nPer-sample detail:")
    for sample in eval_set:
        retrieved = mock_retrieve(sample.query, corpus, top_k=3)
        print(f"  '{sample.query}' ({sample.category}) → docs {retrieved}")

    print("\n✓ Harness complete. Re-run after every retrieval/chunking change.")