from dotenv import load_dotenv
import streamlit as st
import os
import sqlite3
import google.generativeai as genai


load_dotenv()
env_var = os.getenv('GOOGLE_API_KEY')
if env_var is None:
    raise EnvironmentError("GOOGLE_API_KEY not found in environment variables.")
genai.configure(api_key=env_var)

def get_gemini_response(question, prompt):
    try:
        # Initialize the generative model
        model = genai.GenerativeModel("gemini-pro")
        
        # Generate content using the model
        response = model.generate_content(prompt + question)
        
        # Return the response text
        return response.text
    
    except genai.api_core.exceptions.GoogleAPICallError as e:
        print(f"A Google API call error occurred: {e}")
        return None
    except genai.api_core.exceptions.RetryError as e:
        print(f"A retry error occurred: {e}")
        return None
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        return None

# Function to retrieve query from database
def read_sql_query(sql, db):
    try:
        con = sqlite3.connect(db)
        cur = con.cursor()
        cur.execute(sql)
        rows = cur.fetchall()
    except sqlite3.Error as e:
        print(f"An error occurred: {e}")
        rows = []
    finally:
        if con:
            con.close()
        return rows

prompt = """
You are an expert in converting complex English questions to SQL queries!
The SQL database has the name STUDENTS and has the following columns - NAME, CLASS, LEVEL, MARKS.

For example,

Example 1 - How many entries of records are present?
The SQL command will be something like this: SELECT COUNT(*) FROM STUDENTS;

Example 2 - List all student names.
The SQL command will be something like this: SELECT NAME FROM STUDENTS;

Example 3 - Find the number of students in each class.
The SQL command will be something like this: SELECT CLASS, COUNT(*) FROM STUDENTS GROUP BY CLASS;

Example 4 - Retrieve all information for students in level 'LEVEL 100'.
The SQL command will be something like this: SELECT * FROM STUDENTS WHERE LEVEL = 'LEVEL 100';

Example 5 - Get the names of students in class 'AI'.
The SQL command will be something like this: SELECT NAME FROM STUDENTS WHERE CLASS = 'AI';

Example 6 - Count the number of students in level 'LEVEL 300'.
The SQL command will be something like this: SELECT COUNT(*) FROM STUDENTS WHERE LEVEL = 'LEVEL 300';

Example 7 - Find the average marks of students in each class.
The SQL command will be something like this: SELECT CLASS, AVG(MARKS) FROM STUDENTS GROUP BY CLASS;

Example 8 - Retrieve the names and marks of students who have marks greater than 90.
The SQL command will be something like this: SELECT NAME, MARKS FROM STUDENTS WHERE MARKS > 90;

Example 9 - List the names of students in 'Data Science' class and 'Level 100'.
The SQL command will be something like this: SELECT NAME FROM STUDENTS WHERE CLASS = 'Data Science' AND LEVEL = 'Level 100';

Example 10 - Find the student with the highest marks.
The SQL command will be something like this: SELECT NAME, MARKS FROM STUDENTS WHERE MARKS = (SELECT MAX(MARKS) FROM STUDENTS);

also the sql code should not have ``` in beginning or end and sql word in output
"""

st.set_page_config(page_title="I Can Retrieve Any SQL Query")
st.header("Gemini App To Retrieve SQL Data")

question = st.text_input("Input: ", key="input")
submit = st.button("Ask the Question")

# If submit is clicked
if submit:
    response = get_gemini_response(question, prompt)
    if response:
        print(response)
        try:
            response = read_sql_query(response, "students.db")
            st.subheader("The Response is")
            for row in response:
                print(row)
                st.write(row)
        except (sqlite3.OperationalError, sqlite3.IntegrityError) as e:
            st.error(f"A database error occurred: {e}")
        except Exception as e:
            st.error(f"An unexpected error occurred while executing the SQL query: {e}")
    else:
        st.error("Failed to generate a valid SQL query.")