import streamlit as st
import google.generativeai as genai
from pathlib import Path
import sqlite3
import pandas as pd

st.set_page_config(page_title='SQL QUERY GENERATOR')
st.title('Talk To Your DB with GenAI')

secretKey = "AIzaSyAA_R5VXv1qjJ5jDMObkluREA8BxJO67RU"
#from google.colab import userdata
genai.configure(api_key = secretKey)
# Set up the model
generation_config = {
  "temperature": 0.4,
  "top_p": 1,
  "top_k": 32,
  "max_output_tokens": 4096,
}

safety_settings = [
  {
"category": "HARM_CATEGORY_HARASSMENT",
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
  },
  {
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
  },
  {
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
  },
  {
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
  }
]

@st.cache_resource
def load_model(model1,config1,safety1):
    return genai.GenerativeModel(model_name = model1, generation_config = config1, safety_settings = safety1)

model = load_model("gemini-pro",generation_config,safety_settings)


prompt_parts_1 = [
  "You are an expert in converting English questions to SQL code! The SQL database has the name classicmodels and has the following tables - productlines, products, offices, employees, customers, payments, orders and orderdetails.\n\nFor example,\nExample 1 - How many Classic Cars are present?, the SQL command will be something like this\n SELECT COUNT(*) FROM products WHERE productLine = 'Classic Cars';\n\n\nExample 2 - What are the names of the cars having turnable front wheels?\n\nSELECT productName FROM products WHERE productDescription LIKE '%turnable front wheels%';\n\n\n Example 3 - What are the top 5 high performing products in terms of revenue?, the SQL command will be SELECT productName, SUM(quantityOrdered * priceEach) AS totalRevenue FROM orderdetails JOIN products ON products.productCode = orderdetails.productCode GROUP BY productName ORDER BY totalRevenue DESC LIMIT 5;\n\n\n Example 4 - What are the top 5 employees in terms of sales?, the SQL command will be SELECT e.employeeNumber, e.firstName || ' ' || e.lastName AS employeeName, SUM(od.quantityOrdered * od.priceEach) AS totalSales FROM employees e JOIN customers c ON e.employeeNumber = c.salesRepEmployeeNumber JOIN orders o ON c.customerNumber = o.customerNumber JOIN orderdetails od ON o.orderNumber = od.orderNumber GROUP BY e.employeeNumber, employeeName ORDER BY totalSales DESC LIMIT 5; \n\n\nExample 5 - \n\nSELECT productName FROM products WHERE quantityInStock = (SELECT MAX(quantityInStock) FROM products);\n\n\nExample 4 - \n\nSELECT productName FROM products WHERE quantityInStock = (SELECT MAX(quantityInStock) FROM products);\n\n\nDont include ``` and \\n in the output",
]


st.subheader('SHOW TABLE')
input1=st.text_input("Enter table name")
submit1=st.button("Show")
if input1 is not None and submit1:
    conn = sqlite3.connect('data.sqlite')
    cur = conn.cursor()
    query = f"select * from {input1} limit 5"
    cur.execute(query)
    records = cur.fetchall()
    df1 = pd.read_sql_query(query, con=conn)
    conn.close()
    st.dataframe(df1)

st.subheader("GENERATE SQL RESULT")
question=st.text_input("Enter question related to the database")
submit2=st.button("Run")
if question is not None and submit2:
    prompt_parts = [prompt_parts_1[0], question]
    response = model.generate_content(prompt_parts)
    query1 = response.text
    conn1 = sqlite3.connect('data.sqlite')
    cur1 = conn1.cursor()
    cur1.execute(query1)
    records = cur1.fetchall()
    df2 = pd.read_sql_query(query1, con=conn1)
    conn1.close()
    st.dataframe(df2)