import os
import streamlit as st
from dotenv import load_dotenv
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent
from langchain.agents.agent_toolkits import ZapierToolkit
from langchain.agents import AgentType
from langchain.utilities.zapier import ZapierNLAWrapper
import json
from langchain.callbacks import StreamlitCallbackHandler

load_dotenv("../../.env")


agents = os.getenv("ZAPIER_NLA_API_KEY")
api_key = os.getenv("OPENAI_API_KEY")


if api_key is None:
    api_key = st.text_input("Enter your OpenAI API key", type="password")
    os.environ["OPENAI_API_KEY"] = api_key

if agents is None:
    agents = st.text_input("Enter your Zapier NLA API key", type="password")
    os.environ["ZAPIER_NLA_API_KEY"] = agents


# Initialize necessary objects
llm = ChatOpenAI(streaming=True)
zapier = ZapierNLAWrapper()
toolkit = ZapierToolkit.from_zapier_nla_wrapper(zapier)


# Streamlit UI
st.title("Zap those Zapzz")

agent = initialize_agent(
    toolkit.get_tools(), llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
)


if prompt := st.chat_input():
    st.chat_message("user").write(prompt)
    with st.chat_message("assistant"):
        st_callback = StreamlitCallbackHandler(st.container())
        response = agent.run(prompt, callbacks=[st_callback])
        st.write(response)