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

# Load model and tokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
    "onkolahmet/Qwen2-0.5B-Instruct-SQL-generator", 
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("onkolahmet/Qwen2-0.5B-Instruct-SQL-generator")

# # Few-shot examples to include in each prompt
# examples = [
#     {
#         "question": "Get the names and emails of customers who placed an order in the last 30 days.",
#         "sql": "SELECT name, email FROM customers WHERE order_date >= DATE_SUB(CURDATE(), INTERVAL 30 DAY);"
#     },
#     {
#         "question": "Find all employees with a salary greater than 50000.",
#         "sql": "SELECT * FROM employees WHERE salary > 50000;"
#     },
#     {
#         "question": "List all product names and their categories where the price is below 50.",
#         "sql": "SELECT name, category FROM products WHERE price < 50;"
#     },
#     {
#         "question": "How many users registered in the year 2022?",
#         "sql": "SELECT COUNT(*) FROM users WHERE YEAR(registration_date) = 2022;"
#     }
# ]

def generate_sql(question, context=None):
    # Construct prompt with few-shot examples and context if available
    prompt = "Translate natural language questions to SQL queries.\n\n"
    
    # Add table context if available
    if context and context.strip():
        prompt += f"Table Context:\n{context}\n\n"
    
    # # Add few-shot examples
    # for ex in examples:
    #     prompt += f"Q: {ex['question']}\nSQL: {ex['sql']}\n\n"
    
    # Add the current question
    prompt += f"Q: {question}\nSQL:"
    
    # Tokenize and generate
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    # Generate SQL query
    outputs = model.generate(
        inputs.input_ids,
        max_new_tokens=128,
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id
    )
    
    # Extract and decode only the new generation
    sql_query = tokenizer.decode(outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)
    return sql_query.strip()

def clean_sql_output(sql_text):
    """
    Clean and deduplicate SQL queries:
    1. Remove comments
    2. Remove duplicate queries
    3. Extract only the most relevant query
    4. Format properly
    """
    # Remove SQL comments (both single line and multi-line)
    sql_text = re.sub(r'--.*?$', '', sql_text, flags=re.MULTILINE)
    sql_text = re.sub(r'/\*.*?\*/', '', sql_text, flags=re.DOTALL)
    
    # Remove markdown code block syntax if present
    sql_text = re.sub(r'```sql|```', '', sql_text)
    
    # Split into individual queries if multiple exist
    if ';' in sql_text:
        queries = [q.strip() for q in sql_text.split(';') if q.strip()]
    else:
        # If no semicolons, try to identify separate queries by SELECT statements
        sql_text_cleaned = re.sub(r'\s+', ' ', sql_text)
        select_matches = list(re.finditer(r'SELECT\s+', sql_text_cleaned, re.IGNORECASE))
        
        if len(select_matches) > 1:
            queries = []
            for i in range(len(select_matches)):
                start = select_matches[i].start()
                end = select_matches[i+1].start() if i < len(select_matches) - 1 else len(sql_text_cleaned)
                queries.append(sql_text_cleaned[start:end].strip())
        else:
            queries = [sql_text]
    
    # Remove empty queries
    queries = [q for q in queries if q.strip()]
    
    if not queries:
        return ""
    
    # If we have multiple queries, need to deduplicate
    if len(queries) > 1:
        # Normalize queries for comparison (lowercase, remove extra spaces)
        normalized_queries = []
        for q in queries:
            # Use sqlparse to format and normalize
            try:
                formatted = sqlparse.format(
                    q + ('' if q.strip().endswith(';') else ';'), 
                    keyword_case='lower',
                    identifier_case='lower', 
                    strip_comments=True,
                    reindent=True
                )
                normalized_queries.append(formatted)
            except:
                # If sqlparse fails, just do basic normalization
                normalized = re.sub(r'\s+', ' ', q.lower().strip())
                normalized_queries.append(normalized)
        
        # Find unique queries
        unique_queries = []
        unique_normalized = []
        
        for i, norm_q in enumerate(normalized_queries):
            if norm_q not in unique_normalized:
                unique_normalized.append(norm_q)
                unique_queries.append(queries[i])
        
        # Choose the most likely correct query:
        # 1. Prefer queries with SELECT
        # 2. Prefer longer queries (often more detailed)
        # 3. Prefer first query if all else equal
        select_queries = [q for q in unique_queries if re.search(r'SELECT\s+', q, re.IGNORECASE)]
        
        if select_queries:
            # Choose the longest SELECT query (likely most detailed)
            best_query = max(select_queries, key=len)
        elif unique_queries:
            # If no SELECT queries, choose the longest query
            best_query = max(unique_queries, key=len)
        else:
            # Fallback to the first query
            best_query = queries[0]
    else:
        best_query = queries[0]
    
    # Clean up the chosen query
    best_query = best_query.strip()
    if not best_query.endswith(';'):
        best_query += ';'
    
    # Final formatting to ensure consistent spacing
    best_query = re.sub(r'\s+', ' ', best_query)
    
    try:
        # Use sqlparse to nicely format the SQL for display
        formatted_sql = sqlparse.format(
            best_query,
            keyword_case='upper',
            identifier_case='lower',
            reindent=True,
            indent_width=2
        )
        return formatted_sql
    except:
        return best_query

