import openai
import gradio as gr
import pandas as pd
import sqlite3
import os
openai.api_key = os.environ["Secret"]

#OpenAi call
def gpt3(texts):
    response = openai.Completion.create(
      engine="code-davinci-002",
      prompt= texts,
          temperature=0,
          max_tokens=750,
          top_p=1,
          frequency_penalty=0.0,
          presence_penalty=0.0,
          stop = (";", "/*", "</code>")
    )
    x = response.choices[0].text 
    
    return x

# Function to elicit sql response from model



# Function to elicit sql response from model


def greet(prompt, file = None):

    #get the file path from the file object
    file_path = file.name

    # read the file and get the column names
    if file_path:
        if file_path.endswith(".csv"):
            df = pd.read_csv(file_path)
            columns = " ".join(df.columns)
            



        elif file_path.endswith((".xls", ".xlsx")):
            df = pd.read_excel(file_path)
            columns = " ".join(df.columns)
        else:
            return "Invalid file type. Please provide a CSV or Excel file."
        
        # create a SQLite database in memory
        con = sqlite3.connect(":memory:")
        # extract the table name so it can be used in the SQL query
        # in order to get the table name, we need to remove the file extension

        table_name = os.path.splitext(os.path.basename(file_path.name))[0]




        



        # write the DataFrame to a SQL table



        df.to_sql(table_name, con)
    else:
        return "Please upload a file."
    txt= (f'''/*Prompt: {prompt}\nColumns: {columns}\nTable: {table_name}*/ \n —-SQL Code:\n''')
    sql = gpt3(txt)
    
    
    # execute the SQL query
    if con:
        df = pd.read_sql_query(sql, con)
        return sql, df
    else:
        return sql, None
        






#Code to set up Gradio UI
iface = gr.Interface(greet, 
                     inputs = ["text", ("file")], 
                     outputs = ["text",gr.Dataframe(type="pandas")],
                     title="Natural Language to SQL", 
                     description="Enter any prompt and get a SQL statement back! For better results, give it more context")
iface.launch()