import openai 
import streamlit as st
import extra_streamlit_components as stx
import uuid
import time

openai.api_key = st.secrets["OPENAI_API_KEY"]
assistant_id = st.secrets["OPENAI_ASSISTANT_ID"]
client = openai

st.set_page_config(page_title="Assistant API Chat", page_icon=":speech_balloon:")

@st.cache_resource(experimental_allow_widgets=True)
def get_manager():
    return stx.CookieManager()
cookie_manager = get_manager()

st.session_state.thread_id = cookie_manager.get('thread_id')

if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())
if "messages" not in st.session_state:
    st.session_state.messages = []

st.title(":speech_balloon: Assistant API Chat")

with st.sidebar:
    st.caption(f"**Session ID**:  \n  {st.session_state.session_id}")
    st.header("Configuration")
    thread_id = st.text_input("Enter your Thread ID:", 
                              placeholder="Leave empty to start a new thread")
    c1, c2 = st.columns(2)

    if c1.button("Start Chat", use_container_width=True):
        if thread_id:
            st.session_state.thread_id = thread_id
        else:
            thread = client.beta.threads.create(
                metadata={
                    'session_id': st.session_state.session_id,
                }
            )
            st.session_state.thread_id = thread.id
        cookie_manager.set('thread_id', st.session_state.thread_id)

    if c2.button("Clear Chat", use_container_width=True):
        cookie_manager.delete('thread_id')
        st.session_state.thread_id = None

#def process_citations(message):
#    message_content = message.content[0].text
#    annotations = message_content.annotations if hasattr(message_content, 'annotations') else []
#    citations = []
#    full_response = message_content.value + '\n\n' + '\n'.join(citations)
#    return full_response

if st.session_state.thread_id:
    st.session_state.messages = client.beta.threads.messages.list(
        thread_id=st.session_state.thread_id
    )
    st.caption(f"**Thread ID**: {st.session_state.thread_id}")
    for message in reversed(st.session_state.messages.data):
        with st.chat_message(message.role):
            st.markdown(message.content[0].text.value)

    if prompt := st.chat_input():
        with st.chat_message("user"):
            st.markdown(prompt)

        client.beta.threads.messages.create(
            thread_id=st.session_state.thread_id,
            role="user",
            content=prompt
        )

        run = client.beta.threads.runs.create(
            thread_id=st.session_state.thread_id,
            assistant_id=assistant_id,
        )

        while run.status != 'completed':
            time.sleep(1)
            run = client.beta.threads.runs.retrieve(
                thread_id=st.session_state.thread_id,
                run_id=run.id
            )

        st.session_state.messages = client.beta.threads.messages.list(
            thread_id=st.session_state.thread_id
        )

        for message in reversed(st.session_state.messages.data):
            if message.run_id == run.id and message.role == "assistant":
                with st.chat_message("assistant"):
                    st.markdown(message.content[0].text.value)
else:
    st.write("\n")
    st.info("Click on 'Start Chat' to start/continue a thread.", icon='⚠')