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

MODEL_ID = "arinbalyan/qwen-sql-copilot"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",
)

def generate(prompt, history):
    if not prompt.strip():
        return "", history
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.8,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    history.append({"role": "user", "content": prompt})
    history.append({"role": "assistant", "content": response.strip()})
    return "", history

theme = (
    gr.themes.Soft(primary_hue="blue", neutral_hue="slate")
    .set(
        button_primary_background_fill_hover="#2563eb",
        button_primary_text_color="white",
        input_background_fill="#f8fafc",
        input_border_color="#e2e8f0",
        input_border_color_focus="#2563eb",
    )
)

css = """.gradio-container { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important; }
.main-header { text-align: center; padding: 2rem 1rem; background: linear-gradient(135deg, #2563eb 0%, #2563ebdd 100%); color: white; border-radius: 12px; margin-bottom: 1rem; }
.main-header h1 { font-size: 2rem; font-weight: 700; margin-bottom: 0.25rem; }
.main-header p { font-size: 1rem; opacity: 0.9; }
.footer { text-align: center; padding: 1rem; color: #64748b; font-size: 0.875rem; }"""

with gr.Blocks(title="SQL Copilot") as demo:
    with gr.Row():
        with gr.Column():
            gr.HTML(f"""<div class="main-header"><h1>💾 SQL Copilot</h1><p>Generate with AI</p></div>""")
    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(label="Prompt", placeholder='e.g., Show all users older than 30', lines=2)
    with gr.Row():
        submit = gr.Button("Generate", variant="primary", size="lg", scale=1)
        clear = gr.Button("Clear", size="lg", scale=0)
    gr.Examples(examples=["Show all users older than 30", "How many orders were placed last month?", "List employees with their department names", "Find the top 5 products by sales", "Count customers from each city"], inputs=msg, label="Try these")
    gr.HTML(f"""<div class="footer">Qwen2.5-1.5B on sql-create-context</div>""")
    submit.click(generate, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(theme=theme, css=css)