from dotenv import load_dotenv
load_dotenv()

import streamlit as st
import os
import sqlite3
from google import genai

# Initialize Gemini client
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

### Function to convert user input to SQL using Gemini
def get_gemini_sql(user_input):
    # Prepend "Database question" to make intent explicit
    db_input = f"Database question: {user_input}"

    prompt = f"""
You are an expert in converting English questions into SQL queries for an SQLite database.
ONLY generate SQL. Do NOT provide explanations, notes, or general knowledge.
Any response that is not valid SQL should be ignored.

Database table: STUDENT
Columns: NAME, CLASS, SECTION

Examples:
- Question: How many students are studying in Artificial Intelligence class?
  SQL: SELECT NAME FROM STUDENT WHERE CLASS='Artificial Intelligence';
- Question: How many records are in the database?
  SQL: SELECT COUNT(*) FROM STUDENT;

Database Question: {db_input}
SQL Query:
"""

    # Call Gemini
    response = client.models.generate_content(
        model="gemini-2.5-pro",
        contents=prompt
    )

    # Clean up output for SQLite
    sql_query = response.text.strip().replace('"', "'")

    # Validate SQL starts with a common command
    if not sql_query.lower().startswith(("select", "insert", "update", "delete", "count")):
        return None
    return sql_query


### Function to execute SQL query
def read_sql_query(sql, db="students.db"):
    connection = sqlite3.connect(db)
    cur = connection.cursor()
    cur.execute(sql)
    rows = cur.fetchall()
    connection.close()
    return rows


### Streamlit App
st.set_page_config(page_title="Gemini Text-to-SQL")
st.header("Gemini App to Retrieve SQL Data")

user_input = st.text_input("Enter your database question (e.g., 'List students in AI class')")
submit = st.button("Retrieve Data")

if submit:
    # Require meaningful input
    if len(user_input.strip()) < 5:
        st.error("Please enter a clear database-related question, e.g., 'How many students are in Artificial Intelligence?'")
    else:
        sql_query = get_gemini_sql(user_input)

        if sql_query is None:
            st.error("Gemini did not generate a valid SQL query. Please rephrase your question.")
        else:
            st.subheader("Generated SQL Query:")
            st.code(sql_query, language="sql")

            try:
                rows = read_sql_query(sql_query)
                st.subheader("Query Results:")
                if rows:
                    for row in rows:
                        st.write(row)
                else:
                    st.write("No results found.")
            except Exception as e:
                st.error(f"Error executing SQL: {e}")