import os
from dataset_loader import load_dataset

if not os.path.isdir("jlc_chatbot_data"):
    load_dataset()

from jlc_chatbot_data.utils import run_conversation


import streamlit as st
import traceback

st.set_page_config(page_title="Facility Management Chatbot", page_icon=":robot:")


if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "assistant", "content": "How may I help you?"}
    ]

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

if st.session_state.messages[-1]["role"] != "assistant":
    try:
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                answer = run_conversation(st.session_state.messages, prompt)
            st.write(answer.replace("\n", "  \n"))

        message = {"role": "assistant", "content": answer}
        st.session_state.messages.append(message)
    except:
        st.error("An error occurred. Please try again.")
        traceback.print_exc()