import os
import sys 

module_path = ".."
sys.path.append(os.path.abspath(module_path))

import langchain
from langchain.document_loaders import ConfluenceLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from langchain.embeddings import BedrockEmbeddings
from langchain.llms.bedrock import Bedrock
from langchain.vectorstores import FAISS
from langchain.indexes.vectorstore import VectorStoreIndexWrapper
from utils import bedrock



boto3_bedrock = bedrock.get_bedrock_client(
    assumed_role=os.environ.get("BEDROCK_ASSUME_ROLE", None),
    endpoint_url=os.environ.get("BEDROCK_ENDPOINT_URL", None),
    region=os.environ.get("AWS_DEFAULT_REGION", None),
)


class BedrockConfluenceQA:
    def __init__(self, config: dict = {}):
        self.config = config
        self.embedding = None
        self.vectordb = None
        self.llm = None
        self.qa = None
        self.retriever = None
        self.model_id = None

    def init_embeddings(self) -> None:
        # AWS Bedrock Embeddings
        self.embedding = BedrockEmbeddings(client=boto3_bedrock)

    def init_models(self, parameters: dict = {}) -> None:
        self.parameters = parameters
        max_token_count = self.parameters.get("max_token_count", 512)
        temprature = self.parameters.get("temprature", 1)
        top_p = self.parameters.get("top_p", 1)
        top_k = self.parameters.get("top_k", 1)
        model_id = self.parameters.get("model_id", "amazon.titan-tg1-large")
        self.model_id = model_id
        # AWS Bedrock titan
        if "claude" in model_id:
            self.llm = Bedrock(
                model_id=model_id,
                client=boto3_bedrock,
                model_kwargs={
                    "max_tokens_to_sample":max_token_count,
                    "temperature": temprature,
                    "top_k": top_k,
                    "top_p": top_p,
                }
            )
        if "titan" in model_id:
            self.llm = Bedrock(model_id=model_id, client=boto3_bedrock, model_kwargs= {
                "maxTokenCount": max_token_count,
                "temperature": temprature,
                "topP": top_p,
            })
        if "ai21" in model_id:
            self.llm = Bedrock(model_id=model_id, client=boto3_bedrock, model_kwargs= {
                "maxTokens": max_token_count,
                "temperature": temprature,
                "topP": top_p,
            })


    def vector_db_confluence_docs(self, force_reload: bool = False) -> None:
        """
        creates vector db for the embeddings and persists them or loads a vector db from the persist directory
        """
        persist_directory = self.config.get("persist_directory", None)
        confluence_url = self.config.get("confluence_url", None)
        username = self.config.get("username", None)
        api_key = self.config.get("api_key", None)
        space_key = self.config.get("space_key", None)
        if persist_directory and os.path.exists(persist_directory) and not force_reload:
            ## Load from the persist db
            self.vectordb = FAISS.load_local("faiss_index", embeddings=self.embedding)
        else:
            loader = ConfluenceLoader(
                url=confluence_url, username=username, api_key=api_key
            )
            documents = loader.load(space_key=space_key, limit=50)
            ## 2. Split the texts
            text_splitter = RecursiveCharacterTextSplitter(
                # Set a really small chunk size, just to show.
                # Make sure the chunk size does not exceed titan text embeddings max tokens (512)
                chunk_size=1000,
                chunk_overlap=100,
                # separators=["\n", "\n\n"]

            )
            docs = text_splitter.split_documents(documents)
            print(len(docs))

            ## 3. Create Embeddings and add to chroma store
            ##TODO: Validate if self.embedding is not None
            vectorstore_faiss = FAISS.from_documents(
                docs,
                self.embedding,
            )
            VectorStoreIndexWrapper(vectorstore=vectorstore_faiss)
            self.vectordb = vectorstore_faiss
            # vectorstore_faiss_aws.save_local("faiss_index")

    def retreival_qa_chain(self):
        """
        Creates retrieval qa chain using vectordb as retrivar and LLM to complete the prompt
        """
        ##TODO: Use custom prompt
        self.retriever = self.vectordb.as_retriever(search_kwargs={"k": 10})
        # self.qa = RetrievalQA.from_chain_type(llm=self.llm, chain_type="stuff",retriever=self.retriever)
