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 "start_chat" not in st.session_state:
    st.session_state.start_chat = False
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")

def get_all_cookies():
    '''
    WARNING: This uses unsupported feature of Streamlit
    Returns the cookies as a dictionary of kv pairs
    '''
    from streamlit.web.server.websocket_headers import _get_websocket_headers 
    from urllib.parse import unquote

    headers = _get_websocket_headers()
    if headers is None:
        return {}
    
    if 'Cookie' not in headers:
        return {}
    
    cookie_string = headers['Cookie']
    # A sample cookie string: "K1=V1; K2=V2; K3=V3"
    cookie_kv_pairs = cookie_string.split(';')

    cookie_dict = {}
    for kv in cookie_kv_pairs:
        k_and_v = kv.split('=')
        k = k_and_v[0].strip()
        v = k_and_v[1].strip()
        cookie_dict[k] = unquote(v) #e.g. Convert name%40company.com to name@company.com
    return cookie_dict


with st.sidebar:
    st.caption(f"**Session ID**:  \n  {st.session_state.session_id}")
    st.header("Configuration")
    thread_id = st.sidebar.text_input("Enter your Thread ID:", placeholder="Leave empty to start a new thread")
    c1, c2 = st.columns(2)
    st.write(get_all_cookies())
    if c1.button("Start Chat", use_container_width=True):
        
        st.session_state.start_chat = 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):
        st.session_state.start_chat = False
        st.session_state.thread_id = None
        cookie_manager.delete('thread_id')

#def process_message(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("Click on 'Start Chat' to start a new thread.")