"""
Search engine implementation that wraps any LangChain retriever.
This allows using vector stores, databases, or any custom retriever as a search source in LDR.
"""

from typing import Any, Dict, List

from langchain_core.retrievers import BaseRetriever, Document
from loguru import logger

from ..search_engine_base import BaseSearchEngine


class RetrieverSearchEngine(BaseSearchEngine):
    """
    Search engine that uses any LangChain retriever.

    This allows users to plug in any LangChain retriever (vector stores,
    databases, custom implementations) and use it as a search engine in LDR.
    """

    def __init__(
        self,
        retriever: BaseRetriever,
        max_results: int = 10,
        name: str = None,
        **kwargs,
    ):
        """
        Initialize the retriever-based search engine.

        Args:
            retriever: Any LangChain BaseRetriever instance
            max_results: Maximum number of results to return
            name: Display name for this retriever (defaults to retriever class name)
            **kwargs: Additional parameters passed to parent
        """
        super().__init__(max_results=max_results, **kwargs)
        self.retriever = retriever
        self.name = name if name is not None else retriever.__class__.__name__

    def run(
        self, query: str, research_context: Dict[str, Any] | None = None
    ) -> List[Dict[str, Any]]:
        """
        Execute search using the LangChain retriever.

        Args:
            query: Search query
            research_context: Context from previous research to use.

        Returns:
            List of search results in LDR format
        """
        try:
            # Use the retriever to get relevant documents
            docs = self.retriever.invoke(query)

            # Convert LangChain documents to LDR search result format
            results = []
            for i, doc in enumerate(docs[: self.max_results]):
                result = self._convert_document_to_result(doc, i)
                results.append(result)

            logger.info(
                f"Retriever '{self.name}' returned {len(results)} results for query: {query}"
            )
            return results

        except Exception:
            logger.exception("Error in retriever search")
            return []

    def _convert_document_to_result(
        self, doc: Document, index: int
    ) -> Dict[str, Any]:
        """
        Convert a LangChain Document to LDR search result format.

        Args:
            doc: LangChain Document
            index: Result index

        Returns:
            Search result in LDR format
        """
        # Extract metadata
        metadata = doc.metadata or {}

        # Build the result
        result = {
            # Required fields for LDR
            "title": metadata.get("title", f"Document {index + 1}"),
            "url": metadata.get(
                "source",
                metadata.get("url", f"retriever://{self.name}/doc_{index}"),
            ),
            "snippet": doc.page_content[:500] if doc.page_content else "",
            # Optional fields
            "full_content": doc.page_content,
            "author": metadata.get("author", ""),
            "date": metadata.get("date", ""),
            # Include all metadata for flexibility
            "metadata": metadata,
            # Score if available
            "score": metadata.get("score", 1.0),
            # Source information
            "source": self.name,
            "retriever_type": self.retriever.__class__.__name__,
        }

        return result

    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information from the retriever.

        Args:
            query: Search query

        Returns:
            List of preview dictionaries
        """
        try:
            # Use the retriever to get relevant documents
            docs = self.retriever.invoke(query)

            # Convert to preview format
            previews = []
            for i, doc in enumerate(docs[: self.max_results]):
                preview = self._convert_document_to_result(doc, i)
                previews.append(preview)

            logger.info(
                f"Retriever '{self.name}' returned {len(previews)} previews for query: {query}"
            )
            return previews

        except Exception:
            logger.exception("Error getting previews from retriever")
            return []

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        For retrievers, previews already contain full content.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            Same list with full content (already included)
        """
        # For retrievers, the preview already contains the full content
        # Just ensure the 'full_content' field is present
        for item in relevant_items:
            if "full_content" not in item and "snippet" in item:
                item["full_content"] = item["snippet"]
        return relevant_items

    async def arun(self, query: str) -> List[Dict[str, Any]]:
        """
        Async version of search using the retriever.

        Args:
            query: Search query

        Returns:
            List of search results in LDR format
        """
        try:
            # Use async retriever if available
            if hasattr(self.retriever, "aget_relevant_documents"):
                docs = await self.retriever.aget_relevant_documents(query)
            else:
                # Fall back to sync version
                logger.debug(
                    f"Retriever '{self.name}' doesn't support async, using sync version"
                )
                return self.run(query)

            # Convert documents to results
            results = []
            for i, doc in enumerate(docs[: self.max_results]):
                result = self._convert_document_to_result(doc, i)
                results.append(result)

            logger.info(
                f"Retriever '{self.name}' returned {len(results)} async results for query: {query}"
            )
            return results

        except Exception:
            logger.exception("Error in async retriever search")
            return []
