from dotenv import load_dotenv
import os
import gradio as gr
from groq import Groq

load_dotenv()
api = os.getenv("groq_api_key")

def create_prompt(user_query, table_metadata):
    system_prompt = """
You are a SQL query generator specialized in generating SQL queries for a single table at a time.
Your task is to accurately convert natural language queries into SQL statements based on the user's intent and the provided table metadata.

Rules:
- Single Table Only: Use only the table in the metadata.
- Metadata-Based Validation: Use only columns in the metadata.
- User Intent: Support filters, grouping, sorting, etc.
- SQL Syntax: Use standard SQL (DuckDB compatible).
- Output only valid SQL. No extra commentary.

Input:
User Query: {user_query}
Table Metadata: {table_metadata}

Output:
SQL Query (on a single line, nothing else).
    """
    return system_prompt.strip(), f"User Query: {user_query}\nTable Metadata: {table_metadata}"

def generate_output(system_prompt, user_prompt):
    client = Groq(api_key=api)
    chat_completion = client.chat.completions.create(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        model="llama3-70b-8192"
    )
    response = chat_completion.choices[0].message.content.strip()
    return response if response.lower().startswith("select") else "Can't perform the task at the moment."

# NEW: accepts user_query and dynamic table_metadata string
def response(payload):
    user_query = payload.get("question", "")
    table_metadata = payload.get("schema", "")
    system_prompt, user_prompt = create_prompt(user_query, table_metadata)
    return generate_output(system_prompt, user_prompt)

demo = gr.Interface(
    fn=response,
    inputs=gr.JSON(label="Input JSON (question, schema)"),
    outputs="text",
    title="SQL Generator (Groq + LLaMA3)",
    description="Input: question & table metadata. Output: SQL using dynamic schema."
)

demo.launch()