def process_input(question, table_context):
    """Function to process user input through the model and return formatted results"""
    if not question.strip():
        return "Please enter a question."
    
    # Generate SQL from the question and context
    raw_sql = generate_sql(question, table_context)
    
    # Clean the SQL output
    cleaned_sql = clean_sql_output(raw_sql)
    
    if not cleaned_sql:
        return "Sorry, I couldn't generate a valid SQL query. Please try rephrasing your question."
    
    return cleaned_sql

# Sample table context examples for the example selector
example_contexts = [
    # Example 1
    """
CREATE TABLE customers (
  id INT PRIMARY KEY,
  name VARCHAR(100),
  email VARCHAR(100),
  order_date DATE
);
    """,
    
    # Example 2
    """
CREATE TABLE products (
  id INT PRIMARY KEY,
  name VARCHAR(100),
  category VARCHAR(50),
  price DECIMAL(10,2),
  stock_quantity INT
);
    """,
    
    # Example 3
    """
CREATE TABLE employees (
  id INT PRIMARY KEY,
  name VARCHAR(100),
  department VARCHAR(50),
  salary DECIMAL(10,2),
  hire_date DATE
);

CREATE TABLE departments (
  id INT PRIMARY KEY,
  name VARCHAR(50),
  manager_id INT,
  budget DECIMAL(15,2)
);
    """
]

# Sample question examples
example_questions = [
    "Get the names and emails of customers who placed an order in the last 30 days.",
    "Find all products with less than 10 items in stock.",
    "List all employees in the Sales department with a salary greater than 50000.",
    "What is the total budget for departments with more than 5 employees?",
    "Count how many products are in each category where the price is greater than 100."
]

# Create the Gradio interface
with gr.Blocks(title="Text to SQL Converter") as demo:
    gr.Markdown("# Text to SQL Query Converter")
    gr.Markdown("Enter your question and optional table context to generate an SQL query.")
    
    with gr.Row():
        with gr.Column():
            question_input = gr.Textbox(
                label="Your Question",
                placeholder="e.g., Find all products with price less than $50",
                lines=2
            )
            
            table_context = gr.Textbox(
                label="Table Context (Optional)",
                placeholder="Enter your database schema or table definitions here...",
                lines=10
            )
            
            submit_btn = gr.Button("Generate SQL Query")
        
        with gr.Column():
            sql_output = gr.Code(
                label="Generated SQL Query",
                language="sql",
                lines=12
            )
    
    # Examples section
    gr.Markdown("### Try some examples")
    
    example_selector = gr.Examples(
        examples=[
            ["List all products in the 'Electronics' category with price less than $500", example_contexts[1]],
            ["Find the total number of employees in each department", example_contexts[2]],
            ["Get customers who placed orders in the last 7 days", example_contexts[0]],
            ["Count the number of products in each category", example_contexts[1]],
            ["Find the average salary by department", example_contexts[2]]
        ],
        inputs=[question_input, table_context]
    )
    
    # Set up the submit button to trigger the process_input function
    submit_btn.click(
        fn=process_input,
        inputs=[question_input, table_context],
        outputs=sql_output
    )
    
    # Also trigger on pressing Enter in the question input
    question_input.submit(
        fn=process_input,
        inputs=[question_input, table_context],
        outputs=sql_output
    )
    
    # Add information about the model
    gr.Markdown("""
    ### About
    This app uses a fine-tuned language model to convert natural language questions into SQL queries.
    
    - **Model**: [onkolahmet/Qwen2-0.5B-Instruct-SQL-generator](https://huggingface.co/onkolahmet/Qwen2-0.5B-Instruct-SQL-generator)
    - **How to use**: 
      1. Enter your question in natural language
      2. If you have specific table schemas, add them in the Table Context field
      3. Click "Generate SQL Query" or press Enter
      
    Note: The model works best when table context is provided, but can generate generic SQL queries without it.
    """)

# Launch the app
demo.launch()