import streamlit as st
import os
from openai import OpenAI
import time

st.title("Math Tutor Chatbot")

def main1(client, thread_id, prompt):
    assistantId = "asst_pCMn8HOHVthGSn30KEAtbee6"
    
    message = client.beta.threads.messages.create(
        thread_id=thread_id,
        role="user",
        content=prompt
    )
    
    run = client.beta.threads.runs.create(
        thread_id=thread_id,
        assistant_id=assistantId,
        instructions="Please address the user as Rumaisa."
    )
    
    run = client.beta.threads.runs.retrieve(
        thread_id=thread_id,
        run_id=run.id
    )
    
    while run.status != "completed":
        time.sleep(1)
        run = client.beta.threads.runs.retrieve(
            thread_id=thread_id,
            run_id=run.id
        )
                
    messages = client.beta.threads.messages.list(
        thread_id=thread_id
    )
    
    arr = {}
    for message in reversed(messages.data):
        arr[message.role] = message.content[0].text.value
    return arr

def main():
    session_id = st.session_state.get('session_id', None)
    if session_id is None:
        session_id = os.urandom(16).hex()
        st.session_state['session_id'] = session_id

    if 'client' not in st.session_state:
        st.session_state['client'] = None

    with st.form(key='my_form'):
        user_input = st.text_input(label='You:', key='user_input')
        submit_button = st.form_submit_button(label='Send')

    with st.sidebar:
        st.subheader("Add your OPENAI API")
        api = st.text_input("API:", type='password')

    if not api:
        st.warning("Please enter API to continue")
    else:
        if 'client' not in st.session_state or st.session_state['client'] is None:
            st.session_state['client'] = OpenAI(api_key=api)

        client = st.session_state['client']

        if submit_button and user_input:
            if 'thread_id' not in st.session_state:
                thread = client.beta.threads.create()
                st.session_state['thread_id'] = thread.id
                thread_id = thread.id  # Assign thread_id here
            else:
                thread_id = st.session_state['thread_id']
        
            response = main1(client, thread_id, user_input)
            st.write('Chatbot:', response.get('assistant', ''))

if __name__ == '__main__':
    main()