import streamlit as st
import openai
import pandas as pd
from sqlalchemy import create_engine
from langchain.chat_models import ChatOpenAI
from langchain.utilities.sql_database import SQLDatabase
from langchain.tools.sql.toolkit import SQLDatabaseToolkit

# Set OpenAI API Key
openai.api_key = "sk-O7esHSo2XAWm-GXUGXp7_P9l4qXrQMn0CIGzs34ojLT3BlbkFJeXGSSvywppRTAvyT0zZkmZLZsj5cg7XkAkBTh8ZxoA"

# Database connection
DATABASE_URL = "sqlite:///Sakila.db"  # Replace with your DB path or connection string
engine = create_engine(DATABASE_URL)

# Set up LangChain components
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.5)  # OpenAI's Chat model for LLM
db = SQLDatabase(engine)  # Connect LangChain to the database
toolkit = SQLDatabaseToolkit(llm=llm, db=db)  # Create the SQL toolkit

# Streamlit UI setup
st.title("SQL Data Chatbot with LangChain")
st.write("Ask questions about the data, and I will answer them with both a response and an SQL query.")

# Input field for the user question
user_question = st.text_input("Your question:")

# Process the question if provided
if user_question:
    # Generate the SQL query and answer using the toolkit
    try:
        # Execute the question through the SQL toolkit
        answer = toolkit.query(user_question)

        # Display the generated SQL query and answer
        st.subheader("Generated SQL Query")
        st.code(answer.query, language="sql")

        # Display the generated answer
        st.subheader("Answer")
        st.write(answer.result)

        # Execute the SQL query to get results
        with engine.connect() as conn:
            result_df = pd.read_sql_query(answer.query, conn)

        # Show query results if any
        if not result_df.empty:
            st.write("Query Results:")
            st.write(result_df)
        else:
            st.write("No results found for this query.")
    
    except Exception as e:
        st.write(f"Error processing the query: {e}")