from typing import Any, List

import pinecone
from langchain.callbacks.manager import Callbacks
from langchain.docstore.document import Document
from langchain.schema import BaseRetriever

from .secret import Secret
from .vectorstore import PineconeStore


class PineconeRetriever(BaseRetriever):
    def __init__(
        self,
        embedding_api: str,
        embedding_model: str,
        pinecone_api_key: Secret,
        pinecone_environment: str,
        pinecone_index: str,
        namespaces: list = [],
        text_key: str = "content",
    ):
        pinecone.init(api_key=pinecone_api_key, environment=pinecone_environment)
        super(PineconeRetriever, self).__init__()
        self.__dict__["_retriever"] = PineconeStore(
            embedding_api,
            embedding_model,
            pinecone_index,
            namespaces,
            text_key,
        ).as_retriever()

    def _get_relevant_documents(self, query: str, **kwargs: Any) -> List[Document]:
        return self._retriever.get_relevant_documents(query, **kwargs)

    async def _aget_relevant_documents(
        self, query: str, **kwargs: Any
    ) -> List[Document]:
        raise NotImplementedError("aget_relevant_documents not implemented yet")