import os

# Disable Gradio SSR (better for Render)
os.environ["GRADIO_SSR_MODE"] = "False"

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch


MODEL_NAME = "gaussalgo/T5-LM-Large-text2sql-spider"


tokenizer = None
model = None

device = "cuda" if torch.cuda.is_available() else "cpu"


def load_model():

    global tokenizer, model

    if model is None:

        print("Loading model...")

        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME
        )

        model = AutoModelForSeq2SeqLM.from_pretrained(
            MODEL_NAME
        )

        model.to(device)

        model.eval()

        print(f"Model ready on {device}")



def generate_sql(question, context):

    try:

        load_model()


        input_text = f"{question} | {context}"


        inputs = tokenizer(
            input_text,
            return_tensors="pt",
            max_length=512,
            truncation=True
        ).to(device)


        with torch.no_grad():

            outputs = model.generate(
                **inputs,
                max_new_tokens=128,
                num_beams=4,
                early_stopping=True
            )


        sql = tokenizer.decode(
            outputs[0],
            skip_special_tokens=True
        )


        return sql


    except Exception as e:

        return f"Error: {str(e)}"




with gr.Blocks() as demo:


    gr.Markdown(
        "# NL2SQL API\nGenerate SQL from Natural Language"
    )


    with gr.Row():

        question = gr.Textbox(
            label="Question",
            placeholder="Example: Find all users"
        )


        context = gr.Textbox(
            label="Database Schema",
            placeholder="Example: users(id,name,email)"
        )


    output = gr.Textbox(
        label="Generated SQL",
        lines=5
    )


    btn = gr.Button(
        "Generate SQL"
    )


    btn.click(
        fn=generate_sql,
        inputs=[
            question,
            context
        ],
        outputs=output
    )



if __name__ == "__main__":

    demo.launch(
        server_name="0.0.0.0",
        server_port=int(
            os.environ.get("PORT",7860)
        ),
        share=False
    )