"""

    app.py

    Main UI script with streamlit.

    Widget keys:
        'FormSubmitter:intake-Submit': form button key returns value true if clicked (bool) FormSubmitter:{form-key}-{form-submission-key}
"""

import time
import streamlit as st
from agent import Naomi
from data_utils import ChatSession, Contact

contact_options = Contact.__match_args__
intake_form = [
    'name',
    'contact_type',
    'contact',
    'dob',
    'location',
    'intake_submission'
]

# Title of the app
st.title("Chatbot Naomi")

print('Initial Session state', st.session_state)

@st.dialog('Intake form', width='large')
def open_intake_form():
    st.markdown('Fill in your detaisl below to start chat session :)')
    name = st.text_input("Enter your name", key='name')
    contact_col_1, contact_col_2 = st.columns(spec=[0.3, 0.7], vertical_alignment='center')
    contact_col_1.selectbox("Select contact option", contact_options, key='contact_type')
    contact_col_2.text_input('Enter your username', key='contact')
    dob = st.date_input("When is your birthday?", key='dob')
    location = st.text_input('Enter your location', key='location')
    #button = st.button('Submit', use_container_width=True, type='primary')
    # after the button is clicked the page automatically reruns and the workflow starts from the beginning

    if st.button('Submit', use_container_width=True, type='primary', key='intake_submission'):
        print('Session state after submission for user input: ', st.session_state)
        time.sleep(1)
        st.rerun()

@st.fragment()
def open_chat_window(**kwargs):
    # adds to current state (deletes if doesnt)
    st.session_state.update(kwargs)

    if 'naomi' not in st.session_state:
        inputs = dict(
            name=kwargs.get('name', None),
            contact_type=kwargs.get('contact_type', None),
            contact=kwargs.get('contact', None),
            dob=kwargs.get('dob', None),
            location=kwargs.get('location', None)
        )
        st.session_state.naomi = Naomi(**inputs)
    if "messages" not in st.session_state:
        st.session_state.messages = ChatSession()

    st.markdown('Welcome to the chat!')
    msgbox = st.container(height=400, border=False)

    # Display existing chat messages
    for message in st.session_state.messages:
        msgbox.chat_message(message["role"]).write(message['content'])

    if user_input := st.chat_input('Enter your message'):
        # Add user message to chat history
        print(f'State: {st.session_state}\nUser inserted message: {user_input}')

        st.session_state.messages.add_message(role='user', content=user_input)
        # Display user message in chat message container
        msgbox.chat_message("user").write(user_input)
        response = st.session_state.naomi.invoke(st.session_state.messages[-1])
        msgbox.chat_message('assistant').write_stream(st.session_state.naomi.gen(response))
        st.session_state.messages.add_message(role='assistant', content=response)

    undo_button, reset_button = st.columns(2)
    # TODO: Add undo buttn haha
    if undo_button.button('Does Nothing button', use_container_width=True, type='secondary'):
        st.session_state['undo_button'] = True
    if reset_button.button('Reset chat', use_container_width=True, type='primary'):
        st.session_state.clear()
        st.rerun()

def main():
    if 'intake_submission' not in st.session_state:
        if st.button('Start chat session . . .', type='primary', key='open_intake'):
            open_intake_form()
    else:
        if 'end_chat' not in st.session_state:
            st.session_state['name'] = st.session_state['name'].lower().capitalize()
            open_chat_window(**st.session_state)
        else:
            if 'naomi' not in st.session_state:
                st.rerun()
            else:
                st.session_state.naomi.end()

if __name__ == '__main__':
    main()