"""RAG pipeline orchestrator."""

from __future__ import annotations

from dataclasses import dataclass, field

from bifrost_rag.store import Document, VectorStore


@dataclass
class RAGResult:
    """Result of a RAG query."""

    query: str
    answer: str
    sources: list[Document]
    scores: list[float] = field(default_factory=list)


class RAGPipeline:
    """Orchestrates the full RAG pipeline: ingest, retrieve, generate.

    The pipeline connects chunking, embedding, vector storage,
    and retrieval into a single interface.
    """

    def __init__(
        self,
        store: VectorStore,
        top_k: int = 5,
    ):
        self.store = store
        self.top_k = top_k

    def ingest(self, documents: list[Document]) -> int:
        """Ingest documents into the vector store.

        Returns the number of documents ingested.
        """
        self.store.add(documents)
        return len(documents)

    def query(
        self,
        query_embedding: list[float],
        query_text: str = "",
    ) -> RAGResult:
        """Query the RAG pipeline with an embedding.

        Returns retrieved sources. For LLM-generated answers,
        use query_with_llm() instead.
        """
        documents = self.store.query(
            query_embedding=query_embedding,
            top_k=self.top_k,
        )

        # Build context from retrieved documents
        context_parts = [doc.text for doc in documents]
        answer = "\n---\n".join(context_parts) if context_parts else ""

        return RAGResult(
            query=query_text,
            answer=answer,
            sources=documents,
            scores=[1.0] * len(documents),
        )
