import gradio as gr
from PIL import Image
import os
import sqlite3
import google.generativeai as genai
import time

# Initialize Gemini
genai.configure(api_key="YOUR_GEMINI_API_KEY")
genai_model = genai.GenerativeModel('gemini-pro')

class SQLPromptModel:
    def __init__(self, database):
        self.database = database
        self.conn = sqlite3.connect(self.database)
        
    def fetch_table_schema(self, table_name):
        cursor = self.conn.cursor()
        cursor.execute(f"PRAGMA table_info({table_name})")
        schema = cursor.fetchall()
        return schema if schema else None

    def text2sql_gemini(self, schema, user_prompt, inp_prompt=None):
        table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
        
        prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
        Using valid SQLite, write a response that appropriately completes the request for the provided tables.
        ### Instruction: {user_prompt} ### 
        Input: CREATE TABLE sql_pdf({table_columns});
        ### Response: (Return only generated query based on user_prompt , nothing extra)"""

        if inp_prompt is not None:
            prompt = prompt.replace(user_prompt, inp_prompt + " ")
            
        completion = genai_model.generate_content(prompt)
        generated_query = completion.text
        
        # Extract SQL query
        start_index = generated_query.find("SELECT")
        end_index = generated_query.find(";", start_index) + 1
        
        if start_index != -1 and end_index != 0:
            return generated_query[start_index:end_index]
        return generated_query

    def execute_query(self, query):
        cur = self.conn.cursor()
        cur.execute(query)
        columns = [header[0] for header in cur.description]
        rows = [row for row in cur.fetchall()]
        cur.close()
        self.conn.commit()
        return rows, columns

def execute_sql_query(input_prompt):
    database = r"sql_pdf.db"
    sql_model = SQLPromptModel(database)
    
    user_prompt = "Give complete details of properties in India"
    
    for _ in range(3):  # Retry logic
        try:
            table_schema = sql_model.fetch_table_schema("sql_pdf")
            if table_schema:
                if input_prompt.strip():
                    query = sql_model.text2sql_gemini(table_schema, user_prompt, input_prompt)
                else:
                    query = sql_model.text2sql_gemini(table_schema, user_prompt, user_prompt)
                    
                rows, columns = sql_model.execute_query(query)
                return {"Query": query, "Results": rows, "Columns": columns}
            else:
                return {"error": "Table schema not found."}
        except Exception as e:
            print(f"An error occurred: {e}")
            time.sleep(1)
    return {"error": "Failed to execute query after 3 retries."}

# Load the image
image = Image.open(os.path.join(os.path.abspath(''), "house_excel_sheet.png"))

# Create Gradio interface
with gr.Blocks(title="House Database Query") as demo:
    gr.Markdown("# House Database Query System")
    
    # Display the image
    gr.Image(image)
    
    gr.Markdown("""### The database contains information about different properties including their fundamental details. 
                You can query this database using natural language.""")
    
    with gr.Row():
        # Query input and output
        query_input = gr.Textbox(
            lines=2,
            label="Database Query",
            placeholder="Enter your query or choose from examples below. Default: 'Properties in India'"
        )
    
    with gr.Row():
        # Add submit button
        submit_btn = gr.Button("Submit Query", variant="primary")
    
    with gr.Row():
        query_output = gr.JSON(label="Query Results")
    
    # Connect submit button to the query function
    submit_btn.click(
        fn=execute_sql_query,
        inputs=query_input,
        outputs=query_output
    )
    
    # Example queries
    gr.Examples(
        examples=[
            "Properties in France",
            "Properties greater than an acre",
            "Properties with more than 400 bedrooms"
        ],
        inputs=query_input,
        outputs=query_output,
        fn=execute_sql_query
    )

if __name__ == "__main__":
    demo.launch(share=True)