import time
import warnings

import pandas as pd
import streamlit as st
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory

from src.main import get_response
from src.rule_based import simple_chatbot
from src.main import write_answer
warnings.filterwarnings("ignore")


# Streamlit app setup
st.title("FinBuddy")
st.write("I am an AI Assistant helping with answering questions related to financial metrics of MSFT, TSLA, and Apple.")

msgs = StreamlitChatMessageHistory(key="special_app_key")

for msg in msgs.messages:
    st.chat_message(msg.type).write(msg.content)

if prompt := st.chat_input():
    start_time = time.time()
    st.chat_message("human").write(prompt)
    msgs.add_user_message(prompt)

    with st.spinner("Waiting for response..."):
        # Get response from chatbot
        response_text = get_response(prompt)

        # enable below one to get the plots or table
        #response_text = write_answer(response_text)

        # Rule based bot
        #response_text = simple_chatbot(prompt)

    if response_text:
        if 'answer' in response_text:
            st.chat_message("ai").write(response_text['answer'])
            msgs.add_ai_message(response_text['answer'])
        else:
            for key, value in response_text.items():
                msgs.add_ai_message(str(response_text))
        end_time = time.time()
        st.write(f"Total time {end_time - start_time}")
    else:
        st.error("No valid response received from the AI.")