from dotenv import load_dotenv

load_dotenv()
import os
import sqlite3
import google.generativeai as genai
import streamlit as st

genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))


def get_gemini_response(question, prompt):
    model = genai.GenerativeModel("gemini-1.5-flash")
    response = model.generate_content([{"text": prompt}, {"text": question}])
    return response.text


def read_sql_query(sql, db):
    try:
        connection = sqlite3.connect(db)
        cursor = connection.cursor()
        cursor.execute(sql)
        columns = [description[0] for description in cursor.description]
        data = cursor.fetchall()
        connection.close()
        return columns, data
    except Exception as e:
        st.error(f"SQL Error: {e}")
        return [], []


prompt = [
    """
    You are an expert in converting English questions into SQL queries.
    The SQL Database has the name STUDENT and has a table named STUDENT with the following columns:
    NAME, CLASS, SECTION.\n\nExample 1-How many entries of records are present?,the SQL query is: SELECT COUNT(*) FROM STUDENT;
    Example 2-What is the name of the student in class Data Science?,the SQL query is: SELECT NAME FROM STUDENT WHERE CLASS='Data Science';\n\nalso the SQL code should not have ``` in beginning and end and also sql word in output.\n\nNow, convert the question into SQL query.
    """
]

st.set_page_config(
    page_title="SQL Query Generator", page_icon=":guardsman:", layout="wide"
)
st.header("Gemini App to Generate SQL Queries")
question = st.text_input("Enter your question here", key="input")

submit = st.button("Ask the question")

if submit:
    with st.spinner("Generating SQL query..."):
        response = get_gemini_response(question, prompt[0])
        st.write("Generated SQL Query:")
        st.code(response, language="sql")

        columns, result = read_sql_query(response, "student.db")
        if result:
            st.write("Query Results:")
            st.table([dict(zip(columns, row)) for row in result])
            st.success("Query executed successfully!")
        else:
            st.warning("No results found.")