"""Retrieval module for RAG pipelines."""

from __future__ import annotations

from dataclasses import dataclass, field

from bifrost_rag.store import Document, VectorStore


@dataclass
class RetrievalResult:
    """Result of a retrieval operation."""

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

    @property
    def top_document(self) -> Document | None:
        return self.documents[0] if self.documents else None


class Retriever:
    """Retrieves relevant documents from a vector store."""

    def __init__(
        self,
        store: VectorStore,
        top_k: int = 5,
        score_threshold: float | None = None,
    ):
        self.store = store
        self.top_k = top_k
        self.score_threshold = score_threshold

    def retrieve(
        self,
        query_embedding: list[float],
        query_text: str = "",
    ) -> RetrievalResult:
        """Retrieve documents matching the query embedding."""
        documents = self.store.query(
            query_embedding=query_embedding,
            top_k=self.top_k,
        )

        # ChromaDB doesn't return scores directly in the simplified query,
        # so we pass empty scores for now. Production would use distances.
        scores = [1.0] * len(documents)

        return RetrievalResult(
            query=query_text,
            documents=documents,
            scores=scores,
        )
