"""
Data Scientist.: Dr.Eddy Giusepe Chirinos Isidro

Neste script podemos fazer perguntas a diversos documentos. Onde a pergunta 
do usuário sempre deverá mencionar à "Entidade" (o nome do documento) para assim
conseguir interagir com dito documento. 

É um bom começo usar a recuperação de documentos.... mas estaria faltando levar 
em conta o contexto (o hisórico de bate-papo) para assim ter uma melhor resposta
para o usuário.

Nota:
* Ao executar o script será criado 4 diretórios com os nomes do arquivos, os quais 
são armazenado no DB Chroma (cada nome é um índice). 
"""
import os
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.chains.router import MultiRetrievalQAChain
from langchain.llms import OpenAI
from langchain.document_loaders import PyPDFLoader

# Substitua sua chave de API 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']

# Diretório onde os arquivos .txt estão armazenados
docs_dir = 'docs'

# Inicialize os embeddings OpenAI:
embedding = OpenAIEmbeddings()

retrievers = []
retriever_descriptions = []
retriever_names = []

# Inicialize o splitter de texto:
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 500,
    chunk_overlap  = 20,
    length_function = len,
)

# Iterar todos os arquivos .txt e .pdf no diretório:
for filename in os.listdir(docs_dir):
    doc = None
    # Verifique se já existe um Chroma VectorStore persistente para este documento:
    if os.path.exists(filename[:-4]): # "slicing" em Python. Cortando os últimos 4 caracteres do nome do arquivo (".txt", ".jpg" ou ".doc"), para obter o nome do arquivo sem a extensão.
        # Se existir, carregue-o do disco:
        retriever = Chroma(persist_directory=filename[:-4], embedding_function=embedding).as_retriever()
    else:
        # Carregue o documento e divida-o:
        if filename.endswith('.txt'):
            try:
                with open(os.path.join(docs_dir, filename), 'r', encoding='utf-8') as f:
                    doc = f.read()
            except UnicodeDecodeError:
                # Tratar possíveis erros de codificação
                print(f"Ignorando arquivo {filename} devido a erros de codificação.")
                continue
            # Se for um .txt, dividimos o documento:
            doc = text_splitter.create_documents([doc])
        elif filename.endswith('.pdf'):
            loader = PyPDFLoader(os.path.join(docs_dir, filename))
            doc = loader.load_and_split()
            print(doc)
        
        if doc is not None:
            # Crie um novo Chroma VectorStore e salve-o no disco:
            retriever = Chroma.from_documents(documents=doc, embedding=embedding, persist_directory=filename[:-4])
            retriever.persist()
            retriever = retriever.as_retriever()
        
    # Adicione o retriever, seu nome e sua descrição às respectivas listas:
    retrievers.append(retriever)
    retriever_names.append(filename[:-4])
    # PRESTE ATENÇÃO AOS NOMES DOS ARQUIVOS COMO ESTARÃO NAS DESCRIÇÕES:
    retriever_descriptions.append(f"Bom para responder perguntas sobre {filename[:-4]}")


# Inicialize o MultiRetrievalQAChain 🤗:
chain = MultiRetrievalQAChain.from_retrievers(OpenAI(), retriever_names, retriever_descriptions, retrievers, verbose=False)

# Teste-o:
# print(chain.run("Quais são as diferenças entre Newton e Feynman?"))

while True:
    print(chain.run(input("\033[033mO que você gostaria de saber? 🤓\033[m ")))