#!/usr/bin/env python3
# pip install -U langchain langchain-community fastembed chromadb langchain-groq


from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from langchain_groq import ChatGroq
from langchain_community.document_loaders import TextLoader


import cfg


# Create vector database
def create_vector_database():

    loader = TextLoader('2.txt', encoding='utf-8')
    documents = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=100)
    docs = text_splitter.split_documents(documents)

    #len(docs)
    print(f"length of documents loaded: {len(documents)}")
    print(f"total number of document chunks generated :{len(docs)}")
    #docs[0]

    # Initialize Embeddings
    embed_model = FastEmbedEmbeddings(model_name="BAAI/bge-base-en-v1.5")

    # Create and persist a Chroma vector database from the chunked documents
    # vs = Chroma.from_documents(
    #     documents=docs,
    #     embedding=embed_model,
    #     persist_directory="chroma_db_llamaparse1",
    #     collection_name="rag"
    # )

    # load saved vector database from disk
    vs = Chroma(persist_directory="chroma_db_llamaparse1",
                embedding_function=embed_model,
                collection_name="rag")

    print('Vector DB created successfully !')

    return vs,embed_model


def set_custom_prompt():
    """
    Prompt template for QA retrieval for each vectorstore
    """

    custom_prompt_template = """Use the following pieces of information to answer the user's question.
If you don't know the answer, just say that you don't know, don't try to make up an answer.

Context: {context}
Question: {question}

Only return the helpful answer below and nothing else.
Only answer in russian language.
Helpful answer:
"""

    prompt = PromptTemplate(template=custom_prompt_template,
                            input_variables=['context', 'question'])
    return prompt


if __name__ == '__main__':
    groq_api_key = cfg.GROQ_API_KEY[0]

    vs,embed_model = create_vector_database()

    chat_model = ChatGroq(temperature=0,
                          model_name="mixtral-8x7b-32768",
                          api_key=groq_api_key,)

    vectorstore = Chroma(embedding_function=embed_model,
                         persist_directory="chroma_db_llamaparse1",
                         collection_name="rag")
    #
    retriever=vectorstore.as_retriever(search_kwargs={'k': 3})

    prompt = set_custom_prompt()

    qa = RetrievalQA.from_chain_type(llm=chat_model,
                                     chain_type="stuff",
                                     retriever=retriever,
                                     return_source_documents=True,
                                     chain_type_kwargs={"prompt": prompt})

    while 1:
        q = input('> ')
        response = qa.invoke({"query": q})
        print(response['result'])