# Import necessary libraries
import streamlit as st
import os
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain, SequentialChain

# Function to translate and summarize the email
def translate_and_summarize(api_key: str, email: str) -> dict:
    os.environ['OPENAI_API_KEY'] = api_key  # Set the OpenAI API key as an environment variable
    llm = ChatOpenAI()

    # Identify the email's language
    template1 = "Return the language this email is written in:\n{email}.\nONLY return the language it was written in."
    prompt1 = ChatPromptTemplate.from_template(template1)
    chain_1 = LLMChain(llm=llm, prompt=prompt1, output_key="language")

    # Translate the email to English
    template2 = "Translate this email from {language} to English. Here is the email:\n" + email
    prompt2 = ChatPromptTemplate.from_template(template2)
    chain_2 = LLMChain(llm=llm, prompt=prompt2, output_key="translated_email")

    # Provide a summary in English
    template3 = "Create a short summary of this email:\n{translated_email}"
    prompt3 = ChatPromptTemplate.from_template(template3)
    chain_3 = LLMChain(llm=llm, prompt=prompt3, output_key="summary")

    # Provide a reply in English
    template4 = "Reply to the sender of the email giving a plausible reply based on the {summary} and a promise to address issues"
    prompt4 = ChatPromptTemplate.from_template(template4)
    chain_4 = LLMChain(llm=llm, prompt=prompt4, output_key="reply")

    # Provide a translation back to the original language
    template5 = "Translate the {reply} back to the original {language} of the email."
    prompt5 = ChatPromptTemplate.from_template(template5)
    chain_5 = LLMChain(llm=llm, prompt=prompt5, output_key="translated_reply")
    

    seq_chain = SequentialChain(chains=[chain_1, chain_2, chain_3, chain_4, chain_5],
                                input_variables=['email'],
                                output_variables=['language', 'translated_email', 'summary', 'reply', 'translated_reply'],
                                verbose=True)
    return seq_chain(email)

def main():
    st.title("Chain Example: Language, Summary, Translate, Respond & Translate")
    st.write("Translates an email to English, provides a summary, plausible reply and translates back to the sender")

    api_key = st.text_input("Enter your OpenAI API Key:", type="password")  # Set as a password input
    email = st.text_area("Enter the email to translate and summarize:")

    if st.button("Translate"):
        if api_key and email:
            try:
                result = translate_and_summarize(api_key, email)
                st.write(f"**Language:** {result['language']}")
                st.write(f"**Summary:** {result['summary']}")
                st.write(f"**Translated Email:** {result['translated_email']}")
                st.write(f"**Reply in English:** {result['reply']}")
                st.write(f"**Reply in Original Language:** {result['translated_reply']}")
            except Exception as e:
                st.write(f"Error: {e}")
        else:
            st.write("Please provide both the API Key and Email.")

if __name__ == "__main__":
    main()