"""Tests for the RAG evaluation harness."""

import pytest

from bifrost_rag.evaluation import EvaluationResult, RAGEvaluator
from bifrost_rag.store import Document


class TestEvaluationResult:
    def test_creation(self):
        result = EvaluationResult(
            precision=0.8,
            recall=0.6,
            f1=0.685,
            mrr=0.75,
            total_queries=10,
        )
        assert result.precision == 0.8
        assert result.recall == 0.6
        assert result.f1 == 0.685
        assert result.mrr == 0.75

    def test_summary(self):
        result = EvaluationResult(
            precision=0.8, recall=0.6, f1=0.685, mrr=0.75, total_queries=10
        )
        summary = result.summary()
        assert "precision" in summary.lower()
        assert "recall" in summary.lower()


class TestRAGEvaluator:
    def test_creation(self):
        evaluator = RAGEvaluator()
        assert evaluator is not None

    def test_precision_at_k(self):
        evaluator = RAGEvaluator()
        retrieved = [
            Document(id="doc1", text="relevant"),
            Document(id="doc2", text="not relevant"),
            Document(id="doc3", text="relevant"),
        ]
        relevant_ids = {"doc1", "doc3", "doc5"}
        precision = evaluator.precision_at_k(retrieved, relevant_ids, k=3)
        assert precision == pytest.approx(2 / 3)

    def test_precision_at_k_all_relevant(self):
        evaluator = RAGEvaluator()
        retrieved = [
            Document(id="doc1", text="a"),
            Document(id="doc2", text="b"),
        ]
        relevant_ids = {"doc1", "doc2"}
        precision = evaluator.precision_at_k(retrieved, relevant_ids, k=2)
        assert precision == 1.0

    def test_precision_at_k_none_relevant(self):
        evaluator = RAGEvaluator()
        retrieved = [
            Document(id="doc1", text="a"),
            Document(id="doc2", text="b"),
        ]
        relevant_ids = {"doc3", "doc4"}
        precision = evaluator.precision_at_k(retrieved, relevant_ids, k=2)
        assert precision == 0.0

    def test_recall_at_k(self):
        evaluator = RAGEvaluator()
        retrieved = [
            Document(id="doc1", text="a"),
            Document(id="doc2", text="b"),
        ]
        relevant_ids = {"doc1", "doc2", "doc3"}
        recall = evaluator.recall_at_k(retrieved, relevant_ids, k=2)
        assert recall == pytest.approx(2 / 3)

    def test_mrr(self):
        evaluator = RAGEvaluator()
        # First relevant at position 2 (0-indexed position 1)
        retrieved = [
            Document(id="doc_bad", text="a"),
            Document(id="doc_good", text="b"),
            Document(id="doc_also_good", text="c"),
        ]
        relevant_ids = {"doc_good", "doc_also_good"}
        mrr = evaluator.mrr(retrieved, relevant_ids)
        assert mrr == pytest.approx(0.5)  # 1/(1+1) = 1/2

    def test_mrr_first_is_relevant(self):
        evaluator = RAGEvaluator()
        retrieved = [Document(id="doc1", text="a")]
        relevant_ids = {"doc1"}
        mrr = evaluator.mrr(retrieved, relevant_ids)
        assert mrr == 1.0

    def test_mrr_no_relevant(self):
        evaluator = RAGEvaluator()
        retrieved = [Document(id="doc1", text="a")]
        relevant_ids = {"doc_other"}
        mrr = evaluator.mrr(retrieved, relevant_ids)
        assert mrr == 0.0

    def test_evaluate_dataset(self):
        evaluator = RAGEvaluator()
        dataset = [
            {
                "retrieved": [
                    Document(id="doc1", text="a"),
                    Document(id="doc2", text="b"),
                ],
                "relevant_ids": {"doc1"},
            },
            {
                "retrieved": [
                    Document(id="doc3", text="c"),
                    Document(id="doc4", text="d"),
                ],
                "relevant_ids": {"doc3", "doc4"},
            },
        ]
        result = evaluator.evaluate(dataset, k=2)
        assert isinstance(result, EvaluationResult)
        assert result.total_queries == 2
        assert 0.0 <= result.precision <= 1.0
        assert 0.0 <= result.recall <= 1.0
