from langchain.vectorstores import DeepLake
from langchain.llms import BaseLLM
from langchain.docstore.document import Document
from typing import List
from colorama import Fore, Back, Style, init
import json
import time
import openai

init(autoreset=True)

RECURSIVE_PROMPT_LIMIT = 3


class QA:

    def __init__(self, llm: BaseLLM, deeplake_store: DeepLake, num_results: int):
        self.llm = llm
        self.retriever = deeplake_store.as_retriever()
        self.retriever.search_kwargs['distance_metric'] = 'cos'
        self.retriever.search_kwargs["fetch_k"] = 100
        self.retriever.search_kwargs['maximal_marginal_relevance'] = False
        self.num_results = num_results
        self.retriever.search_kwargs['k'] = num_results

    def create_prompt(self, query_str: str, similar_chunks: List[Document]) -> str:
        """Build the final prompt string using query and similar chunks"""
        similar_chunk_str = '\n'.join([chunk.page_content for chunk in similar_chunks])
        # TODO: Try structured json prompt
        final_prompt = f"You will be asked a question based on the following code snippets, \n {similar_chunk_str}\n " \
                       f"You may need to combine the above snippets according to their line numbers to answer the " \
                       f"following question.  The question is: {query_str}"
        return final_prompt

    def get_relevant_documents(self, query_str: str, documents: List[Document]) -> List[Document]:

        chunk_dict = {index: value for index, value in enumerate([chunk.page_content for chunk in documents])}
        chunk_json = json.dumps(chunk_dict)
        recursive_prompt = f"You will be given a json string where the keys are indexes and the values are snippets of " \
                           f"code from a repo with some contextual information above them. You will also be given a " \
                           f"question to answer based on these snippets. The snippets of code were " \
                           f"generated by finding the snippets from the repo with the highest semantic similarity " \
                           f"to the question string.  Therefore even though these snippets are semantically similar to " \
                           f"the question string, they may or may not be useful for answering the question.  Your job is" \
                           f"to decide which snippets are useful for answering the question and which are not useful. " \
                           f"Return a json string where the keys are the snippet indexes and values are true or false " \
                           f"based on whether that snippet is useful for answering the given question string." \
                           f"The json string containing the code snippets is {chunk_json} and the " \
                           f"question string is {query_str}."

        try:
            resp = self.llm(recursive_prompt)
        except openai.error.RateLimitError as e:
            print("sleeping for 10 seconds because of rate limit error")
            time.sleep(10)

        resp_dict = json.loads(resp)

        relevant_chunks = []

        # Extract the relevant documents
        for key in resp_dict:
            if resp_dict[key]:
                relevant_chunks.append(documents[int(key)])

        return relevant_chunks

    def deduplicate_documents(self, documents: List[Document]) -> List[Document]:
        seen_meta = set()
        dedup_list = []

        for doc in documents:
            doc_meta = f"{doc.metadata['source']} - lines {doc.metadata['starting_line']} - {doc.metadata['ending_line']}"
            if doc_meta not in seen_meta:
                dedup_list.append(doc)
                seen_meta.add(doc_meta)

        return dedup_list

    def get_relevance_retrieval(self, query_str: str) -> List[Document]:
        """Recursively generate relevant documents by asking the LLM which similar documents are relevant"""

        # start with getting a list of semantically similar documents
        similar_chunks = self.retriever.get_relevant_documents(query_str)

        # Get documents deemed relevant to question according to LLM
        relevant_documents = self.get_relevant_documents(query_str, similar_chunks)

        # Get most similar documents to relevant documents
        chunks_similar_to_relevant_docs = []
        self.retriever.search_kwargs['k'] = 10
        for doc in relevant_documents:
            similar_to_relevant_chunks = self.retriever.get_relevant_documents(doc.page_content)
            chunks_similar_to_relevant_docs.extend(similar_to_relevant_chunks)

        self.retriever.search_kwargs['k'] = self.num_results
        return self.deduplicate_documents(relevant_documents+chunks_similar_to_relevant_docs)

    def get_resp(self, query_str: str) -> str:
        """Given a string, get similar chunks and construct a prompt feed it to LLM and return response"""
        relevant_chunks = self.get_relevance_retrieval(query_str)
        print(Fore.RED + "Relevant files:")
        for chunk in relevant_chunks:
            print(Fore.RED +
                  f"{chunk.metadata['source']} - lines {chunk.metadata['starting_line']} - {chunk.metadata['ending_line']}")
        qa_prompt = self.create_prompt(query_str,relevant_chunks)
        print("Computing response...")

        time.sleep(5)
        return self.llm(qa_prompt)
