import gradio as gr
import pandas as pd
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Load NL2SQL model from Hugging Face (no API key needed)
model_name = "PaulGan1/t5-small-spider-sql"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Global dataframe
df = None

def load_data(file):
    global df
    df = pd.read_csv(file.name)
    return f"✅ Data uploaded successfully! Columns: {', '.join(df.columns)}"

def generate_sql(question):
    if df is None:
        return "⚠️ Please upload a CSV file first."
    
    # Create prompt
    columns = ", ".join(df.columns)
    prompt = f"translate English to SQL: {question} | table: {columns}"
    
    # Generate SQL
    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
    outputs = model.generate(**inputs, max_length=128)
    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Execute SQL
    try:
        result = pd.read_sql_query(sql_query, con=df)
        return f"🧠 SQL Query:\n{sql_query}\n\n📊 Result:\n{result.head()}"
    except Exception:
        return f"🧠 SQL Query:\n{sql_query}\n\n⚠️ Unable to execute SQL. (Demo only)"

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Natural Language to SQL Query Generator")
    file_input = gr.File(label="Upload your CSV file")
    upload_output = gr.Textbox(label="Upload Status")
    question = gr.Textbox(label="Ask your question in natural language:")
    sql_output = gr.Textbox(label="Generated SQL Query & Output", lines=10)

    file_input.change(load_data, inputs=file_input, outputs=upload_output)
    question.submit(generate_sql, inputs=question, outputs=sql_output)

demo.launch()