"""LangChain retriever for aqoon.

Usage::

    from aqoon.integrations.langchain import AqoonRetriever

    retriever = AqoonRetriever(api_key="aqn_your_key")
    docs = retriever.invoke("machine learning basics")

    # In a chain:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_core.runnables import RunnablePassthrough

    chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
    )
"""

from __future__ import annotations

from typing import Optional

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever


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

    api_key: str
    base_url: str = "https://aqoon.ai"
    collection: str | None = None
    k: int = 5

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun | None = None,
    ) -> 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 [
            Document(
                page_content=r["content"],
                metadata={
                    "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"]
        ]
