import gradio as gr
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Load tokenizer and model from local folder
model_path = "./"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = T5Tokenizer.from_pretrained(model_path)
model = T5ForConditionalGeneration.from_pretrained(model_path).to(device)
model.eval()

def generate_sql(schema, user_query, date_info=""):
    prompt = f"""Given the following database schema and requirements, generate a PostgreSQL query:

{schema}

User Query: "{user_query}"

IMPORTANT REQUIREMENTS:
- Always filter by user_id = $1 for security
- The current year is 2025. You are working in this year!
- CRITICAL: Use ONLY the dates provided in the input parameters. Do NOT infer or change dates on your own!
- Single date provided: {date_info if date_info else '[none provided]'}
- If date range is provided, use DATE(created_at) BETWEEN 'startDate' AND 'endDate'
- If single date is provided, use DATE(created_at) = 'YYYY-MM-DD'
- NEVER use hardcoded years like 2024 - always use the provided dates exactly as given
- Generate ONLY simple SELECT, INSERT, UPDATE, DELETE statements.
- NO WITH clauses, NO CTEs, NO complex subqueries.
- Use RETURNING * for INSERT, UPDATE, DELETE when applicable.

SQL Query:"""

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True).to(device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=256)
    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_sql

# Gradio UI
iface = gr.Interface(
    fn=generate_sql,
    inputs=[
        gr.Textbox(label="Database Schema", lines=12, placeholder="CREATE TABLE ..."),
        gr.Textbox(label="User Query", placeholder="How much did I spend on food last week?"),
        gr.Textbox(label="Date Info (optional)", placeholder="2025-06-12 or 2025-06-01 to 2025-06-07")
    ],
    outputs=gr.Textbox(label="Generated SQL Query"),
    title="HISAB AI - Natural Language to SQL",
    description="Enter your schema, user query and date (optional). Model will output SQL query."
)

iface.launch()