"""Class definition for QueryBot."""
import contextvars
import hashlib
from pathlib import Path
from typing import List, Union

from langchain.callbacks.base import BaseCallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.schema import AIMessage, HumanMessage, SystemMessage
from llama_index import (
    GPTVectorStoreIndex,
    LLMPredictor,
    ServiceContext,
    load_index_from_storage,
)
from llama_index.node_parser import SimpleNodeParser
from llama_index.retrievers import VectorIndexRetriever
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index.storage.storage_context import StorageContext
from llama_index.vector_stores import SimpleVectorStore

from llamabot.doc_processor import magic_load_doc, split_document
from llamabot.recorder import autorecord

# from loguru import logger


prompt_recorder_var = contextvars.ContextVar("prompt_recorder")


class QueryBot:
    """QueryBot is a bot that lets us use GPT4 to query documents."""

    def __init__(
        self,
        system_message: str,
        model_name="gpt-4-32k",
        temperature=0.0,
        doc_path: Union[str, Path] = None,
        saved_index_path: Union[str, Path] = None,
        chunk_size: int = 2000,
        chunk_overlap: int = 0,
        stream=True,
    ):
        """Initialize QueryBot.

        Pass in either the doc_path or saved_index_path to initialize the QueryBot.

        NOTE: QueryBot is not designed to have memory!

        The default text splitter is the TokenTextSplitter from LangChain.
        The default index that we use is the GPTVectorStoreIndex from LlamaIndex.
        We also default to using GPT4 with temperature 0.0.

        :param system_message: The system message to send to the chatbot.
        :param model_name: The name of the OpenAI model to use.
        :param temperature: The model temperature to use.
            See https://platform.openai.com/docs/api-reference/completions/create#completions/create-temperature
            for more information.
        :param doc_path: A path to document to use for the chatbot.
        :param saved_index_path: The path to the saved index to use for the chatbot.
        :param chunk_size: The chunk size to use for the LlamaIndex TokenTextSplitter.
        :param chunk_overlap: The chunk overlap to use for the LlamaIndex TokenTextSplitter.
        :param stream: Whether to stream the chatbot or not.
        """

        chat = ChatOpenAI(
            model_name=model_name,
            temperature=temperature,
            streaming=True,
            verbose=True,
            callback_manager=BaseCallbackManager(
                handlers=[StreamingStdOutCallbackHandler()] if stream else []
            ),
        )
        llm_predictor = LLMPredictor(llm=chat)
        service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)

        # Initialize index or load it from disk.
        if saved_index_path is None:
            index = GPTVectorStoreIndex(nodes=[], service_context=service_context)
        else:
            index = GPTVectorStoreIndex.load_from_disk(
                saved_index_path, service_context=service_context
            )

        # Update index with a new document.
        if doc_path is not None:
            index = make_or_load_index(doc_path, chunk_size, chunk_overlap)

        # Set object attributes.
        self.system_message = system_message
        self.index = index
        self.doc_path = doc_path
        self.chat = chat
        self.chat_history = [
            SystemMessage(content=system_message),
            SystemMessage(
                content="""Do not hallucinate content.
If you cannot answer something, respond by saying that you don't know.
"""
            ),
        ]
        # Store a mapping of query to source nodes.
        self.source_nodes: dict = {}
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def __call__(
        self,
        query: str,
        similarity_top_k=10,
    ) -> AIMessage:
        """Call the QueryBot.

        :param query: The query to send to the document index.
        :param similarity_top_k: The number of documents to return from the index.
            These documents are added to the context of the chat history
            and then used to synthesize the response.
        :raises ValueError: if the index is not initialized.
        :return: The response to the query generated by GPT4.
        """
        if self.index is None:
            raise ValueError(
                "You need to provide a document for querybot to index against!"
            )
        # Step 1: Get documents from the index that are deemed to be matching the query.
        # logger.info(f"Querying index for top {similarity_top_k} documents...")

        retriever = VectorIndexRetriever(
            index=self.index,
            similarity_top_k=similarity_top_k,
        )

        source_nodes = retriever.retrieve(query)
        source_texts = [n.node.text for n in source_nodes]

        # Step 2: Construct a faux message history to work with.
        faux_chat_history = [SystemMessage(content=self.system_message)]
        faux_chat_history.append(
            SystemMessage(content="Here is the context you will be working with:")
        )
        for text in source_texts:
            faux_chat_history.append(SystemMessage(content=text))

        faux_chat_history.append(
            SystemMessage(content="Based on this context, answer the following query:")
        )

        faux_chat_history.append(HumanMessage(content=query))

        # Step 3: Send the chat history through the model
        response = self.chat(faux_chat_history)

        # Step 4: Record only the human response and the GPT response but not the original.
        self.chat_history.append(HumanMessage(content=query))
        self.chat_history.append(response)

        # Step 5: Record the source nodes of the query.
        self.source_nodes[query] = source_nodes

        autorecord(query, response.content)

        # Step 6: Return the response.
        return response

    def save(self, path: Union[str, Path]):
        """Save the QueryBot index to disk.

        :param path: The path to save the QueryBot index.
        """
        path = Path(path)
        if not path.suffix == ".json":
            path = path.with_suffix(".json")
        self.index.save_to_disk(path)

    def insert(self, file: Path):
        """Insert a document into the index.

        :param file: The path to the document to insert.
        """

        self.index = insert_documents_into_index(
            [file], self.index, self.chunk_size, self.chunk_overlap
        )