#         prompt_template = """Human: Use the following pieces of context to provide a concise answer to the question at the end. If you don't know the answer, just  say that you don't know, don't try to make up an answer.
#         {context}

#         Question: {question}
#         Assistant:"""

#         prompt_template = """Human: Please use the context below to craft a succinct response to the following question. If you don't have the information, it's okay to state that you're unaware instead of inventing an answer.
#         {context}

#         Question: {question}
#         Assistant:"""

        prompt_template = """Human: Utilize the context provided to formulate a comprehensive response to the following question. If you're uncertain about the answer, it's perfectly fine to acknowledge that you're unsure rather than providing speculative information.
            {context}

            Question: {question}
            Assistant:"""
    
        ## used for the bulk answers generation
        prompt_template = """# INSTRUCTION
            Answer any question about onboarding or company-related topics at LogicWorks acting as a onboarding manager. If you don't have the information, it's okay to state that you're unaware instead of inventing an answer.
            Utilize the context provided to formulate a comprehensive response to the following question. If you don't have the information, it's okay to state that you're unaware instead of inventing an answer.

            # CONTEXT
            {context}

            # QUESTION
            {question}

            Assistant:
        """


        prompt_template = """User: Answer the question based only on the information provided between ##. If you don't know the answer, just  say that you don't know, don't try to make up an answer.
        #
        {context}
        #

        Question: {question}
        Assistant:"""

        PROMPT = PromptTemplate(
            template=prompt_template, input_variables=["context", "question"]
        )
    
        self.qa = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.retriever,
            return_source_documents=True,
            chain_type_kwargs={"prompt": PROMPT},
        )

    def answer_confluence(self, question: str) -> str:
        """
        Answer the question
        """
        answer = self.qa({"query": question})
        return answer

