"""
HieromemRetriever - LangChain BaseRetriever implementation.

Retrieves documents from HIEROMEM using hybrid vector + graph retrieval.
"""
from typing import Any, Dict, List, Optional
import httpx

try:
    from langchain.schema import Document
    from langchain.schema.retriever import BaseRetriever
    from langchain.callbacks.manager import CallbackManagerForRetrieverRun
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    BaseRetriever = object
    Document = object
    CallbackManagerForRetrieverRun = object


class HieromemRetriever(BaseRetriever):  # type: ignore[misc]
    """
    LangChain retriever backed by HIEROMEM.

    Usage:
        from hieromem.integrations.langchain import HieromemRetriever

        retriever = HieromemRetriever(
            hieromem_url="http://localhost:8000",
            api_key="sk-...",
            namespace="knowledge_base"
        )

        docs = retriever.get_relevant_documents("What is HIEROMEM?")
    """

    hieromem_url: str = "http://localhost:8000"
    api_key: Optional[str] = None
    namespace: str = "default"
    k: int = 5
    mode: str = "graph_hybrid"
    alpha_dense: float = 1.0
    beta_symbolic: float = 50.0
    min_confidence: float = 0.0

    _client: Optional[httpx.Client] = None

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"

    def __init__(self, **kwargs: Any):
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is required for HieromemRetriever. "
                "Install with: pip install langchain"
            )
        super().__init__(**kwargs)
        self._client = httpx.Client(
            base_url=self.hieromem_url,
            headers={"Authorization": f"Bearer {self.api_key}"} if self.api_key else {},
            timeout=30.0,
        )

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: Optional[CallbackManagerForRetrieverRun] = None,
    ) -> List[Document]:
        """
        Retrieve relevant documents from HIEROMEM.

        Args:
            query: The search query
            run_manager: Callback manager (optional)

        Returns:
            List of LangChain Document objects
        """
        if self._client is None:
            return []

        try:
            response = self._client.post(
                "/memory/retrieve",
                json={
                    "query": query,
                    "k": self.k,
                    "namespace": self.namespace,
                    "mode": self.mode,
                    "alpha_dense": self.alpha_dense,
                    "beta_symbolic": self.beta_symbolic,
                    "min_confidence": self.min_confidence,
                },
            )
            response.raise_for_status()
            episodes = response.json().get("episodes", [])

            return self._episodes_to_documents(episodes)

        except httpx.HTTPError:
            return []

    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: Optional[CallbackManagerForRetrieverRun] = None,
    ) -> List[Document]:
        """Async version of document retrieval."""
        async with httpx.AsyncClient(
            base_url=self.hieromem_url,
            headers={"Authorization": f"Bearer {self.api_key}"} if self.api_key else {},
            timeout=30.0,
        ) as client:
            try:
                response = await client.post(
                    "/memory/retrieve",
                    json={
                        "query": query,
                        "k": self.k,
                        "namespace": self.namespace,
                        "mode": self.mode,
                        "alpha_dense": self.alpha_dense,
                        "beta_symbolic": self.beta_symbolic,
                        "min_confidence": self.min_confidence,
                    },
                )
                response.raise_for_status()
                episodes = response.json().get("episodes", [])
                return self._episodes_to_documents(episodes)
            except httpx.HTTPError:
                return []

    def _episodes_to_documents(self, episodes: List[Dict[str, Any]]) -> List[Document]:
        """Convert HIEROMEM episodes to LangChain Documents."""
        documents = []
        for ep in episodes:
            meta = ep.get("meta", {})
            content = meta.get("text", "")

            # Build metadata for Document
            doc_meta = {
                "episode_id": ep.get("id"),
                "trust": ep.get("trust"),
                "timestamp": ep.get("timestamp"),
                "namespace": self.namespace,
                "score": meta.get("score"),
            }

            # Include graph info if present
            if ep.get("graph"):
                doc_meta["graph_nodes"] = len(ep["graph"].get("nodes", []))
                doc_meta["graph_edges"] = len(ep["graph"].get("edges", []))

            documents.append(Document(page_content=content, metadata=doc_meta))

        return documents

    def __del__(self) -> None:
        """Cleanup HTTP client."""
        if self._client:
            self._client.close()
