import streamlit as st
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
import os 

os.environ["OPENAI_API_KEY"] = os.getenv("k1")

st.title("Language Processing Application :robot_face:")

# Sidebar for task selection
task = st.sidebar.selectbox("Choose a task:", ["Translation", "Summarization"])

# Tabs for multiple pages
tab1, tab2 = st.tabs(["Page 1", "Page 2"])

with tab1:
    if task == "Translation":
        st.header("Translation Task")
        options1 = ["English", "Telugu", "Hindi", "French", "German", "Russian", "Spanish"]
        input_language = st.selectbox("Input Language: ", options1)

        options2 = ["Hindi", "Telugu", "Spanish", "English", "German", "Russian", "French"]
        output_language = st.selectbox("Output Language: ", options2)

        text = st.text_input("Text Input: ")

        if st.button("Submit Translation"):
            llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

            prompt = ChatPromptTemplate.from_messages([("system", "you are a good assistant for translation from {il} to {ol}"),
                                                       ("human", "{i}")])

            chain = prompt | llm

            response = chain.invoke({"il": input_language, "ol": output_language, "i": text})

            st.write("Response: ")
            st.write(response.content)

    elif task == "Summarization":
        st.header("Summarization Task")
        text = st.text_area("Text Input for Summarization: ")

        if st.button("Submit Summarization"):
            llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

            prompt = ChatPromptTemplate.from_messages([("system", "you are a good assistant for summarization"),
                                                       ("human", "Summarize the following text: {i}")])

            chain = prompt | llm

            response = chain.invoke({"i": text})

            st.write("Response: ")
            st.write(response.content)

with tab2:
    if task == "Translation":
        st.header("Translation Task")
        # Repeat the same structure for additional pages if needed
        options1 = ["English", "Telugu", "Hindi", "French", "German", "Russian", "Spanish"]
        input_language = st.selectbox("Input Language: ", options1, key='input_lang2')

        options2 = ["Hindi", "Telugu", "Spanish", "English", "German", "Russian", "French"]
        output_language = st.selectbox("Output Language: ", options2, key='output_lang2')

        text = st.text_input("Text Input: ", key='text_input2')

        if st.button("Submit Translation", key='submit_translation2'):
            llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

            prompt = ChatPromptTemplate.from_messages([("system", "you are a good assistant for translation from {il} to {ol}"),
                                                       ("human", "{i}")])

            chain = prompt | llm

            response = chain.invoke({"il": input_language, "ol": output_language, "i": text})

            st.write("Response: ")
            st.write(response.content)

    elif task == "Summarization":
        st.header("Summarization Task")
        text = st.text_area("Text Input for Summarization: ", key='text_area2')

        if st.button("Submit Summarization", key='submit_summarization2'):
            llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

            prompt = ChatPromptTemplate.from_messages([("system", "you are a good assistant for summarization"),
                                                       ("human", "Summarize the following text: {i}")])

            chain = prompt | llm

            response = chain.invoke({"i": text})

            st.write("Response: ")
            st.write(response.content)