"""Velixar retriever for LangChain."""

from typing import List, Optional

import httpx
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from pydantic import Field

from ._http import create_client, request_with_retry


class VelixarRetriever(BaseRetriever):
    """Retrieve memories from Velixar AI's cognitive memory API.

    Example:
        from velixar_langchain import VelixarRetriever

        retriever = VelixarRetriever(
            api_key="vx_...",
            workspace_id="ws_abc123",
        )
        docs = retriever.invoke("What does the user prefer?")
    """

    api_key: str
    workspace_id: str
    base_url: str = "https://api.velixarai.com"
    top_k: int = 5
    min_salience: float = 0.0
    include_org: bool = False
    _client: Optional[httpx.Client] = None

    class Config:
        arbitrary_types_allowed = True
        underscore_attrs_are_private = True

    def _get_client(self) -> httpx.Client:
        if self._client is None:
            self._client = create_client(self.base_url, self.api_key)
        return self._client

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        resp = request_with_retry(
            self._get_client(), "GET", "/v1/memory/search",
            params={
                "q": query,
                "workspace_id": self.workspace_id,
                "limit": self.top_k,
                "include_org": str(self.include_org).lower(),
            },
        )
        resp.raise_for_status()
        data = resp.json()

        docs = []
        for mem in data.get("memories", []):
            salience = mem.get("salience", 0)
            if salience < self.min_salience:
                continue
            docs.append(Document(
                page_content=mem.get("content", ""),
                metadata={
                    "memory_id": mem.get("id", ""),
                    "tier": mem.get("tier", 2),
                    "salience": salience,
                    "tags": mem.get("tags", []),
                    "score": mem.get("score", 0),
                    "source": "velixar",
                },
            ))
        return docs
