"""Evaluation harness for RAG pipelines."""

from __future__ import annotations

from dataclasses import dataclass

from bifrost_rag.store import Document


@dataclass
class EvaluationResult:
    """Aggregate evaluation metrics for a RAG system."""

    precision: float
    recall: float
    f1: float
    mrr: float
    total_queries: int

    def summary(self) -> str:
        return (
            f"RAG Evaluation ({self.total_queries} queries)\n"
            f"  Precision@K: {self.precision:.3f}\n"
            f"  Recall@K:    {self.recall:.3f}\n"
            f"  F1:          {self.f1:.3f}\n"
            f"  MRR:         {self.mrr:.3f}"
        )


class RAGEvaluator:
    """Evaluation harness for measuring RAG retrieval quality."""

    def precision_at_k(
        self,
        retrieved: list[Document],
        relevant_ids: set[str],
        k: int,
    ) -> float:
        """Compute precision at K."""
        top_k = retrieved[:k]
        if not top_k:
            return 0.0
        relevant_count = sum(1 for doc in top_k if doc.id in relevant_ids)
        return relevant_count / len(top_k)

    def recall_at_k(
        self,
        retrieved: list[Document],
        relevant_ids: set[str],
        k: int,
    ) -> float:
        """Compute recall at K."""
        if not relevant_ids:
            return 0.0
        top_k = retrieved[:k]
        relevant_count = sum(1 for doc in top_k if doc.id in relevant_ids)
        return relevant_count / len(relevant_ids)

    def mrr(
        self,
        retrieved: list[Document],
        relevant_ids: set[str],
    ) -> float:
        """Compute Mean Reciprocal Rank."""
        for i, doc in enumerate(retrieved):
            if doc.id in relevant_ids:
                return 1.0 / (i + 1)
        return 0.0

    def evaluate(
        self,
        dataset: list[dict],
        k: int = 5,
    ) -> EvaluationResult:
        """Evaluate over a dataset of queries.

        Each dataset entry should have:
            - "retrieved": list[Document]
            - "relevant_ids": set[str]
        """
        if not dataset:
            return EvaluationResult(
                precision=0.0, recall=0.0, f1=0.0, mrr=0.0, total_queries=0
            )

        precisions = []
        recalls = []
        mrrs = []

        for entry in dataset:
            retrieved = entry["retrieved"]
            relevant_ids = entry["relevant_ids"]

            p = self.precision_at_k(retrieved, relevant_ids, k)
            r = self.recall_at_k(retrieved, relevant_ids, k)
            m = self.mrr(retrieved, relevant_ids)

            precisions.append(p)
            recalls.append(r)
            mrrs.append(m)

        avg_p = sum(precisions) / len(precisions)
        avg_r = sum(recalls) / len(recalls)
        avg_mrr = sum(mrrs) / len(mrrs)

        f1 = (2 * avg_p * avg_r / (avg_p + avg_r)) if (avg_p + avg_r) > 0 else 0.0

        return EvaluationResult(
            precision=avg_p,
            recall=avg_r,
            f1=f1,
            mrr=avg_mrr,
            total_queries=len(dataset),
        )
