import streamlit as st
from backend import get_answer, get_clarification_answer, load_documents, prepare_documents
import logging

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

st.title("Chatbot Osservatori")

# Initialize session state if not already set
if 'step' not in st.session_state:
    st.session_state.step = 1
    st.session_state.user_input = ""
    st.session_state.clarification = ""
    st.session_state.clarify_answer = ""
    st.session_state.answer = ""
    st.session_state.sources = ""
    st.session_state.db = None

# Display the appropriate UI based on the step
if st.session_state.step == 1:
    user_input = st.text_input("Chiedici qualcosa:", key='user_input_input')
    if st.button("Invia", key='button_step1'):
        if user_input:
            st.session_state.user_input = user_input
            with st.spinner('Interrogando il modello...'):
                st.session_state.clarify_answer = get_clarification_answer(user_input)
            st.session_state.step = 2

if st.session_state.step == 2:
    st.write(st.session_state.clarify_answer)
    clarification_input = st.text_input("", key='clarification_input', value=st.session_state.clarification)
    
    if st.button("Invia", key='button_step2'):
        st.session_state.clarification = clarification_input
        documents = None
        with st.spinner('Caricando i documenti rilevanti...'):
            if "blockchain" in clarification_input.lower(): 
                documents = load_documents('documents/blockchain')
            elif "metaverse" in clarification_input.lower():
                documents = load_documents('documents/metaverso')
            elif "payment" in clarification_input.lower():
                documents = load_documents('documents/payment')
            else:
                st.write('Per favore, usa il nome corretto degli osservatori che vuoi interrogare. Opzioni valide: "Blockchain", "Metaverse", "Payment".')
                documents = None
        
        if documents:
            with st.spinner('Preparando i documenti...'):
                st.session_state.db = prepare_documents(documents)
            with st.spinner('Interrogando il modello...'):    
                st.session_state.answer, st.session_state.sources = get_answer(st.session_state.user_input, st.session_state.db)
            st.session_state.step = 3

if st.session_state.step == 3:
    st.write(st.session_state.answer)
    st.write("Fonti: " + st.session_state.sources)
    
    new_question = st.text_input("Fai un'altra domanda:", key='new_question_input')
    if st.button("Invia", key='button_step3'):
        if new_question:
            st.session_state.user_input = new_question
            with st.spinner('Interrogando il modello...'):
                st.session_state.answer, st.session_state.sources = get_answer(new_question, st.session_state.db)
            st.rerun()    

    if st.button("Reset", key='button_reset'):
        st.session_state.step = 1
        st.session_state.user_input = ""
        st.session_state.clarification = ""
        st.session_state.clarify_answer = ""
        st.session_state.answer = ""
        st.session_state.sources = ""
        st.session_state.db = None