import gradio as gr
import openai
import os

# OpenRouter API Key
OPENROUTER_API_KEY = "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"
OPENROUTER_MODEL = "sophosympatheia/rogue-rose-103b-v0.2:free"

# Initialize OpenAI client with OpenRouter base URL
print(f"Using API Key: {OPENROUTER_API_KEY}")
openai_client = openai.OpenAI(
    api_key=OPENROUTER_API_KEY,
    base_url="https://openrouter.ai/api/v1"  # OpenRouter API endpoint
)

# Few-shot examples for text-to-SQL conversion
few_shot_examples = [
    {
        "input": "Show all customers from the USA.",
        "output": "SELECT * FROM customers WHERE country = 'USA';"
    },
    {
        "input": "Find the total sales for each product category.",
        "output": "SELECT product_category, SUM(sales) AS total_sales FROM sales GROUP BY product_category;"
    },
    {
        "input": "List all orders placed in 2023.",
        "output": "SELECT * FROM orders WHERE YEAR(order_date) = 2023;"
    }
]

def text_to_sql(query):
    # Construct the prompt with few-shot examples
    prompt = "Convert the following natural language queries to SQL:\n\n"
    for example in few_shot_examples:
        prompt += f"Input: {example['input']}\nOutput: {example['output']}\n\n"
    prompt += f"Input: {query}\nOutput:"

    print("Sending query to OpenRouter API...")
    try:
        response = openai_client.chat.completions.create(
            model=OPENROUTER_MODEL,
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful assistant. Your task is to convert natural language queries into SQL queries. "
                               "Use the provided examples as a guide. If the query cannot be converted into SQL, say 'I cannot convert this query into SQL.'"
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ]
        )
        print("Received response from OpenRouter API.")
        return response.choices[0].message.content
    except Exception as e:
        print(f"Error calling OpenRouter API: {e}")
        return f"Error: {e}"

# Gradio UI
def gradio_ui():
    with gr.Blocks() as demo:
        gr.Markdown("## Text-to-SQL Converter. Enter a natural language query and get the corresponding SQL query!")
        query_input = gr.Textbox(label="Enter your query")
        submit_btn = gr.Button("Convert to SQL")
        output = gr.Textbox(label="SQL Query")
        
        submit_btn.click(text_to_sql, inputs=[query_input], outputs=[output])
    
    return demo

demo = gradio_ui()

print("Launching Gradio UI...")
demo.launch()