import streamlit as st
import os
from vanna import VannaBase
import pandas as pd
from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.exc import SQLAlchemyError

# Load environment variables
load_dotenv()

# Initialize Vanna AI
vanna_api_key = os.getenv("VANNA_API_KEY")
if not vanna_api_key:
    st.error("VANNA_API_KEY is not set in the environment variables. Please set it and restart the application.")
    st.stop()

vn = VannaBase(api_key=vanna_api_key)

# Check if DATABASE_URL is set
database_url = os.getenv("DATABASE_URL")
if not database_url:
    st.error("DATABASE_URL is not set in the environment variables. Please set it and restart the application.")
    st.stop()

# Try to connect to the database
try:
    engine = create_engine(database_url)
    with engine.connect() as connection:
        st.success("Successfully connected to the database.")
    vn.connect_to_postgres(database_url)
except SQLAlchemyError as e:
    st.error(f"Failed to connect to the database: {str(e)}")
    st.stop()

st.title("Data Explorer")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Chat input
if prompt := st.chat_input("Ask about your data"):
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display user message
    with st.chat_message("user"):
        st.markdown(prompt)
    
    try:
        # Generate SQL query
        sql_query = vn.generate_sql(prompt)
        
        # Execute SQL query and get results
        df = vn.run_sql(sql_query)
        
        # Display assistant response
        with st.chat_message("assistant"):
            st.markdown(f"Here's the SQL query I generated:\n```sql\n{sql_query}\n```")
            st.markdown("And here are the results:")
            st.dataframe(df)
        
        # Add assistant message to chat history
        st.session_state.messages.append({
            "role": "assistant",
            "content": f"Here's the SQL query I generated:\n```sql\n{sql_query}\n```\n\nAnd here are the results:\n{df.to_markdown()}"
        })
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")

# Sidebar with additional information
st.sidebar.header("About")
st.sidebar.info(
    "This is a data exploration tool using Streamlit and Vanna AI. "
    "Ask questions about your data in natural language, and the app will "
    "generate SQL queries and display the results."
)