"""LlamaIndex retriever for aqoon.

Usage::

    from aqoon.integrations.llamaindex import AqoonRetriever

    retriever = AqoonRetriever(api_key="aqn_your_key")
    nodes = retriever.retrieve("machine learning basics")

    # In a query engine:
    from llama_index.core.query_engine import RetrieverQueryEngine

    query_engine = RetrieverQueryEngine.from_args(retriever)
    response = query_engine.query("What is supervised learning?")
"""

from __future__ import annotations

from llama_index.core.retrievers import BaseRetriever
from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode


class AqoonRetriever(BaseRetriever):
    """Retrieve nodes 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
        super().__init__()

    def _retrieve(self, query_bundle: QueryBundle) -> list[NodeWithScore]:
        from aqoon import Aqoon

        with Aqoon(api_key=self._api_key, base_url=self._base_url) as client:
            response = client.search(
                query_bundle.query_str,
                collection=self._collection,
                limit=self._k,
            )

        return [
            NodeWithScore(
                node=TextNode(
                    text=r["content"],
                    metadata={
                        "title": r["title"],
                        "document_id": r["document_id"],
                        "collection_name": r.get("collection_name", ""),
                        "tags": ",".join(r.get("tags", [])),
                    },
                ),
                score=r["score"],
            )
            for r in response["results"]
        ]
