"""

    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

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

print('Initial Session state', st.session_state)
contact_options = ['Instagram', 'Email', 'Number']
intake_form = [
    'candidate_name',
    'candidate_contact_type',
    'candidate_contact',
    'candidate_dob',
    'candidate_location',
    'intake_submission'
]
if "messages" not in st.session_state:
    st.session_state.messages = []

@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='candidate_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='candidate_contact_type')
    contact_col_2.text_input('Enter your username', key='candidate_contact')
    dob = st.date_input("When is your birthday?", key='candidate_dob')
    location = st.text_input('Enter your location', key='candidate_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)
    naomi = Naomi(**kwargs)

    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.append({"role": "user", "content": user_input})
        # Display user message in chat message container
        msgbox.chat_message("user").write(user_input)
        response = msgbox.chat_message('assistant').write_stream(naomi.respond(st.session_state.messages))

        # Append assistant's response to the messages history
        st.session_state.messages.append({"role": "assistant", "content": response})

    undo_button, reset_button = st.columns(2)
    if undo_button.button('Undo message', 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:
        st.session_state['candidate_name'] = st.session_state['candidate_name'].lower().capitalize()
        open_chat_window(**st.session_state)

        # render_agent_mood()

if __name__ == '__main__':
    main()