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

Link de estudo --> https://youtu.be/cY-0TRj-teI

Link de estudo (adicional) --> https://medium.com/@woyera/how-to-chat-with-your-pdf-using-python-llama-2-41df80c4e674

https://github.com/langchain-ai/langchain/discussions/9158
"""
from langchain.document_loaders import YoutubeLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import ChatVectorDBChain
import os

# Pip install langchain, openai, chromadb, youtube-transcript-api and pytube

#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']

# Carregue a transcrição do youtube:
# SE O SEU VÍDEO DO YOUTUBE FOR MUITO CURTO, VOCÊ PODE RECEBER UM ERRO DO CHROMADB OU SOLICITAR MAIS ELEMENTOS DO QUE EXISTEM NO ÍNDICE
loader = YoutubeLoader.from_youtube_url("https://www.youtube.com/watch?v=Q8eajxcS6dQ&t=21s",
                                        add_video_info=True,
                                        language=["pt", "id"],
                                        translation="pt")

# O chromadb está recebendo mensagens de erro com os metadados do youtube, então estamos usando apenas o texto da transcrição e não os metadados
documents = loader.load()
documents[0].metadata
documents[0].page_content
del documents[0].metadata['publish_date']


# Divida os documentos em pedaços de 1000 caracteres:
text_splitter = RecursiveCharacterTextSplitter(
    # Defina um tamanho de Chunk bem pequeno, apenas para mostrar.
    chunk_size = 1000,
    chunk_overlap  = 20,
    length_function = len,
)
documents = text_splitter.split_documents(documents)

persist_directory = "db"
# Crie um armazenamento de vetores a partir dos documentos
embeddings = OpenAIEmbeddings()

# verifique o diretório persist e crie se ele não existir
if not os.path.exists(persist_directory):
    vectorstore = Chroma.from_documents(documents, embeddings, persist_directory=persist_directory)
else:
    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)

from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)

# Criamos um prompt de bate-papo:
system_template="""Use os seguintes trechos de contexto e histórico de bate-papo para responder às perguntas dos usuários.
                   Se você não sabe a resposta, apenas diga que não sabe, não tente inventar uma resposta.
----------------
{context}

"""
messages = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}")
]

prompt = ChatPromptTemplate.from_messages(messages) # Não estou usando o prompt


# Antigo (Deprecated):
# qa = ChatVectorDBChain.from_llm(ChatOpenAI(temperature=0), vectorstore, qa_prompt=prompt)

from langchain.chains import ConversationalRetrievalChain
qa = ConversationalRetrievalChain.from_llm(ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0),
                                           vectorstore.as_retriever(search_kwargs={'k':2}),
                                           return_source_documents=True,
                                           verbose=False
                                           )
                                           


# Não consegui fazer com que o histórico de bate-papo funcionasse corretamente. provavelmente está confiando demais no contexto ou ficando confuso com muita informação:
chat_history = []

# Usamos um loop while para que possamos continuar fazendo perguntas:
while True:
    query = input("Você: ")
    result = qa({"question": query, "chat_history": chat_history})
    # append to the chat history
    # chat_history.append((query, result["answer"]))
    # chat_history = [(query, result["answer"])]
    # stylize the answer in bold
    print("\033[1m" + "Bot: " + result["answer"] + "\033[0m")
    # print(result)