langchain.embeddings.BedrockEmbeddings langchain.vectorstores.FAISS.load_local langchain.chains.RetrievalQA.from_chain_type langchain.text_splitter.RecursiveCharacterTextSplitter langchain.indexes.vectorstore.VectorStoreIndexWrapper langchain.document_loaders.ConfluenceLoader langchain.vectorstores.FAISS.from_documents langchain.llms.bedrock.Bedrock langchain.prompts.PromptTemplate
[((58, 86), 'os.path.abspath', 'os.path.abspath', (['module_path'], {}), '(module_path)\n', (73, 86), False, 'import os\n'), ((606, 649), 'os.environ.get', 'os.environ.get', (['"""BEDROCK_ASSUME_ROLE"""', 'None'], {}), "('BEDROCK_ASSUME_ROLE', None)\n", (620, 649), False, 'import os\n'), ((668, 712), 'os.environ.get', 'os.environ.get', (['"""BEDROCK_ENDPOINT_URL"""', 'None'], {}), "('BEDROCK_ENDPOINT_URL', None)\n", (682, 712), False, 'import os\n'), ((725, 767), 'os.environ.get', 'os.environ.get', (['"""AWS_DEFAULT_REGION"""', 'None'], {}), "('AWS_DEFAULT_REGION', None)\n", (739, 767), False, 'import os\n'), ((1135, 1174), 'langchain.embeddings.BedrockEmbeddings', 'BedrockEmbeddings', ([], {'client': 'boto3_bedrock'}), '(client=boto3_bedrock)\n', (1152, 1174), False, 'from langchain.embeddings import BedrockEmbeddings\n'), ((6643, 6728), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'template': 'prompt_template', 'input_variables': "['context', 'question']"}), "(template=prompt_template, input_variables=['context',\n    'question'])\n", (6657, 6728), False, 'from langchain.prompts import PromptTemplate\n'), ((6770, 6935), 'langchain.chains.RetrievalQA.from_chain_type', 'RetrievalQA.from_chain_type', ([], {'llm': 'self.llm', 'chain_type': '"""stuff"""', 'retriever': 'self.retriever', 'return_source_documents': '(True)', 'chain_type_kwargs': "{'prompt': PROMPT}"}), "(llm=self.llm, chain_type='stuff', retriever=\n    self.retriever, return_source_documents=True, chain_type_kwargs={\n    'prompt': PROMPT})\n", (6797, 6935), False, 'from langchain.chains import RetrievalQA\n'), ((1689, 1861), 'langchain.llms.bedrock.Bedrock', 'Bedrock', ([], {'model_id': 'model_id', 'client': 'boto3_bedrock', 'model_kwargs': "{'max_tokens_to_sample': max_token_count, 'temperature': temprature,\n    'top_k': top_k, 'top_p': top_p}"}), "(model_id=model_id, client=boto3_bedrock, model_kwargs={\n    'max_tokens_to_sample': max_token_count, 'temperature': temprature,\n    'top_k': top_k, 'top_p': top_p})\n", (1696, 1861), False, 'from langchain.llms.bedrock import Bedrock\n'), ((2068, 2217), 'langchain.llms.bedrock.Bedrock', 'Bedrock', ([], {'model_id': 'model_id', 'client': 'boto3_bedrock', 'model_kwargs': "{'maxTokenCount': max_token_count, 'temperature': temprature, 'topP': top_p}"}), "(model_id=model_id, client=boto3_bedrock, model_kwargs={\n    'maxTokenCount': max_token_count, 'temperature': temprature, 'topP': top_p}\n    )\n", (2075, 2217), False, 'from langchain.llms.bedrock import Bedrock\n'), ((2326, 2465), 'langchain.llms.bedrock.Bedrock', 'Bedrock', ([], {'model_id': 'model_id', 'client': 'boto3_bedrock', 'model_kwargs': "{'maxTokens': max_token_count, 'temperature': temprature, 'topP': top_p}"}), "(model_id=model_id, client=boto3_bedrock, model_kwargs={'maxTokens':\n    max_token_count, 'temperature': temprature, 'topP': top_p})\n", (2333, 2465), False, 'from langchain.llms.bedrock import Bedrock\n'), ((3068, 3101), 'os.path.exists', 'os.path.exists', (['persist_directory'], {}), '(persist_directory)\n', (3082, 3101), False, 'import os\n'), ((3192, 3250), 'langchain.vectorstores.FAISS.load_local', 'FAISS.load_local', (['"""faiss_index"""'], {'embeddings': 'self.embedding'}), "('faiss_index', embeddings=self.embedding)\n", (3208, 3250), False, 'from langchain.vectorstores import FAISS\n'), ((3286, 3358), 'langchain.document_loaders.ConfluenceLoader', 'ConfluenceLoader', ([], {'url': 'confluence_url', 'username': 'username', 'api_key': 'api_key'}), '(url=confluence_url, username=username, api_key=api_key)\n', (3302, 3358), False, 'from langchain.document_loaders import ConfluenceLoader\n'), ((3518, 3584), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(100)'}), '(chunk_size=1000, chunk_overlap=100)\n', (3548, 3584), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((4079, 4121), 'langchain.vectorstores.FAISS.from_documents', 'FAISS.from_documents', (['docs', 'self.embedding'], {}), '(docs, self.embedding)\n', (4099, 4121), False, 'from langchain.vectorstores import FAISS\n'), ((4181, 4235), 'langchain.indexes.vectorstore.VectorStoreIndexWrapper', 'VectorStoreIndexWrapper', ([], {'vectorstore': 'vectorstore_faiss'}), '(vectorstore=vectorstore_faiss)\n', (4204, 4235), False, 'from langchain.indexes.vectorstore import VectorStoreIndexWrapper\n')]