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



st.set_page_config(page_title="Text2SQLRAG", layout="wide")
st.title("Text2SQLRAG: Conversational SQL Generator")

with st.sidebar:
    
    st.markdown("""
    # Example Question
    - Calculate the average invoice amount
    - Find all artists whose names start with 'A'
    - Find all playlists containing tracks composed by Bach
    - Count how many different artists are in each playlist
    - Find the 5 tracks that appear on the most invoices
    """)

text2sql = Text2SQLRAG()
chromadb.api.client.SharedSystemClient.clear_system_cache()

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    role = message.get("role", "assistant")
    with st.chat_message(role):
        st.markdown(message.get("text", ""))
        
        if "reasoning" in message:
            with st.expander("Reasoning", expanded=False):
                st.markdown(message["reasoning"])
        
        if "sql_query" in message:
            with st.expander("SQL Query", expanded=False):
                st.code(message["sql_query"], language="sql")
        if "dataframe" in message:
            st.dataframe(message["dataframe"])

input_text = st.chat_input("Chat with your bot here...")

if input_text:
    with st.chat_message("user"):
        st.markdown(input_text)
    
    st.session_state.messages.append({"role": "user", "text": input_text})
    
    response = text2sql.run(input_text)
    sql_query = response.query
    reasoning = response.reasoning
    df = execute_query_and_return_df(sql_query)
    
    with st.chat_message("assistant"):
        if sql_query:
            with st.expander("SQL Query", expanded=True):
                st.code(sql_query, language="sql")
            
            with st.expander("Reasoning", expanded=True):
                st.markdown(reasoning)
        
        if df is not None and not df.empty:
            st.dataframe(df)
        else:
            st.error("Error executing query or no data returned.")
    
    st.session_state.messages.append({
        "role": "assistant",
        "reasoning": reasoning,
        "sql_query": sql_query,
        "dataframe": df,
    })