"""Tests for the RAG pipeline orchestrator."""

from bifrost_rag.pipeline import RAGPipeline, RAGResult
from bifrost_rag.store import ChromaStore, Document


class TestRAGResult:
    def test_creation(self):
        result = RAGResult(
            query="test query",
            answer="test answer",
            sources=[Document(id="doc1", text="source text")],
            scores=[0.95],
        )
        assert result.query == "test query"
        assert result.answer == "test answer"
        assert len(result.sources) == 1

    def test_no_answer(self):
        result = RAGResult(query="test", answer="", sources=[], scores=[])
        assert result.answer == ""
        assert len(result.sources) == 0


class TestRAGPipeline:
    def test_creation(self):
        store = ChromaStore(collection_name="test_pipeline")
        pipeline = RAGPipeline(store=store)
        assert pipeline is not None

    def test_ingest_documents(self):
        store = ChromaStore(collection_name="test_ingest")
        pipeline = RAGPipeline(store=store)
        docs = [
            Document(id="doc1", text="Python is great", embedding=[1.0, 0.0]),
            Document(id="doc2", text="Java is verbose", embedding=[0.0, 1.0]),
        ]
        count = pipeline.ingest(docs)
        assert count == 2
        assert store.count() == 2

    def test_query_returns_rag_result(self):
        store = ChromaStore(collection_name="test_query_result")
        docs = [
            Document(id="doc1", text="Python is a language", embedding=[1.0, 0.0, 0.0]),
            Document(id="doc2", text="Java is a language", embedding=[0.0, 1.0, 0.0]),
        ]
        store.add(docs)

        pipeline = RAGPipeline(store=store, top_k=1)
        result = pipeline.query(query_embedding=[1.0, 0.0, 0.0])
        assert isinstance(result, RAGResult)
        assert len(result.sources) >= 1

    def test_pipeline_config(self):
        store = ChromaStore(collection_name="test_config")
        pipeline = RAGPipeline(store=store, top_k=3)
        assert pipeline.top_k == 3
