import streamlit as st
from lib import Text2SQLRAG
from utils import execute_query_and_return_df


st.set_page_config(page_title="Text2SQLRAG")
st.title("Text2SQLRAG")

# Create an instance of Text2SQLRAG
text2sql = Text2SQLRAG()

# Initialize session state for storing chat messages
if "messages" not in st.session_state:
    st.session_state.messages = []
    
    # Display conversation history from session state
for message in st.session_state.messages:
    role = message.get("role", "assistant")
    with st.chat_message(role):
        if "output" in message:
            with st.expander("Reasoning", expanded=True):
                st.markdown(message["reasoning"])
        if "sql_query" in message and message["sql_query"]:
            with st.expander("SQL Query", expanded=True):
                st.code(message["sql_query"])

# Get user input
input_text = st.chat_input("Chat with your bot here...")

if input_text:
    # Display user input
    with st.chat_message("user"):
        st.markdown(input_text)

    # Add user input to chat history
    st.session_state.chat_history.append({"role": "user", "text": input_text})

    # Get chatbot response
    response = text2sql.run(input_text)
    sql_query = response.query
    reasoning = response.reasoning
    df = execute_query_and_return_df(sql_query)

    if sql_query:
        with st.expander("SQL Query", expanded=True):
            st.code(sql_query)

        with st.expander("Reasoning", expanded=True):
            st.write(reasoning)

    if df is not None:
        st.dataframe(df)
    else:
        st.error("Error executing query")
    
    # Append assistant response to session state
    st.session_state.messages.append(
        {
            "role": "assistant",
            "reasoning": reasoning,
            "sql_query": sql_query,
        }
    )