"""

"""
from langchain.document_loaders import UnstructuredFileLoader, TextLoader
from langchain.document_loaders import UnstructuredURLLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
# INSTALL CHROMADB pip install chromadb
from langchain.vectorstores import Chroma
from langchain.chains.question_answering import load_qa_chain
from langchain import OpenAI, VectorDBQA
from langchain.prompts import PromptTemplate
import os

# Defina a sua chave OpenAI:
#import openai
import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # read local .env file
#openai.api_key  = os.environ['OPENAI_API_KEY']


# Definimos um diretório persist para o chromaDB para que DB seja armazenada LOCALMENTE:
persist_directory = 'db'
embeddings = OpenAIEmbeddings()

# Verifique se o banco de dados já existe:
if not os.path.exists(persist_directory):
    # load book.txt encode and decode in ascii and write to file to avoid encoding errors
    with open("book.txt", "r",encoding="utf-8") as f:
        book = f.read().encode("ascii", "ignore").decode("ascii")
    with open("book_ascii.txt", "w") as f:
        f.write(book)

    print("Loading book_ascii.txt")
    # loader = UnstructuredFileLoader("book_ascii.txt")


    loader = TextLoader('book_ascii.txt')
    documents = loader.load()

    # print(docs[0].page_content[:])

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    docs = text_splitter.split_documents(documents)
    # print(texts[0].page_content)


    
    print("Embedding book_ascii.txt")
    db = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
else:
    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)


# CUSTOM PROMPT
prompt_template = """Use os seguintes trechos de contexto para responder à pergunta no final,
                     resumindo o contexto. Se você não sabe a resposta, apenas diga que não sabe,
                     não tente inventar uma resposta.

{context}

Pergunta: {question}
Resposta:
"""

PROMPT = PromptTemplate(template=prompt_template,
                        input_variables=["context", "question"]
                       )

chain_type_kwargs = {"prompt": PROMPT}

# Aqui estamos usando um Banco de Vetores como um BANCO DE DADOS e não documentos pesquisados por similaridade, pois isso é feito na Chain:
qa = VectorDBQA.from_chain_type(llm=OpenAI(temperature=0.0),
                                chain_type="stuff",
                                vectorstore=db,
                                return_source_documents=False,
                                chain_type_kwargs=chain_type_kwargs
                               )

from langchain.callbacks import get_openai_callback
while True:
    with get_openai_callback() as cb:
        query = input("Digite sua query: ")
        result = qa(query)
        print(result['result'])
        # print(result['source_documents'])
        print("Tokens usados: ", cb.total_tokens)