def insert_documents_into_index(
    doc_path: List[Union[str, Path]],
    index: GPTVectorStoreIndex,
    chunk_size: int,
    chunk_overlap: int,
):
    """Insert documents into the index.

    :param doc_path: A list of paths to the documents to insert.
    :param index: The index to insert the documents into.
    :param chunk_size: The chunk size to use for the LangChain TokenTextSplitter.
    :param chunk_overlap: The chunk overlap to use for the LangChain TokenTextSplitter.
    :returns: The updated index.
    """
    # logger.info(f"Inserting {path} into the index...")
    documents = magic_load_doc(doc_path)
    for document in documents:
        chunks = split_document(document, chunk_size, chunk_overlap)
        for chunk in chunks:
            index.insert(chunk)
    return index


def compute_file_hash(fpath: Path) -> str:
    """Compute the hash of a file.

    :param fpath: The path to the file to compute the hash of.
    :returns: The hash of the file.
    """
    file_content = fpath.read_bytes()
    return hashlib.sha256(file_content).hexdigest()


def make_service_context():
    """Make a service context for the QueryBot.

    :returns: A service context.
    """
    chat = ChatOpenAI(
        model_name="gpt-4-32k",
        temperature=0.0,
        streaming=True,
        verbose=True,
        callback_manager=BaseCallbackManager([StreamingStdOutCallbackHandler()]),
    )
    llm_predictor = LLMPredictor(llm=chat)
    service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)
    return service_context


def get_persist_dir(file_hash: str):
    """Get the persist directory for a given file hash.

    :param file_hash: The file hash to use for the persist directory.
    :returns: The persist directory.
    """
    persist_dir = Path.home() / ".llamabot" / "cache" / file_hash
    persist_dir.mkdir(parents=True, exist_ok=True)
    return persist_dir


def load_index(persist_dir, service_context):
    """Load an index from disk.

    :param persist_dir: The directory to load the index from.
    :param service_context: The service context to use for the index.
    :returns: The index.
    """
    storage_context = StorageContext.from_defaults(
        docstore=SimpleDocumentStore.from_persist_dir(persist_dir=persist_dir),
        vector_store=SimpleVectorStore.from_persist_dir(persist_dir=persist_dir),
        index_store=SimpleIndexStore.from_persist_dir(persist_dir=persist_dir),
    )
    index = load_index_from_storage(storage_context, service_context=service_context)
    return index


def make_index(docs, file_hash, persist_dir, service_context):
    """Make an index from a list of documents.

    :param docs: A list of documents to index.
    :param file_hash: The hash of the file to use as the index id.
    :param persist_dir: The directory to persist the index to.
    :param service_context: The service context to use for the index.
    :returns: The index.
    """
    # create parser and parse document into nodes
    parser = SimpleNodeParser()
    nodes = parser.get_nodes_from_documents(docs)

    # create (or load) docstore and add nodes
    storage_context = StorageContext.from_defaults(
        docstore=SimpleDocumentStore(),
        vector_store=SimpleVectorStore(),
        index_store=SimpleIndexStore(),
    )
    storage_context.docstore.add_documents(nodes)

    index = GPTVectorStoreIndex(
        nodes,
        storage_context=storage_context,
        index_id=file_hash,
        service_context=service_context,
    )
    index.storage_context.persist(persist_dir=persist_dir)
    return index


def make_or_load_index(doc_path, chunk_size=2000, chunk_overlap=0):
    """Make or load an index for a document.

    :param doc_path: The path to the document to make or load an index for.
    :param chunk_size: The chunk size to use for the LangChain TokenTextSplitter.
    :param chunk_overlap: The chunk overlap to use for the LangChain TokenTextSplitter.
    :returns: The index.
    """
    file_hash = compute_file_hash(doc_path)
    service_context = make_service_context()
    persist_dir = get_persist_dir(file_hash)

    document = magic_load_doc(doc_path)
    split_docs = split_document(document[0], chunk_size, chunk_overlap)

    # Check that the persist directory exists and that we have made a docstore,
    # which is a sentinel test for the rest of the index.
    if persist_dir.exists() and (persist_dir / "docstore.json").exists():
        index = load_index(persist_dir, service_context=service_context)
    else:
        persist_dir.mkdir(exist_ok=True, parents=True)
        index = make_index(
            split_docs, file_hash, persist_dir, service_context=service_context
        )
    return index
