import os

import openai
import streamlit as st
from streamlit_chat import message

st.set_page_config(
    page_title="ChatGPT",
    page_icon=":robot:"
)

st.header("ChatGPT")

if 'generated' not in st.session_state:
    st.session_state['generated'] = []
if 'past' not in st.session_state:
    st.session_state['past'] = []
if 'messages' not in st.session_state:
    st.session_state['messages'] = [
        {"role": "system", "content": "You are a helpful assistant that translates English to Chinese."}]


def query(question):
    st.session_state['messages'].append({"role": "user", "content": question})

    openai.api_key = os.environ['OPENAI_API_KEY']
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        temperature=1.2,
        messages=st.session_state['messages'],
        max_tokens=2000
    )

    answer = response['choices'][0]['message']['content']
    st.session_state['messages'].append({"role": "assistant", "content": answer})
    return answer


with st.form("my_form"):
    # Create a text input for the first field
    input_text = st.text_input("You: ", "Hello, how are you?", key="input")

    # Every form must have a submit button.
    # c1, c2 = st.columns([2, 2])
    submitted = st.form_submit_button("🤖 Submit")

if submitted and input_text:
    output = query(input_text, )
    if output:
        st.session_state.past.append(input_text)
        st.session_state.generated.append(output)

if st.session_state['generated']:

    for i in range(len(st.session_state['generated']) - 1, -1, -1):
        message(st.session_state["generated"][i], key=str(i))
        message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')