"""Provides an adapter for OpenSearch vector store integration."""

from collections.abc import Sequence
from typing import Any

from typing_extensions import override

try:
    from langchain_community.vectorstores import OpenSearchVectorSearch
except (ImportError, ModuleNotFoundError):
    raise ImportError("please `pip install langchain-community opensearch-py`")

from langchain_core.documents import Document

from langchain_graph_retriever._conversion import METADATA_EMBEDDING_KEY
from langchain_graph_retriever.adapters.langchain import LangchainAdapter


class OpenSearchAdapter(LangchainAdapter[OpenSearchVectorSearch]):
    """
    Adapter to traverse OpenSearch vector stores.

    This adapter enables similarity search and document retrieval using an
    OpenSearch vector store. It supports both "lucene" and "faiss" engines.

    Parameters
    ----------
    vector_store  : OpenSearchVectorSearch
        The OpenSearch vector store instance.
    """

    def __init__(self, vector_store: OpenSearchVectorSearch):
        if vector_store.engine not in ["lucene", "faiss"]:
            msg = (
                f"Invalid engine for Traversal: '{self.vector_store.engine}'"
                " please instantiate the Open Search Vector Store with"
                " either the 'lucene' or 'faiss' engine"
            )
            raise ValueError(msg)
        super().__init__(vector_store)

    def _build_filter(
        self, filter: dict[str, str] | None = None
    ) -> list[dict[str, Any]] | None:
        """
        Build a filter query for OpenSearch based on metadata.

        Parameters
        ----------
        filter : dict[str, str], optional
            Metadata filter to apply.

        Returns
        -------
        list[dict[str, Any]] | None
            Filter query for OpenSearch.
        """
        if filter is None:
            return None
        return [
            {
                "terms" if isinstance(value, list) else "term": {
                    f"metadata.{key}.keyword": value
                }
            }
            for key, value in filter.items()
        ]

    @override
    def _similarity_search_with_embedding_by_vector(
        self,
        embedding: list[float],
        k: int = 4,
        filter: dict[str, str] | None = None,
        **kwargs: Any,
    ) -> list[Document]:
        if filter is not None:
            # use an efficient_filter to collect results that
            # are near the embedding vector until up to 'k'
            # documents that match the filter are found.
            kwargs["efficient_filter"] = {
                "bool": {"must": self._build_filter(filter=filter)}
            }

        if k == 0:
            return []

        docs = self.vector_store.similarity_search_by_vector(
            embedding=embedding,
            k=k,
            metadata_field="*",
            **kwargs,
        )

        # if metadata=="*" on the search, then the document
        # embedding vector and text are included in the
        # document metadata in the returned document.
        #
        # The actual document metadata is moved down into a
        # sub "metadata" key.
        return [
            Document(
                id=doc.id,
                page_content=doc.page_content,
                metadata={
                    METADATA_EMBEDDING_KEY: doc.metadata["vector_field"],
                    **doc.metadata["metadata"],
                },
            )
            for doc in docs
        ]

    @override
    def _get(self, ids: Sequence[str], /, **kwargs: Any) -> list[Document]:
        try:
            from opensearchpy.exceptions import NotFoundError
        except (ImportError, ModuleNotFoundError):
            msg = "please `pip install opensearch-py`."
            raise ImportError(msg)

        docs: list[Document] = []
        for id in ids:
            try:
                hit = self.vector_store.client.get(
                    index=self.vector_store.index_name,
                    id=id,
                    _source_includes=["text", "metadata", "vector_field"],
                    **kwargs,
                )
                docs.append(
                    Document(
                        page_content=hit["_source"]["text"],
                        metadata={
                            METADATA_EMBEDDING_KEY: hit["_source"]["vector_field"],
                            **hit["_source"]["metadata"],
                        },
                        id=hit["_id"],
                    )
                )
            except NotFoundError:
                pass
        return docs
