"""Base vector store index.

An index that that is built on top of an existing vector store.

"""

from abc import abstractmethod
from typing import Any, Dict, Generic, List, Optional, Sequence, Set, Tuple, TypeVar

from llama_index.data_structs.data_structs import IndexStruct, Node
from llama_index.embeddings.base import BaseEmbedding
from llama_index.indices.base import DOCUMENTS_INPUT, BaseGPTIndex
from llama_index.langchain_helpers.chain_wrapper import LLMPredictor
from llama_index.langchain_helpers.text_splitter import TextSplitter
from llama_index.prompts.default_prompts import DEFAULT_TEXT_QA_PROMPT
from llama_index.prompts.prompts import QuestionAnswerPrompt
from llama_index.schema import BaseDocument
from llama_index.utils import get_new_id

BID = TypeVar("BID", bound=IndexStruct)


class BaseGPTVectorStoreIndex(BaseGPTIndex[BID], Generic[BID]):
    """Base GPT Vector Store Index.

    Args:
        text_qa_template (Optional[QuestionAnswerPrompt]): A Question-Answer Prompt
            (see :ref:`Prompt-Templates`).
        embed_model (Optional[BaseEmbedding]): Embedding model to use for
            embedding similarity.
    """

    def __init__(
        self,
        documents: Optional[Sequence[DOCUMENTS_INPUT]] = None,
        index_struct: Optional[BID] = None,
        text_qa_template: Optional[QuestionAnswerPrompt] = None,
        llm_predictor: Optional[LLMPredictor] = None,
        embed_model: Optional[BaseEmbedding] = None,
        text_splitter: Optional[TextSplitter] = None,
        **kwargs: Any,
    ) -> None:
        """Initialize params."""
        self.text_qa_template = text_qa_template or DEFAULT_TEXT_QA_PROMPT
        super().__init__(
            documents=documents,
            index_struct=index_struct,
            llm_predictor=llm_predictor,
            embed_model=embed_model,
            text_splitter=text_splitter,
            **kwargs,
        )

    def _get_node_embedding_tups(
        self, nodes: List[Node], existing_node_ids: Set
    ) -> List[Tuple[str, Node, List[float]]]:
        """Get tuples of id, node, and embedding.

        Allows us to store these nodes in a vector store.
        Embeddings are called in batches.

        """
        id_to_node_map: Dict[str, Node] = {}
        id_to_embed_map: Dict[str, List[float]] = {}

        for n in nodes:
            new_id = get_new_id(existing_node_ids.union(id_to_node_map.keys()))
            if n.embedding is None:
                self._embed_model.queue_text_for_embeddding(new_id, n.get_text())
            else:
                id_to_embed_map[new_id] = n.embedding

            id_to_node_map[new_id] = n

        # call embedding model to get embeddings
        result_ids, result_embeddings = self._embed_model.get_queued_text_embeddings()
        for new_id, text_embedding in zip(result_ids, result_embeddings):
            id_to_embed_map[new_id] = text_embedding

        result_tups = []
        for id, embed in id_to_embed_map.items():
            result_tups.append((id, id_to_node_map[id], embed))
        return result_tups

    @abstractmethod
    def _add_document_to_index(
        self,
        index_struct: BID,
        document: BaseDocument,
    ) -> None:
        """Add document to index."""

    def _build_fallback_text_splitter(self) -> TextSplitter:
        # if not specified, use "smart" text splitter to ensure chunks fit in prompt
        return self._prompt_helper.get_text_splitter_given_prompt(
            self.text_qa_template, 1
        )

    def _build_index_from_documents(
        self,
        documents: Sequence[BaseDocument],
    ) -> BID:
        """Build index from documents."""
        index_struct = self.index_struct_cls()
        for d in documents:
            self._add_document_to_index(index_struct, d)
        return index_struct

    def _insert(self, document: BaseDocument, **insert_kwargs: Any) -> None:
        """Insert a document."""
        self._add_document_to_index(self._index_struct, document)
