import re
import logging
import pinecone
from environs import Env
from observability import span
from langchain.docstore.document import Document
from langchain.embeddings.base import Embeddings
from langchain.vectorstores.base import VectorStore
from typing import Any, Iterable, List, Optional, Callable


logger = logging.getLogger(__name__)

def remove_images_from_markdown(markdown_text):
    pattern = r"!\[.*?\]\(.*?\)"
    replaced_text = re.sub(pattern, "", markdown_text)
    return replaced_text

def remove_links_from_markdown(markdown_text):
    pattern = r"\[(.*?)\]\(.*?\)"
    replaced_text = re.sub(pattern, r"\1", markdown_text)
    return replaced_text

class PineconeStore(VectorStore):
    def __init__(self, env: Env, embedding: Embeddings, namespaces=[], text_key='content'):
        super(PineconeStore)
        self.embedding = embedding
        self.index = pinecone.Index(env.str('DOCUMENT_INDEX_NAME'))
        self.namespaces = namespaces
        self._text_key = text_key
    
    @classmethod
    def from_texts(
        cls,
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        **kwargs: Any,
    ):
        """Return VectorStore initialized from texts and embeddings."""
        raise NotImplementedError

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        **kwargs: Any,
    ) -> List[str]:
        raise NotImplementedError

    @span(name="select documents")
    def similarity_search(
        self, query: str, k: int = 16, **kwargs: Any
    ) -> List[Document]:
        """Return docs most similar to query."""
        query_embedding = self.embedding.embed_query(query)
        
        matches = []
        for ns in self.namespaces:
            ms = self.index.query(
                vector=query_embedding,
                top_k=k,
                include_metadata=True, include_values=False,
                namespace=ns,
            ).matches
            for m in ms:
                m.metadata['namespace'] = ns
            matches += ms
        relevant = sorted(matches, key=lambda x: x.score, reverse=True)[:k]
        
        docs = []
        for r in relevant:
            metadata = r['metadata']
            if self._text_key in metadata:
                text = metadata.pop(self._text_key)
                # remove images
                text = remove_images_from_markdown(text)
                # remove links
                text = remove_links_from_markdown(text)
                docs.append(Document(page_content=text, metadata=metadata))
            else:
                logger.warning(
                    f"Found document with no `{self._text_key}` key. Skipping."
                )
        return docs