"""Tests for the retrieval module."""

from bifrost_rag.retrieval import RetrievalResult, Retriever
from bifrost_rag.store import ChromaStore, Document


class TestRetrievalResult:
    def test_creation(self):
        result = RetrievalResult(
            query="test query",
            documents=[Document(id="doc1", text="result")],
            scores=[0.95],
        )
        assert result.query == "test query"
        assert len(result.documents) == 1
        assert result.scores[0] == 0.95

    def test_empty_result(self):
        result = RetrievalResult(query="no results", documents=[], scores=[])
        assert len(result.documents) == 0

    def test_top_document(self):
        docs = [
            Document(id="doc1", text="first"),
            Document(id="doc2", text="second"),
        ]
        result = RetrievalResult(query="test", documents=docs, scores=[0.8, 0.9])
        assert result.top_document.id == "doc1"

    def test_top_document_empty(self):
        result = RetrievalResult(query="test", documents=[], scores=[])
        assert result.top_document is None


class TestRetriever:
    def test_creation(self):
        store = ChromaStore(collection_name="test_retriever")
        retriever = Retriever(store=store, top_k=5)
        assert retriever.top_k == 5

    def test_retrieve_with_embedding(self):
        store = ChromaStore(collection_name="test_retrieve_embed")
        docs = [
            Document(id="doc1", text="Python guide", embedding=[1.0, 0.0, 0.0]),
            Document(id="doc2", text="Java guide", embedding=[0.0, 1.0, 0.0]),
            Document(id="doc3", text="Rust guide", embedding=[0.0, 0.0, 1.0]),
        ]
        store.add(docs)

        retriever = Retriever(store=store, top_k=2)
        result = retriever.retrieve(query_embedding=[1.0, 0.0, 0.0])
        assert len(result.documents) == 2
        assert result.documents[0].id == "doc1"

    def test_score_threshold(self):
        store = ChromaStore(collection_name="test_threshold")
        docs = [
            Document(id="doc1", text="very relevant", embedding=[1.0, 0.0]),
            Document(id="doc2", text="not relevant", embedding=[0.0, 1.0]),
        ]
        store.add(docs)

        retriever = Retriever(store=store, top_k=10, score_threshold=0.5)
        result = retriever.retrieve(query_embedding=[1.0, 0.0])
        # Should filter low-scoring results
        assert len(result.documents) >= 1
