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

Aqui aprenderemos a construir um aplicativo Web com Streamlit UI e usaremos 
4 funções Python como Tools do LangChain customizadas. 
"""
from langchain import OpenAI
from langchain.agents import initialize_agent, Tool
from langchain.chains.conversation.memory import ConversationBufferMemory

import streamlit as st

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


# Defina uma função para calcular o enésimo número de Fibonacci:
def fib(n):
    if n <= 1:
        return n
    else:
        return(fib(n-1) + fib(n-2))
    
# Defina uma função que classifique a string de entrada em ordem alfabética:
def sort_string(string):
    return ''.join(sorted(string))
    
# Defina uma função para transformar uma palavra em uma palavra criptografada:
def encrypt(word):
    encrypted_word = ""
    for letter in word:
        encrypted_word += chr(ord(letter) + 1)
    return encrypted_word

# Defina uma função para transformar uma palavra em uma palavra descriptografada:
def decrypt(word):
    decrypted_word = ""
    for letter in word:
        decrypted_word += chr(ord(letter) - 1)
    return decrypted_word

# 'return direct' significa que a saída da função será retornada após o máximo de iterações ser atingido?:
tools = [
    Tool(
        name = "Fibonacci",
        func= lambda n: str(fib(int(n))),
        description="Use quando quiser calcular o enésimo número de Fibonacci",
        # return_direct=True
    ),
    Tool(
        name = "Sort String",
        func= lambda string: sort_string(string),
        description="Use quando quiser ordenar uma string em ordem alfabética",
        # return_direct=True
    ),
    Tool(
        name = "Encrypt",
        func= lambda word: encrypt(word),
        description="Use quando quiser criptografar uma palavra",
        # return_direct=True
    ),
    Tool(
        name = "Decrypt",
        func= lambda word: decrypt(word),
        description="Use quando quiser descriptografar uma palavra",
        # return_direct=True
    )
]

if "agent_memory" not in st.session_state:
    st.session_state["agent_memory"] =  ConversationBufferMemory(memory_key="chat_history") # Você pode usar outros tipos de memória também!

llm=OpenAI(temperature=0,
           verbose=True
          )

agent_chain = initialize_agent(tools,  llm, agent="conversational-react-description", memory=st.session_state["agent_memory"], verbose=True) # verbose=True para ver o processo de pensamento do agente


st.header("🤗 :red[Langchain chatbot com agent/tools e memória] 🤗") # short term memory. no embedded vectorstore
# while True:
#     user_input = input("You: ")
#     agent_chain.run(input=user_input)
# initialize the memory buffer
if "memory" not in st.session_state:
    st.session_state["memory"] = ""

user_input = st.text_input("Você: ")
# streamlit button
if st.button("Enviar"):
    
    st.markdown(agent_chain.run(input=user_input))
    # print the memory buffer
    # add conversation history to the memory buffer
    st.session_state["memory"] += st.session_state["agent_memory"].buffer
    print(st.session_state["memory"])
    # write the memory buffer to a file