import streamlit as st
import openai
import pandas as pd
from sqlalchemy import create_engine
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# Set OpenAI API Key directly here if not set as an environment variable
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_api_key=OPENAI_API_KEY)  # Pass API key here

# Custom prompt to generate SQL query
prompt_template = """
You are an SQL expert. Write an SQL query to answer the following question:
Question: {question}
SQL Query:
"""

prompt = PromptTemplate(input_variables=["question"], template=prompt_template)
llm_chain = LLMChain(llm=llm, prompt=prompt)

# 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 using the custom LLM chain
    try:
        # Generate SQL query with the language model
        sql_query = llm_chain.run({"question": user_question}).strip()

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

        # Execute SQL query to get the results
        with engine.connect() as conn:
            result_df = pd.read_sql_query(sql_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}")