import spaces
import gradio as gr
from transformers import pipeline
import torch

DEVICE = 0 if torch.cuda.is_available() else -1

generator = pipeline(
    "text-generation",
    model="sirunchained/text-to-sql-model-v2",
    device=DEVICE,
)

@spaces.GPU
@torch.inference_mode()
def predict_sql(text_input, schema_input=""):
    # Construct the user content including schema if provided
    if schema_input and schema_input.strip():
        user_content = f"# Schema\n{schema_input}\n\n# Text\n{text_input}"
    else:
        user_content = f"# Text\n{text_input}"

    chat_template_input = [
        {"role": "user", "content": user_content}
    ]

    # Apply chat template and generate SQL
    formatted_input = generator.tokenizer.apply_chat_template(
        chat_template_input,
        tokenize=False,
        add_generation_prompt=True,
    )

    # Generate output using the fine-tuned model
    output = generator(
        formatted_input,
        max_new_tokens=256,
    )

    # Extract the generated SQL query
    generated_text = output[0]["generated_text"]

    # Remove the prompt from the generated output
    sql_query = generated_text[len(formatted_input):].strip()

    # Further clean up any unwanted tokens like '<end_of_turn>'
    sql_query = sql_query.replace("<end_of_turn>", "").strip()

    return sql_query

# Sample queries for the Gradio interface including SELECT, INSERT, UPDATE, and DELETE operations
samples = [
    # SELECT queries
    ["Count how many rooms have a daily rate higher than 300.", "rooms(id, daily_rate)"],
    ["List all products that have been ordered in the month of June 2026.",
     "customers(id, name, email, country, registration_date), products(id, name, category, price, stock_quantity), orders(id, customer_id, order_date, total_amount, status), order_items(id, order_id, product_id, quantity, unit_price), reviews(id, customer_id, product_id, rating, review_text, review_date)"],
    ["Find customers who have placed more than 2 orders.", "customers(id, name), orders(id, customer_id)"],
    ["What is the average price of products in the 'Electronics' category?", "products(id, name, category, price)"],
    
    # INSERT queries
    ["Insert a new customer named John Doe with email john@example.com from USA.", 
     "customers(id, name, email, country, registration_date)"],
    ["Add a new product called 'Wireless Mouse' in 'Electronics' category with price 29.99 and stock 150.", 
     "products(id, name, category, price, stock_quantity)"],
    ["Insert a new order for customer_id 5 with total amount 199.99 and status 'pending'.", 
     "orders(id, customer_id, order_date, total_amount, status)"],
    
    # UPDATE queries
    ["Update the price of product 'Laptop' to 899.99.", 
     "products(id, name, category, price, stock_quantity)"],
    ["Change the status of order with id 123 to 'shipped'.", 
     "orders(id, customer_id, order_date, total_amount, status)"],
    ["Update the email of customer named 'Alice Smith' to alice.new@email.com.", 
     "customers(id, name, email, country)"],
    ["Increase the salary of all employees in the 'Sales' department by 10%.", 
     "employees(id, name, department, salary)"],
    
    # DELETE queries
    ["Delete all products that have 0 stock quantity.", 
     "products(id, name, category, price, stock_quantity)"],
    ["Remove customer with email 'inactive@example.com'.", 
     "customers(id, name, email, country)"],
    ["Delete all orders that are older than 2 years and have status 'cancelled'.", 
     "orders(id, customer_id, order_date, total_amount, status)"],
    ["Remove all reviews with rating less than 2 stars.", 
     "reviews(id, customer_id, product_id, rating, review_text, review_date)"],
]

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_sql,
    inputs=[
        gr.Textbox(
            lines=2,
            placeholder="Enter your natural language query here...",
            label="Natural Language Query",
        ),
        gr.Textbox(
            lines=3,
            placeholder="Optional: Enter your schema here (e.g., customers(id, name, email), orders(id, customer_id, order_date))",
            label="Database Schema",
        ),
    ],
    outputs="text",
    title="Text-to-SQL Model Demo",
    description=(
        "Enter a natural language query and optionally provide the database schema. "
        "The model will generate a SQL query. Supports SELECT, INSERT, UPDATE, and DELETE operations.\n"
    ),
    examples=samples,
)

iface.queue()

# Launch the interface
iface.launch()