import streamlit as st
import psycopg2
import  google.generativeai as genai

# Function to connect to PostgreSQL
def connect_to_db():
    try:
        conn = psycopg2.connect(
            dbname="postgres",
            user="postgres",
            password="database",
            host="localhost",
            port="5432"
        )
        return conn
    except Exception as e:
        st.error(f"Error connecting to the database: {e}")
        return None

# Function to generate SQL query from natural language using Google Gemini
def generate_sql_from_text(prompt):
    try:
        # Initialize the Gemini client
        genai.configure(api_key="AIzaSyCzQ1u1RfAOEgYy2RVyFGGGSzEsiKEtWMk")

        model = genai.GenerativeModel('gemini-pro')
        
        # Send the prompt to the Gemini model
        response = model.generate_text(prompt)
        
        # Extract the SQL query from the response
        sql_query = response["generated_text"]
        
        return sql_query
    except Exception as e:
        st.error(f"Error generating SQL query: {e}")
        return None

# Function to execute SQL query and return results
def execute_sql_query(conn, query):
    try:
        cursor = conn.cursor()
        cursor.execute(query)
        results = cursor.fetchall()
        cursor.close()
        return results
    except Exception as e:
        st.error(f"Error executing SQL query: {e}")
        return None

# Streamlit UI
st.title("Natural Language to SQL with Google Gemini")

# Input prompt from the user
user_prompt = st.text_input("Enter your query in natural language:")

# Connect to the PostgreSQL database
conn = connect_to_db()

if conn and user_prompt:
    # Generate SQL query from natural language prompt
    st.write("Generating SQL query using Google Gemini...")
    sql_query = generate_sql_from_text(user_prompt)

    if sql_query:
        st.write(f"Generated SQL Query: {sql_query}")

        # Execute the generated SQL query
        results = execute_sql_query(conn, sql_query)

        # Display the results
        if results:
            st.write("Query Results:")
            st.dataframe(results)

    # Close the database connection
    conn.close()