import logging
import re
from typing import Any, Callable, Iterable, List, Optional

import pinecone
import requests
from environs import Env
from langchain.docstore.document import Document
from langchain.embeddings.base import Embeddings
from langchain.vectorstores.base import VectorStore

from exceptions import EmbeddingError

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,
        embedding_api: str,
        embedding_model: str,
        pinecone_index: str,
        namespaces=[],
        text_key="content",
    ):
        super(PineconeStore).__init__()
        self.index = pinecone.Index(pinecone_index)
        self.namespaces = namespaces
        self._text_key = text_key
        self.embedding_api = embedding_api
        self.embedding_model = embedding_model

    @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 embed(self, query: str) -> List[float]:
        try:
            r = requests.post(
                f"{self.embedding_api}/embeddings",
                json={
                    "input": query,
                    "model": self.embedding_model,
                },
            )
            return r.json()["data"][0]["embedding"]
        except Exception as e:
            raise EmbeddingError(str(e)) from e

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

    def similarity_search(
        self, query: str, k: int = 16, **kwargs: Any
    ) -> List[Document]:
        """Return docs most similar to query."""
        query_embedding = self.embed(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