"""Haystack retriever component for aqoon.

Usage::

    from aqoon.integrations.haystack import AqoonRetriever

    retriever = AqoonRetriever(api_key="aqn_your_key")
    result = retriever.run(query="machine learning basics")
    docs = result["documents"]

    # In a pipeline:
    from haystack import Pipeline

    pipe = Pipeline()
    pipe.add_component("retriever", AqoonRetriever(api_key="aqn_..."))
    pipe.add_component("llm", ...)
    pipe.connect("retriever.documents", "llm.documents")
"""

from __future__ import annotations

from haystack import Document, component


@component
class AqoonRetriever:
    """Retrieve documents from aqoon knowledge base.

    Args:
        api_key: aqoon API key (starts with ``aqn_``).
        base_url: aqoon instance URL. Defaults to ``https://aqoon.ai``.
        collection: Optional collection slug to scope the search.
        k: Number of results to return. Defaults to 5.
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://aqoon.ai",
        collection: str | None = None,
        k: int = 5,
    ) -> None:
        self.api_key = api_key
        self.base_url = base_url
        self.collection = collection
        self.k = k

    @component.output_types(documents=list[Document])
    def run(self, query: str) -> dict[str, list[Document]]:
        from aqoon import Aqoon

        with Aqoon(api_key=self.api_key, base_url=self.base_url) as client:
            response = client.search(query, collection=self.collection, limit=self.k)

        return {
            "documents": [
                Document(
                    content=r["content"],
                    meta={
                        "title": r["title"],
                        "document_id": r["document_id"],
                        "collection_name": r.get("collection_name", ""),
                        "tags": r.get("tags", []),
                        "score": r["score"],
                    },
                )
                for r in response["results"]
            ],
        }
