from groq import Groq
from pydantic import BaseModel
import json
import gradio as gr

class ValidationStatus(BaseModel):
    is_valid: bool
    syntax_errors: list[str]

class SQLQueryGeneration(BaseModel):
    query: str
    query_type: str
    tables_used: list[str]
    estimated_complexity: str
    execution_notes: list[str]
    validation_status: ValidationStatus
    table_schema: str
    sample_data: str
    execution_results: str
    optimization_notes: list[str]

def generate_sql_query(api_key, user_query):
    """Generate SQL query from natural language using GROQ API"""
    try:
        if not api_key:
            return "Error: Please enter your GROQ API key", "", "", "", "", ""
        
        if not user_query:
            return "Error: Please enter a query description", "", "", "", "", ""
        
        client = Groq(api_key=api_key)
        
        response = client.chat.completions.create(
            model="moonshotai/kimi-k2-instruct-0905",
            messages=[
                {
                    "role": "system",
                    "content": """You are a SQL expert. Generate structured SQL queries from natural language descriptions with proper syntax validation and metadata.
After generating the SQL query, you must:
1. Create a sample SQL table schema based on the natural language description, including all necessary columns with appropriate data types
2. Populate the table with realistic sample data that demonstrates the query's functionality
3. Execute the generated SQL query against the sample table
4. Display the SQL table structure and data clearly
5. Show the query execution results
Always present your response in this order:
- Generated SQL query with syntax explanation
- Table schema (CREATE TABLE statement)
- Sample data (INSERT statements or table visualization)
- Query execution results
- Any relevant notes about assumptions made or query optimization suggestions""",
                },
                {
                    "role": "user", 
                    "content": user_query
                },
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": "sql_query_generation",
                    "schema": SQLQueryGeneration.model_json_schema()
                }
            }
        )
        
        sql_query_generation = SQLQueryGeneration.model_validate(
            json.loads(response.choices[0].message.content)
        )
        
        # Format validation status
        validation_text = f"Valid: {sql_query_generation.validation_status.is_valid}\n"
        if sql_query_generation.validation_status.syntax_errors:
            validation_text += "Errors:\n" + "\n".join(
                f"- {error}" for error in sql_query_generation.validation_status.syntax_errors
            )
        else:
            validation_text += "No syntax errors found"
        
        # Format metadata
        metadata = f"""Query Type: {sql_query_generation.query_type}
Tables Used: {', '.join(sql_query_generation.tables_used)}
Complexity: {sql_query_generation.estimated_complexity}

Execution Notes:
{chr(10).join(f"- {note}" for note in sql_query_generation.execution_notes)}

Optimization Notes:
{chr(10).join(f"- {note}" for note in sql_query_generation.optimization_notes)}"""
        
        return (
            sql_query_generation.query,
            metadata,
            sql_query_generation.table_schema,
            sql_query_generation.sample_data,
            sql_query_generation.execution_results,
            validation_text
        )
        
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        return error_msg, "", "", "", "", ""

# Create Gradio interface
with gr.Blocks(title="SQL Query Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🗄️ Natural Language to SQL Query Generator
        Convert your natural language descriptions into structured SQL queries with validation and execution results.
        """
    )
    
    with gr.Row():
        with gr.Column():
            api_key_input = gr.Textbox(
                label="GROQ API Key",
                type="password",
                placeholder="Enter your GROQ API key here...",
                info="Your API key is not stored and only used for this session"
            )
            
            query_input = gr.Textbox(
                label="Natural Language Query",
                placeholder="e.g., Find all the students who scored more than 90 out of 100",
                lines=3,
                value="Find all the students who scored more than 90 out of 100"
            )
            
            generate_btn = gr.Button("Generate SQL Query", variant="primary", size="lg")
            
            gr.Examples(
                examples=[
                    ["Find all the students who scored more than 90 out of 100"],
                    ["Get the top 5 customers by total purchase amount"],
                    ["List all employees hired in the last 6 months"],
                    ["Find products with price between $50 and $100"],
                    ["Show average salary by department"]
                ],
                inputs=query_input,
                label="Example Queries"
            )
    
    with gr.Row():
        with gr.Column():
            sql_output = gr.Code(
                label="Generated SQL Query",
                language="sql",
                lines=5
            )
            
            metadata_output = gr.Textbox(
                label="Query Metadata",
                lines=8
            )
            
            validation_output = gr.Textbox(
                label="Validation Status",
                lines=3
            )
    
    with gr.Row():
        with gr.Column():
            schema_output = gr.Code(
                label="Table Schema",
                language="sql",
                lines=8
            )
        
        with gr.Column():
            sample_data_output = gr.Code(
                label="Sample Data",
                language="sql",
                lines=8
            )
    
    with gr.Row():
        execution_output = gr.Textbox(
            label="Execution Results",
            lines=10
        )
    
    generate_btn.click(
        fn=generate_sql_query,
        inputs=[api_key_input, query_input],
        outputs=[
            sql_output,
            metadata_output,
            schema_output,
            sample_data_output,
            execution_output,
            validation_output
        ]
    )
    
    gr.Markdown(
        """
        ---
        ### How to use:
        1. Enter your GROQ API key (get one from [console.groq.com](https://console.groq.com))
        2. Type your natural language query description
        3. Click "Generate SQL Query" to see the results
        
        The app will provide:
        - A validated SQL query
        - Table schema and sample data
        - Execution results
        - Optimization suggestions
        """
    )

if __name__ == "__main__":
    demo.launch(share=True)