from dotenv import load_dotenv
load_dotenv() # load all the envireonment variables

import streamlit as st
import os
import sqlite3

import google.generativeai as genai 

# Configure genai key
genai.configure(api_key = os.getenv("GOOGLE_API_KEY"))

# Function to load google gemini model
def get_gemini_response(question, prompt):
    model = genai.GenerativeModel("gemini-pro")
    response = model.generate_content([prompt, question])
    return response.text

# Function to retrieve query from database
def read_sql_query(query, db):
    conn = sqlite3.connect(db)
    cur = conn.cursor()
    cur.execute(query)
    rows = cur.fetchall()
    conn.commit()           #Imp
    conn.close()
    for row in rows:
        print(row)
    return rows

# Define your prompt
prompt = [
    """
    You are an expert in converting English questions to SQL query!
    The SQL database has the name STUDENT and has the following columns - NAME, CLASS, SECTIOn \n\n
    For Example: \n
    Example1 - How many entires of records are present?, the SQL command will be something like this 
    SELECT COUNT(*) FROM STUDENT;
    also all the sql code should not have ''' in the begining or end and sql word in output. \n
    Exmaple2 - tell me all the students study in Data Science class, the SQL command will be something like this
    SELECT * FROM STUDENT where CLASS = "Data Science";
    also all the sql code should not have ''' in the begining or end and sql word in output. \n
    """
]

# Strealit app
st.set_page_config(page_title = "I can retreive 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:
    resp = get_gemini_response(question=question , prompt=prompt[0])
    print(resp)
    response = read_sql_query(query = resp, db="student.db")
    st.subheader("The Response is")
    for row in response:
        print(row)
        st.header(row)


# Can you provide a list of students categorized as "first class" if their marks are greater than 60 and 
# categorized as "second class" if their marks are lesser than 60
        
# Give me the 2nd last ranker name in terms of marks